From c0729610569b0193eb6294e00c4c2dd4a0593b17 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 17 Apr 2023 16:31:54 +0100
Subject: [PATCH 001/428] Fix loop setup when a dimension has size zero

For example, if there are no neutral species.
---
 src/looping.jl | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/looping.jl b/src/looping.jl
index 6449f8cdd..08018cdd8 100644
--- a/src/looping.jl
+++ b/src/looping.jl
@@ -127,6 +127,12 @@ function get_local_range(sub_block_rank, sub_block_size, dim_size)
     # because the root process might have slightly more work to do in general.
     # This calculation is not at all optimized, but is not going to take long, and is
     # only done in initialization, so it is more important to be simple and robust.
+
+    if dim_size == 0
+        # No processor includes a grid point
+        return 1:0
+    end
+
     remaining = dim_size
     done = false
     n_points_for_proc = zeros(mk_int, sub_block_size)

From 54484a0d7f9e655001b77ec20b10de2fd3269127 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 17 Apr 2023 16:32:42 +0100
Subject: [PATCH 002/428] Workaround for NetCDF not handling Bool variables

Save bools as type `Char` in NetCDF files, because NetCDF does not
support `Bool`. We don't use `Char` for anything, so this should not
cause a problem.
---
 src/file_io_netcdf.jl | 11 +++++++++--
 src/load_data.jl      |  6 +++++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 47c24d86c..634d3dfe0 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -67,8 +67,15 @@ function write_single_value!(file_or_group::NCDataset, name,
         end
         dims = Tuple(c.name for c in coords)
     end
-    var = defVar(file_or_group, name, type, dims, attrib=attributes)
-    var[:] = value
+    if isa(value, Bool)
+        # As a hack, write bools to NetCDF as Char, as NetCDF does not support bools (?),
+        # and we do not use Char for anything else
+        var = defVar(file_or_group, name, Char, dims, attrib=attributes)
+        var[:] = Char(value)
+    else
+        var = defVar(file_or_group, name, type, dims, attrib=attributes)
+        var[:] = value
+    end
 
     return nothing
 end
diff --git a/src/load_data.jl b/src/load_data.jl
index bb208ffb5..4bdbc11ac 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -89,7 +89,11 @@ function load_variable(file_or_group::NCDataset, name::String)
     # This overload deals with cases where fid is a NetCDF `Dataset` (which could be a
     # file or a group).
     try
-        return file_or_group[name].var[:]
+        var = file_or_group[name].var[:]
+        if isa(var, Char)
+            var = (var == Char(true))
+        end
+        return var
     catch
         println("An error occured while loading $name")
         rethrow()

From b1903380ffa845da0c4a8f1a78c1e9dfb399097e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 May 2023 13:08:32 +0100
Subject: [PATCH 003/428] Workaround for errors when creating size-0
 shared-memory arrays

With some MPI implementations, creating a size-0 shared-memory array may
cause an error. Shared memory is not necessary, as these arrays can
never be accessed, so work around the error by just creating a regular
Julia array.
---
 src/communication.jl | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/src/communication.jl b/src/communication.jl
index 706476610..9ca3e9643 100644
--- a/src/communication.jl
+++ b/src/communication.jl
@@ -314,6 +314,19 @@ function allocate_shared(T, dims)
     bs = block_size[]
     n = prod(dims)
 
+    if n == 0
+        # Special handling as some MPI implementations cause errors when allocating a
+        # size-zero array
+        array = Array{T}(undef, dims...)
+
+        @debug_shared_array begin
+            # If @debug_shared_array is active, create DebugMPISharedArray instead of Array
+            array = DebugMPISharedArray(array)
+        end
+
+        return array
+    end
+
     if br == 0
         # Allocate points on rank-0 for simplicity
         n_local = n

From d6879f68505466cf5f9af43f56318071de768627 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 16 May 2023 12:23:46 +0100
Subject: [PATCH 004/428] Update expected results for a couple of failing
 'long' tests

Differences do not look significant. The difference probably just comes
from changed rounding errors in not-too-well-resolved test cases.
---
 test/sound_wave_tests.jl | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/test/sound_wave_tests.jl b/test/sound_wave_tests.jl
index f3af07242..23f9e3732 100644
--- a/test/sound_wave_tests.jl
+++ b/test/sound_wave_tests.jl
@@ -260,13 +260,15 @@ function run_test_set_finite_difference()
 
     # n_i>>n_n T_e=1
     @long run_test(test_input_finite_difference, 2*π*1.4467, -2*π*0.6020,
-                   [-0.0010684224665919893, -0.0010505277216983934,
-                    -0.0010288041337547594, -0.0010033394223312585,
-                    -0.0009742364063434105, -0.0009416125854969064];
+                   [-0.001068422466592656, -0.001050527721698838, -0.0010288041337549816,
+                    -0.0010033394223323698, -0.0009742364063442995,
+                    -0.000941612585497573];
                    initial_density1=0.9999, initial_density2=0.0001)
     @long run_test(test_input_finite_difference, 2*π*1.4467, -2*π*0.6020,
-                   [-0.0010684224665919893, -0.0010505277216983934, -0.0010288041337547594, -0.0010033394223312585, -0.0009742364063434105, -0.0009416125854969064]; initial_density1=0.9999, initial_density2=0.0001,
-                   charge_exchange_frequency=2*π*2.0)
+                   [-0.001068422466592656, -0.001050527721698838, -0.0010288041337549816,
+                    -0.0010033394223323698, -0.0009742364063442995,
+                    -0.000941612585497573]; initial_density1=0.9999,
+                    initial_density2=0.0001, charge_exchange_frequency=2*π*2.0)
 
     # n_i<<n_n T_e=1
     @long run_test(test_input_finite_difference, 2*π*1.3954, -2*π*0.6815,
@@ -494,15 +496,14 @@ function run_test_set_chebyshev()
 
     # n_i>>n_n T_e=1
     @long run_test(test_input_chebyshev, 2*π*1.4467, -2*π*0.6020,
-                   [-0.00010000500033334732, 0.00046539975104573, 0.0007956127502551822,
-                    0.0008923624901797472, 0.0008994953327500175,
-                    0.0008923624901797472]; initial_density1=0.9999,
-                    initial_density2=0.0001)
+                   [-0.00010000500033334732, 0.0004653997510461739, 0.0007956127502558478,
+                    0.0008923624901804128, 0.0008994953327500175, 0.0008923624901804128];
+                   initial_density1=0.9999, initial_density2=0.0001)
     @long run_test(test_input_chebyshev, 2*π*1.4467, -2*π*0.6020,
-                   [-0.00010000500033334732, 0.00046539975104573, 0.0007956127502551822,
-                    0.0008923624901797472, 0.0008994953327500175,
-                    0.0008923624901797472]; initial_density1=0.9999,
-                   initial_density2=0.0001, charge_exchange_frequency=2*π*2.0)
+                   [-0.00010000500033334732, 0.0004653997510461739, 0.0007956127502558478,
+                    0.0008923624901804128, 0.0008994953327500175, 0.0008923624901804128];
+                   initial_density1=0.9999, initial_density2=0.0001,
+                   charge_exchange_frequency=2*π*2.0)
 
     # n_i<<n_n T_e=1
     @long run_test(test_input_chebyshev, 2*π*1.3954, -2*π*0.6815,

From 9d33e53264035fa13725b5585c62dafce7254d86 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 17 May 2023 17:06:53 +0100
Subject: [PATCH 005/428] Option to speed up debug checks by avoiding stack
 traces

Only record the stacktrace showing where each array was allocated when
`@debug_track_array_allocate_location` is active.

Also move `@debug_block_synchronize` and `@debug_shared_array_allocate`
to debug level 4, so that at debug level 2, no stack traces are created
and at level 3 a minimum number are created.

The two changes above should significantly speed up debug level 2, while
still allowing errors in shared-memory usage to be detected. Now use
debug level 2 for the CI tests.
---
 .github/workflows/debug_checks.yml | 10 +++++-----
 src/communication.jl               | 19 ++++++++++++++-----
 src/debugging.jl                   | 13 ++++++++-----
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/debug_checks.yml b/.github/workflows/debug_checks.yml
index 9f5497fef..6f0601822 100644
--- a/.github/workflows/debug_checks.yml
+++ b/.github/workflows/debug_checks.yml
@@ -35,7 +35,7 @@ jobs:
   #        # Hard code the debug level so that we can run without using the
   #        # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
   #        # moment.
-  #        sed -i -e "s/_debug_level = get_options.*/_debug_level = 3/" src/debugging.jl
+  #        sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
 
   #        pip3 install --user matplotlib
   #        julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary()'
@@ -78,7 +78,7 @@ jobs:
           # Hard code the debug level so that we can run without using the
           # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
           # moment.
-          sed -i -e "s/_debug_level = get_options.*/_debug_level = 3/" src/debugging.jl
+          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
 
           pip3 install --user matplotlib
           julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary()'
@@ -121,7 +121,7 @@ jobs:
           # Hard code the debug level so that we can run without using the
           # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
           # moment.
-          sed -i -e "s/_debug_level = get_options.*/_debug_level = 3/" src/debugging.jl
+          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
 
           pip3 install --user matplotlib
           julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary()'
@@ -164,7 +164,7 @@ jobs:
           # Hard code the debug level so that we can run without using the
           # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
           # moment.
-          sed -i -e "s/_debug_level = get_options.*/_debug_level = 3/" src/debugging.jl
+          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
 
           pip3 install --user matplotlib
           julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary()'
@@ -207,7 +207,7 @@ jobs:
           # Hard code the debug level so that we can run without using the
           # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
           # moment.
-          sed -i -e "s/_debug_level = get_options.*/_debug_level = 3/" src/debugging.jl
+          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
 
           pip3 install --user matplotlib
           julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary()'
diff --git a/src/communication.jl b/src/communication.jl
index 9ca3e9643..957bd0587 100644
--- a/src/communication.jl
+++ b/src/communication.jl
@@ -228,7 +228,9 @@ end
         is_read .= false
         is_written = Array{Bool}(undef, dims)
         is_written .= false
-        creation_stack_trace = string([string(s, "\n") for s in stacktrace()]...)
+        creation_stack_trace = @debug_track_array_allocate_location_ifelse(
+                                   string([string(s, "\n") for s in stacktrace()]...),
+                                   "")
         @debug_detect_redundant_block_synchronize begin
             # Initialize as `true` so that the first call to _block_synchronize() with
             # @debug_detect_redundant_block_synchronize activated does not register the
@@ -514,10 +516,17 @@ end
                     # shared-memory arrays, so would cause segfaults.
                     return false
                 else
-                    error("Shared memory array written at $i from multiple ranks "
-                          * "between calls to _block_synchronize(). Array was "
-                          * "created at:\n"
-                          * array.creation_stack_trace)
+                    if array.creation_stack_trace != ""
+                        error("Shared memory array written at $i from multiple ranks "
+                              * "between calls to _block_synchronize(). Array was "
+                              * "created at:\n"
+                              * array.creation_stack_trace)
+                    else
+                        error("Shared memory array written at $i from multiple ranks "
+                              * "between calls to _block_synchronize(). Enable "
+                              * "`debug_track_array_allocate_location` to track where "
+                              * "array was created.")
+                    end
                 end
             elseif n_writes == 1 && n_reads > 0
                 if global_is_written[i, block_rank[] + 1]
diff --git a/src/debugging.jl b/src/debugging.jl
index 337cfa19b..82594b2f2 100644
--- a/src/debugging.jl
+++ b/src/debugging.jl
@@ -23,16 +23,19 @@ macronames = [
     ("debug_error_stop_all", 1,
     "Use MPI.Allgather to stop all processes following an error on any process."),
 
-    ("debug_block_synchronize", 2,
-     "Check _block_synchronize() was called from the same place on every process."),
-
     ("debug_shared_array", 2,
      "Check for incorrect reads/writes to shared-memory arrays"),
 
-    ("debug_shared_array_allocate", 3,
+    ("debug_track_array_allocate_location", 3,
+     "Record where every array was allocated."),
+
+    ("debug_shared_array_allocate", 4,
      "Check that allocate_shared() was called from the same place on every process."),
 
-    ("debug_detect_redundant_block_synchronize", 4,
+    ("debug_block_synchronize", 4,
+     "Check _block_synchronize() was called from the same place on every process."),
+
+    ("debug_detect_redundant_block_synchronize", 5,
      "Check if any _block_synchronize() call could have been skipped without resulting "
      * "in an error.")
 ]

From 78fcd8a4eb754b985102ba0dbb0145f36e77b9dc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 11:30:54 +0100
Subject: [PATCH 006/428] Reduce timeout for examples workflow

Workflow now only runs 10 timesteps for each case, so does not take too
long.
---
 .github/workflows/examples.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml
index 375f4b1d0..763751f7d 100644
--- a/.github/workflows/examples.yml
+++ b/.github/workflows/examples.yml
@@ -10,7 +10,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macOS-latest]
       fail-fast: false
-    timeout-minutes: 45
+    timeout-minutes: 25
 
     steps:
       - uses: actions/checkout@v2

From bdc020900d86e49a31506c6fb43f9551c93d6b45 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 11:37:10 +0100
Subject: [PATCH 007/428] Remove no longer used `enforce_moment_constraints!()`

The approach used in this function is no longer used, only
hard_force_moment_constraints!() and
hard_force_moment_constraints_neutral!() are used now.
---
 src/moment_constraints.jl | 124 +-------------------------------------
 src/time_advance.jl       |   3 -
 2 files changed, 1 insertion(+), 126 deletions(-)

diff --git a/src/moment_constraints.jl b/src/moment_constraints.jl
index e130c805c..81dd05cdf 100644
--- a/src/moment_constraints.jl
+++ b/src/moment_constraints.jl
@@ -10,129 +10,7 @@ using ..initial_conditions: enforce_zero_incoming_bc!
 using ..looping
 using ..velocity_moments: integrate_over_vspace, update_qpar!
 
-export enforce_moment_constraints!, hard_force_moment_constraints!,
-       hard_force_moment_constraints_neutral!
-
-"""
-    enforce_moment_constraints!(fvec_new, fvec_old, vpa, z, r, composition, moments, dummy_sr)
-
-Force moment constraints to be true for fvec_new if they were true for fvec_old. Should
-always be a small correction because the time step is 'small' so fvec_new should only
-violate the constraints by a small amount.
-"""
-function enforce_moment_constraints!(fvec_new, fvec_old, vpa, z, r, composition, moments, dummy_sr)
-    # pre-calculate avgdens_ratio so that we don't read fvec_new.density[:,is] on every
-    # process in the next loop - that would be an error because different processes
-    # write to fvec_new.density[:,is]
-    # This loop needs to be @loop_s_r because it fills the (not-shared)
-    # dummy_sr buffer to be used within the @loop_s_r below, so the values
-    # of is looped over by this process need to be the same.
-    # Need to call _block_synchronize() even though loop type does not change because
-    # all spatial ranks read fvec_new.density, but it will be written below.
-    if any(moments.particle_number_conserved)
-        _block_synchronize()
-    end
-    @loop_s_r is ir begin
-        if moments.particle_number_conserved[is]
-            @views @. z.scratch = fvec_old.density[:,ir,is] - fvec_new.density[:,ir,is]
-            @views dummy_sr[ir,is] = integral(z.scratch, z.wgts)/integral(fvec_old.density[:,ir,is], z.wgts)
-        end
-    end
-    # Need to call _block_synchronize() even though loop type does not change because
-    # all spatial ranks read fvec_new.density, but it will be written below.
-    if any(moments.particle_number_conserved)
-        _block_synchronize()
-    end
-
-    @loop_s is begin
-        # add a small correction to the density for each species to ensure that
-        # that particle number is conserved if it should be;
-        # ionisation collisions and net particle flux out of the domain due to, e.g.,
-        # a wall BC break particle conservation, in which cases it should not be enforced.
-        if moments.particle_number_conserved[is]
-            @loop_r ir begin
-                avgdens_ratio = dummy_sr[ir,is]
-                @loop_z iz begin
-                    # update the density with the above factor to ensure particle conservation
-                    fvec_new.density[iz,ir,is] += fvec_old.density[iz,ir,is] * avgdens_ratio
-                    # update the thermal speed, as the density has changed
-                    moments.vth[iz,ir,is] = sqrt(2.0*fvec_new.ppar[iz,ir,is]/fvec_new.density[iz,ir,is])
-                end
-            end
-        end
-        @loop_r ir begin
-            if moments.evolve_upar && is ∈ composition.ion_species_range && z.bc == "wall"
-                # Enforce zero-incoming boundary condition on the old distribution
-                # function with the new parallel flow, then force the updated old
-                # distribution function to obey the integral constraints exactly, which
-                # should be a small correction here as the boundary condition should
-                # only modify a few points due to the small change in upar.
-                # This procedure ensures that fvec_old obeys both the new boundary
-                # conditions and the moment constraints, so that when it is used to
-                # update f_new, f_new also does.
-                # Note fvec_old is never used after this function, so it is OK to modify
-                # it in-place.
-
-                # define a zero that accounts for finite precision
-                zero = 1.0e-10
-
-                @views enforce_zero_incoming_bc!(
-                    fvec_old.pdf[:,:,ir,is], vpa, fvec_new.density[:,ir,is],
-                    fvec_new.upar[:,ir,is], fvec_new.ppar[:,ir,is], moments.evolve_upar,
-                    moments.evolve_ppar, zero)
-                # Correct fvec_old.pdf in case applying new bc messed up moment
-                # constraints
-                @views hard_force_moment_constraints!(fvec_old.pdf[:,1,ir,is], moments,
-                                                      vpa)
-                @views hard_force_moment_constraints!(fvec_old.pdf[:,end,ir,is],
-                                                      moments, vpa)
-            end
-            @loop_z iz begin
-                # Create views once to save overhead
-                fnew_view = @view(fvec_new.pdf[:,iz,ir,is])
-                fold_view = @view(fvec_old.pdf[:,iz,ir,is])
-
-                # first calculate all of the integrals involving the updated pdf fvec_new.pdf
-                density_integral = integrate_over_vspace(fnew_view, vpa.wgts)
-                if moments.evolve_upar
-                    upar_integral = integrate_over_vspace(fnew_view, vpa.grid, vpa.wgts)
-                end
-                if moments.evolve_ppar
-                    ppar_integral = integrate_over_vspace(fnew_view, vpa.grid, 2, vpa.wgts)
-                end
-                # update the pdf to account for the density-conserving correction
-                @. fnew_view += fold_view * (1.0 - density_integral)
-                if moments.evolve_upar
-                    if !moments.evolve_ppar
-                        upar_coefficient = upar_integral /
-                            integrate_over_vspace(fold_view, vpa.grid, 2, vpa.wgts)
-                    else
-                        vpa3_moment = integrate_over_vspace(fold_view, vpa.grid, 3, vpa.wgts)
-                        vpa4_moment = integrate_over_vspace(fold_view, vpa.grid, 4, vpa.wgts)
-                        ppar_coefficient = (-ppar_integral + 0.5*density_integral +
-                                            2.0*upar_integral*vpa3_moment) /
-                                            (vpa4_moment - 0.25 - 2.0*vpa3_moment*vpa3_moment)
-                        upar_coefficient = 2.0 * (upar_integral + ppar_coefficient * vpa3_moment)
-                    end
-                    # update the pdf to account for the momentum-conserving correction
-                    @. fnew_view -= upar_coefficient * vpa.grid * fold_view
-                    if moments.evolve_ppar
-                        # update the pdf to account for the energy-conserving correction
-                        #@. fnew_view += ppar_coefficient * (vpa.grid^2 - 0.5) * fold_view
-                        # Until julia-1.8 is released, prefer x*x to x^2 to avoid
-                        # extra allocations when broadcasting.
-                        @. fnew_view += ppar_coefficient * (vpa.grid * vpa.grid - 0.5) * fold_view
-                    end
-                end
-            end
-        end
-    end
-    # the pdf, density and thermal speed have been changed so the corresponding parallel heat flux must be updated
-    moments.charged.qpar_updated .= false
-    update_qpar!(moments.qpar, moments.charged.qpar_updated, fvec_new.density, fvec_new.upar,
-                 moments.vth, fvec_new.pdf, vpa, z, r, composition,
-                 moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
-end
+export hard_force_moment_constraints!, hard_force_moment_constraints_neutral!
 
 """
     hard_force_moment_constraints!(f, moments, vpa)
diff --git a/src/time_advance.jl b/src/time_advance.jl
index db45559d5..7313d6e2b 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -328,7 +328,6 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
                 neutral_z_advect, nothing, nothing, neutral_vz_advect, r, z, vzeta, vr,
                 vz, composition, geometry, scratch_dummy)
             begin_sn_r_z_region()
-            #enforce_moment_constraints!(new_scratch, scratch[1], vpa, z, r, composition, moments, scratch_dummy_sr)
             @loop_sn_r_z isn ir iz begin
                 @views hard_force_moment_constraints_neutral!(
                     pdf.neutral.norm[:,:,:,iz,ir,isn], moments, vz)
@@ -1119,7 +1118,6 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 
     if moments.evolve_density && moments.enforce_conservation
         begin_s_r_z_region()
-        #enforce_moment_constraints!(new_scratch, scratch[1], vpa, z, r, composition, moments, scratch_dummy_sr)
         @loop_s_r_z is ir iz begin
             @views hard_force_moment_constraints!(new_scratch.pdf[:,:,iz,ir,is], moments,
                                                   vpa)
@@ -1192,7 +1190,6 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 
         if moments.evolve_density && moments.enforce_conservation
             begin_sn_r_z_region()
-            #enforce_moment_constraints!(new_scratch, scratch[1], vpa, z, r, composition, moments, scratch_dummy_sr)
             @loop_sn_r_z isn ir iz begin
                 @views hard_force_moment_constraints_neutral!(
                     new_scratch.pdf_neutral[:,:,:,iz,ir,isn], moments, vz)

From ceaa4d5f4c0f69370a9489f6ed4daecb01bb4f26 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 11:39:10 +0100
Subject: [PATCH 008/428] Remove ion_species_range, neutral_species_range from
 species_composition

These ranges are no longer needed, as ion and neutral arrays are stored
separately, so ions and neutrals are not covered by the same species
index.
---
 src/input_structs.jl         | 4 ----
 src/moment_kinetics_input.jl | 5 +++--
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/input_structs.jl b/src/input_structs.jl
index c8c8aa64b..3b47a8a8f 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -241,10 +241,6 @@ mutable struct species_composition
     # if false -- wall bc uses true Knudsen cosine to specify neutral pdf leaving the wall
     # if true -- use a simpler pdf that is easier to integrate
     use_test_neutral_wall_pdf::Bool
-    # Species indices that represent ions
-    ion_species_range::UnitRange{mk_int}
-    # Species indices that represent neutrals
-    neutral_species_range::UnitRange{mk_int}
     # electron temperature used for Boltzmann response
     T_e::mk_float
     # wall temperature used if 'wall' BC selected for z coordinate; normalised by electron temperature
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index f3a2abd51..8761fa3a2 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -773,8 +773,9 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     # ratio of the electron particle mass to the ion particle mass
     me_over_mi = 1.0/1836.0
     composition = species_composition(n_species, n_ion_species, n_neutral_species,
-        electron_physics, use_test_neutral_wall_pdf, 1:n_ion_species, n_ion_species+1:n_species, T_e, T_wall,
-        phi_wall, Er_constant, epsilon_offset, use_vpabar_in_mms_dfni, alpha_switch, mn_over_mi, me_over_mi, allocate_float(n_species))
+        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, Er_constant,
+        epsilon_offset, use_vpabar_in_mms_dfni, alpha_switch, mn_over_mi, me_over_mi,
+        allocate_float(n_species))
     
     species_charged = Array{species_parameters_mutable,1}(undef,n_ion_species)
     species_neutral = Array{species_parameters_mutable,1}(undef,n_neutral_species)

From ad71667adbfc6bba8d4ae85fcf108851d6edfedb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 13:11:13 +0100
Subject: [PATCH 009/428] (Re-)add numerical dissipation for neutral pdf

Ensure consistency with previous 1D1V implementation.
---
 src/input_structs.jl         |   1 +
 src/moment_kinetics_input.jl |   6 +-
 src/numerical_dissipation.jl | 145 ++++++++++++++++++++++++++++++++++-
 src/time_advance.jl          |   8 ++
 4 files changed, 156 insertions(+), 4 deletions(-)

diff --git a/src/input_structs.jl b/src/input_structs.jl
index 3b47a8a8f..c5fbdc9d1 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -70,6 +70,7 @@ mutable struct advance_info
     manufactured_solns_test::Bool
     r_diffusion::Bool #flag to control how r bc is imposed when r diffusion terms are present
     vpa_diffusion::Bool #flag to control how vpa bc is imposed when vpa diffusion terms are present
+    vz_diffusion::Bool #flag to control how vz bc is imposed when vz diffusion terms are present
 end
 
 """
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 8761fa3a2..ad2026ddd 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -234,7 +234,7 @@ function mk_input(scan_input=Dict())
     vpa.fd_option = get(scan_input, "vpa_finite_difference_option", "third_order_upwind")
 
     num_diss_params = setup_numerical_dissipation(
-        get(scan_input, "numerical_dissipation", Dict{String,Any}()))
+        get(scan_input, "numerical_dissipation", Dict{String,Any}()), true)
 
     # overwrite some default parameters related to the vperp grid
     # ngrid is the number of grid points per element
@@ -313,8 +313,10 @@ function mk_input(scan_input=Dict())
         vzeta.discretization = get(scan_input, "vzeta_discretization", "chebyshev_pseudospectral")
     end
 
+    is_1V = (vperp.ngrid == vperp.nelement_global == 1 && vzeta.ngrid ==
+             vzeta.nelement_global == 1 && vr.ngrid == vr.nelement_global == 1)
     num_diss_params = setup_numerical_dissipation(
-        get(scan_input, "numerical_dissipation", Dict{String,Any}()))
+        get(scan_input, "numerical_dissipation", Dict{String,Any}()), is_1V)
     
     #########################################################################
     ########## end user inputs. do not modify following code! ###############
diff --git a/src/numerical_dissipation.jl b/src/numerical_dissipation.jl
index b98d9a78f..7b3093e59 100644
--- a/src/numerical_dissipation.jl
+++ b/src/numerical_dissipation.jl
@@ -17,13 +17,22 @@ Base.@kwdef struct numerical_dissipation_parameters
     vpa_boundary_buffer_damping_rate::mk_float = -1.0
     vpa_boundary_buffer_diffusion_coefficient::mk_float = -1.0
     vpa_dissipation_coefficient::mk_float = -1.0
+    vz_dissipation_coefficient::mk_float = -1.0
     z_dissipation_coefficient::mk_float = -1.0
     r_dissipation_coefficient::mk_float = -1.0
     moment_dissipation_coefficient::mk_float = -1.0
     force_minimum_pdf_value::Union{Nothing,mk_float} = nothing
 end
 
-function setup_numerical_dissipation(input_section::Dict)
+function setup_numerical_dissipation(input_section::Dict, is_1V)
+    if is_1V && "vpa_dissipation_coefficient" ∈ keys(input_section)
+        # Set default for vz_dissipation_coefficient the same as
+        # vpa_dissipation_coefficient for 1V case
+        input_section["vz_dissipation_coefficient"] =
+            get(input_section, "vz_dissipation_coefficient",
+                input_section["vpa_dissipation_coefficient"])
+    end
+
     input = Dict(Symbol(k)=>v for (k,v) in input_section)
 
     return numerical_dissipation_parameters(; input...)
@@ -231,13 +240,14 @@ vpa_dissipation_coefficient = 0.1
 """
 function vpa_dissipation!(f_out, f_in, vpa, spectral::T_spectral, dt,
         num_diss_params::numerical_dissipation_parameters) where T_spectral
-    begin_s_r_z_vperp_region()
 
     diffusion_coefficient = num_diss_params.vpa_dissipation_coefficient
     if diffusion_coefficient <= 0.0
         return nothing
     end
 
+    begin_s_r_z_vperp_region()
+
     # if T_spectral <: Bool
     #     # Scale diffusion coefficient like square of grid spacing, so convergence will
     #     # be second order accurate despite presence of numerical dissipation.
@@ -383,6 +393,137 @@ function r_dissipation!(f_out, f_in, r, r_spectral::T_spectral, dt,
     return nothing
 end
 
+"""
+Add diffusion in the vz direction to suppress oscillations for neutrals
+
+Disabled by default.
+
+The diffusion coefficient is set in the input TOML file by the parameter
+```
+[numerical_dissipation]
+vz_dissipation_coefficient = 0.1
+```
+"""
+function vz_dissipation_neutral!(f_out, f_in, vz, spectral::T_spectral, dt,
+        num_diss_params::numerical_dissipation_parameters) where T_spectral
+
+    diffusion_coefficient = num_diss_params.vz_dissipation_coefficient
+    if diffusion_coefficient <= 0.0
+        return nothing
+    end
+
+    begin_sn_r_z_vzeta_vr_region()
+
+    @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+        vz.scratch2 .= 1.0 # placeholder for Q in d / d vpa ( Q d f / d vpa)
+        @views second_derivative!(vz.scratch, f_in[:,ivr,ivzeta,iz,ir,isn], vz.scratch2, vz, spectral)
+        @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch
+    end
+
+    return nothing
+end
+
+"""
+Add diffusion in the z direction to suppress oscillations for neutrals
+
+Disabled by default.
+
+The diffusion coefficient is set in the input TOML file by the parameter
+```
+[numerical_dissipation]
+z_dissipation_coefficient = 0.1
+```
+
+Note that the current distributed-memory compatible
+implementation does not impose a penalisation term
+on internal or external element boundaries
+
+"""
+function z_dissipation_neutral!(f_out, f_in, z, z_spectral::T_spectral, dt,
+        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+
+    diffusion_coefficient = num_diss_params.z_dissipation_coefficient
+    if diffusion_coefficient <= 0.0
+        return nothing
+    end
+
+    begin_sn_r_vzeta_vr_vz_region()
+
+    # calculate d / d z ( Q d f / d z ) using distributed memory compatible routines
+    # first compute d f / d z using centred reconciliation and place in dummy array #1
+    derivative_z!(scratch_dummy.buffer_vpavperpzrs_1, f_in,
+                  scratch_dummy.buffer_vpavperprs_1, scratch_dummy.buffer_vpavperprs_2,
+                  scratch_dummy.buffer_vpavperprs_3,scratch_dummy.buffer_vpavperprs_4,
+                  z_spectral,z)
+    # form Q d f / d r and place in dummy array #2
+    @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+        Q = 1.0 # placeholder for geometrical or velocity space dependent metric coefficient
+        @. scratch_dummy.buffer_vpavperpzrs_2[ivz,ivr,ivzeta,:,ir,isn] =  Q * scratch_dummy.buffer_vpavperpzrs_1[ivz,ivr,ivzeta,:,ir,isn]
+    end
+    # compute d / d z ( Q d f / d z ) using centred reconciliation and place in dummy array #1
+    derivative_z!(scratch_dummy.buffer_vpavperpzrs_1, scratch_dummy.buffer_vpavperpzrs_2,
+                  scratch_dummy.buffer_vpavperprs_1, scratch_dummy.buffer_vpavperprs_2,
+                  scratch_dummy.buffer_vpavperprs_3,scratch_dummy.buffer_vpavperprs_4,
+                  z_spectral,z)
+    # advance f due to diffusion_coefficient * d / d z ( Q d f / d z )
+    @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+        @views @. f_out[ivz,ivr,ivzeta,:,ir,isn] += dt * diffusion_coefficient * scratch_dummy.buffer_vpavperpzrs_1[ivz,ivr,ivzeta,:,ir,isn]
+    end
+
+    return nothing
+end
+
+"""
+Add diffusion in the r direction to suppress oscillations for neutrals
+
+Disabled by default.
+
+The diffusion coefficient is set in the input TOML file by the parameter
+```
+[numerical_dissipation]
+r_dissipation_coefficient = 0.1
+
+```
+
+Note that the current distributed-memory compatible
+implementation does not impose a penalisation term
+on internal or external element boundaries
+
+"""
+function r_dissipation_neutral!(f_out, f_in, r, r_spectral::T_spectral, dt,
+        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+
+    diffusion_coefficient = num_diss_params.r_dissipation_coefficient
+    if diffusion_coefficient <= 0.0
+        return nothing
+    end
+
+    begin_sn_z_vzeta_vr_nz_region()
+
+    # calculate d / d r ( Q d f / d r ) using distributed memory compatible routines
+    # first compute d f / d r using centred reconciliation and place in dummy array #1
+    derivative_r!(scratch_dummy.buffer_vpavperpzrs_1, f_in,
+                  scratch_dummy.buffer_vpavperpzs_1, scratch_dummy.buffer_vpavperpzs_2,
+                  scratch_dummy.buffer_vpavperpzs_3,scratch_dummy.buffer_vpavperpzs_4,
+                  r_spectral,r)
+    # form Q d f / d r and place in dummy array #2
+    @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
+        Q = 1.0 # placeholder for geometrical or velocity space dependent metric coefficient
+        @. scratch_dummy.buffer_vpavperpzrs_2[ivz,ivr,ivzeta,iz,:,isn] =  Q * scratch_dummy.buffer_vpavperpzrs_1[ivz,ivr,ivzeta,iz,:,isn]
+    end
+    # compute d / d r ( Q d f / d r ) using centred reconciliation and place in dummy array #1
+    derivative_r!(scratch_dummy.buffer_vpavperpzrs_1, scratch_dummy.buffer_vpavperpzrs_2,
+                  scratch_dummy.buffer_vpavperpzs_1, scratch_dummy.buffer_vpavperpzs_2,
+                  scratch_dummy.buffer_vpavperpzs_3,scratch_dummy.buffer_vpavperpzs_4,
+                  r_spectral,r)
+    # advance f due to diffusion_coefficient * d / d r ( Q d f / d r )
+    @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
+        @views @. f_out[ivz,ivr,ivzeta,iz,:,isn] += dt * diffusion_coefficient * scratch_dummy.buffer_vpavperpzrs_1[ivz,ivr,ivzeta,iz,:,isn]
+    end
+
+    return nothing
+end
+
 """
     force_minimum_pdf_value!(f, num_diss_paras::numerical_dissipation_parameters)
 
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 7313d6e2b..b2ccfa7f0 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -41,6 +41,8 @@ using ..ionization: ionization_collisions_1V!, ionization_collisions_3V!, consta
 using ..numerical_dissipation: vpa_boundary_buffer_decay!,
                                vpa_boundary_buffer_diffusion!, vpa_dissipation!,
                                z_dissipation!, r_dissipation!,
+                               vz_dissipation_neutral!, z_dissipation_neutral!,
+                               r_dissipation_neutral!,
                                vpa_boundary_force_decreasing!, force_minimum_pdf_value!,
                                force_minimum_pdf_value_neutral!
 using ..source_terms: source_terms!, source_terms_neutral!, source_terms_manufactured!
@@ -1623,6 +1625,12 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                        num_diss_params, scratch_dummy)
         r_dissipation!(fvec_out.pdf, fvec_in.pdf, r, r_spectral, dt,
                        num_diss_params, scratch_dummy)
+        vz_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, vz,
+                                vz_spectral, dt, num_diss_params)
+        z_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, z, z_spectral,
+                               dt, num_diss_params, scratch_dummy)
+        r_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, r, r_spectral,
+                               dt, num_diss_params, scratch_dummy)
     end
     # End of advance for distribution function
 

From e27c0f04bf2f4e6180fa8060382c5fbf0b948884 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 13:12:53 +0100
Subject: [PATCH 010/428] Fix merge of advance.*_diffusion flags, 'zero'
 boundary conditions

As implemented in tho 2D2V branch, the 'zero' boundary condition should
apply a Dirichlet boundary condition only at the upwind boundary when
there is advection but no diffusion, or at both boundaries when there is
diffusion. This commit enforces this behaviour in the merged branch, and
applies it consistently to all r, vpa, and vz diffusion. The z direction
is not included as if there is a wall boundary the outflow is very
strong, and we definitely do not want to impose a Dirichlet boundary
condition there.
---
 src/initial_conditions.jl | 43 ++++++++++++++++++++++-----------------
 src/time_advance.jl       | 38 ++++++++++++++++++++++++----------
 2 files changed, 51 insertions(+), 30 deletions(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 0595f4dfe..322266242 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -937,14 +937,15 @@ also enforce boundary conditions in z on all separately evolved velocity space m
 """
 function enforce_boundary_conditions!(f, f_r_bc, density, upar, ppar, moments, vpa_bc,
         z_bc, r_bc, vpa, vperp, z, r, vpa_adv, z_adv, r_adv, composition, scratch_dummy,
-        r_diffusion)
+        r_diffusion, vpa_diffusion)
 
     begin_s_r_z_vperp_region()
     @loop_s_r_z_vperp is ir iz ivperp begin
         # enforce the vpa BC
         # use that adv.speed independent of vpa
         @views enforce_v_boundary_condition_local!(f[:,ivperp,iz,ir,is], vpa_bc,
-                                                   vpa_adv[is].speed[:,ivperp,iz,ir])
+                                                   vpa_adv[is].speed[:,ivperp,iz,ir],
+                                                   vpa_diffusion)
     end
     begin_s_r_vperp_vpa_region()
     # enforce the z BC on the evolved velocity space moments of the pdf
@@ -962,10 +963,10 @@ end
 
 function enforce_boundary_conditions!(fvec_out::scratch_pdf, moments, f_r_bc, vpa_bc,
         z_bc, r_bc, vpa, vperp, z, r, vpa_adv, z_adv, r_adv, composition, scratch_dummy,
-        r_diffusion)
+        r_diffusion, vpa_diffusion)
     enforce_boundary_conditions!(fvec_out.pdf, f_r_bc, fvec_out.density, fvec_out.upar,
         fvec_out.ppar, moments, vpa_bc, z_bc, r_bc, vpa, vperp, z, r, vpa_adv, z_adv,
-        r_adv, composition, scratch_dummy, r_diffusion)
+        r_adv, composition, scratch_dummy, r_diffusion, vpa_diffusion)
 end
 
 """
@@ -1007,11 +1008,11 @@ function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc:
         # impose bc on both sides of the domain to accomodate a diffusion operator d^2 / d r^2
         @loop_s_z_vperp_vpa is iz ivperp ivpa begin
             ir = 1 # r = -L/2 -- check that the point is on lowest rank
-            if (adv[is].speed[ir,ivpa,ivperp,iz] > zero || r_diffusion) && r.irank == 0
+            if r.irank == 0 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] > zero)
                 f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,1,is]
             end
             ir = r.n # r = L/2 -- check that the point is on highest rank
-            if (adv[is].speed[ir,ivpa,ivperp,iz] < -zero || r_diffusion) && r.irank == r.nrank - 1
+            if r.irank == r.nrank - 1 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] < -zero)
                 f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,end,is]
             end
         end
@@ -1081,7 +1082,7 @@ enforce boundary conditions on neutral particle distribution function
 function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
         boundary_distributions, density_neutral, uz_neutral, pz_neutral, moments,
         density_ion, upar_ion, Er, r_adv, z_adv, vzeta_adv, vr_adv, vz_adv, r, z, vzeta,
-        vr, vz, composition, geometry, scratch_dummy)
+        vr, vz, composition, geometry, scratch_dummy, r_diffusion, vz_diffusion)
 
     if vzeta.n_global > 1 && vzeta.bc != "none"
         begin_sn_r_z_vr_vz_region()
@@ -1089,7 +1090,8 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
             # enforce the vz BC
             @views enforce_v_boundary_condition_local!(f_neutral[ivz,ivr,:,iz,ir,isn],
                                                        vzeta.bc,
-                                                       vzeta_adv[isn].speed[ivz,ivr,:,iz,ir])
+                                                       vzeta_adv[isn].speed[ivz,ivr,:,iz,ir],
+                                                       false)
         end
     end
     if vr.n_global > 1 && vr.bc != "none"
@@ -1098,7 +1100,8 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
             # enforce the vz BC
             @views enforce_v_boundary_condition_local!(f_neutral[ivz,:,ivzeta,iz,ir,isn],
                                                        vr.bc,
-                                                       vr_adv[isn].speed[ivz,:,ivzeta,iz,ir])
+                                                       vr_adv[isn].speed[ivz,:,ivzeta,iz,ir],
+                                                       false)
         end
     end
     if vz.n_global > 1 && vz.bc != "none"
@@ -1107,7 +1110,8 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
             # enforce the vz BC
             @views enforce_v_boundary_condition_local!(f_neutral[:,ivr,ivzeta,iz,ir,isn],
                                                        vz.bc,
-                                                       vz_adv[isn].speed[:,ivr,ivzeta,iz,ir])
+                                                       vz_adv[isn].speed[:,ivr,ivzeta,iz,ir],
+                                                       vz_diffusion)
         end
     end
     # f_initial contains the initial condition for enforcing a fixed-boundary-value condition
@@ -1122,7 +1126,7 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
                                     r_adv, vz, vr, vzeta, z, r, composition,
                                     scratch_dummy.buffer_vzvrvzetazsn_1, scratch_dummy.buffer_vzvrvzetazsn_2,
                                     scratch_dummy.buffer_vzvrvzetazsn_3, scratch_dummy.buffer_vzvrvzetazsn_4,
-                                    scratch_dummy.buffer_vzvrvzetazrsn_1)
+                                    scratch_dummy.buffer_vzvrvzetazrsn_1, r_diffusion)
     end
 end
 
@@ -1130,7 +1134,7 @@ function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6},
         f_r_bc::AbstractArray{mk_float,6}, adv::T, vz, vr, vzeta, z, r, composition,
         end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5},
         buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5},
-        buffer_dfn::AbstractArray{mk_float,6}) where T #f_initial,
+        buffer_dfn::AbstractArray{mk_float,6}, r_diffusion) where T #f_initial,
 
     bc = r.bc
     nr = r.n
@@ -1162,12 +1166,12 @@ function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6},
         @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
             ir = 1 # r = -L/2
             # incoming particles and on lowest rank
-            if adv[isn].speed[ir,ivz,ivr,ivzeta,iz] > zero && r.irank == 0
+            if r.irank == 0 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] > zero)
                 f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,1,isn]
             end
             ir = nr # r = L/2
             # incoming particles and on highest rank
-            if adv[isn].speed[ir,ivz,ivr,ivzeta,iz] < -zero && r.irank == r.nrank - 1
+            if r.irank == r.nrank - 1 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] < -zero)
                 f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,end,isn]
             end
         end
@@ -1855,25 +1859,26 @@ end
 impose the prescribed vpa boundary condition on f
 at every z grid point
 """
-function enforce_vpa_boundary_condition!(f, bc, src)
+function enforce_vpa_boundary_condition!(f, bc, src, v_diffusion)
     nz = size(f,2)
     nr = size(f,3)
     for ir ∈ 1:nr
         for iz ∈ 1:nz
-            enforce_v_boundary_condition_local!(view(f,:,iz,ir), bc, src.speed[:,iz,ir])
+            enforce_v_boundary_condition_local!(view(f,:,iz,ir), bc, src.speed[:,iz,ir],
+                                                v_diffusion)
         end
     end
 end
 
 """
 """
-function enforce_v_boundary_condition_local!(f, bc, speed)
+function enforce_v_boundary_condition_local!(f, bc, speed, v_diffusion)
     if bc == "zero"
-        if speed[1] > 0.0
+        if v_diffusion || speed[1] > 0.0
             # 'upwind' boundary
             f[1] = 0.0
         end
-        if speed[end] < 0.0
+        if v_diffusion || speed[end] < 0.0
             # 'upwind' boundary
             f[end] = 0.0
         end
diff --git a/src/time_advance.jl b/src/time_advance.jl
index b2ccfa7f0..06f37388d 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -171,8 +171,8 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     # indicate which parts of the equations are to be advanced concurrently.
     # if no splitting of operators, all terms advanced concurrently;
     # else, will advance one term at a time.
-    advance = setup_advance_flags(moments, composition, t_input, collisions, rk_coefs, r,
-                                  vperp, vpa, vzeta, vr, vz)
+    advance = setup_advance_flags(moments, composition, t_input, collisions,
+                                  num_diss_params, rk_coefs, r, vperp, vpa, vzeta, vr, vz)
 
 
     begin_serial_region()
@@ -308,7 +308,7 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             pdf.charged.norm, boundary_distributions.pdf_rboundary_charged,
             moments.charged.dens, moments.charged.upar, moments.charged.ppar, moments,
             vpa.bc, z.bc, r.bc, vpa, vperp, z, r, vpa_advect, z_advect, r_advect,
-            composition, scratch_dummy, advance.r_diffusion)
+            composition, scratch_dummy, advance.r_diffusion, advance.vpa_diffusion)
         # Ensure normalised pdf exactly obeys integral constraints if evolving moments
         begin_s_r_z_region()
         @loop_s_r_z is ir iz begin
@@ -328,7 +328,8 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
                 moments.neutral.dens, moments.neutral.uz, moments.neutral.pz, moments,
                 moments.charged.dens, moments.charged.upar, fields.Er, neutral_r_advect,
                 neutral_z_advect, nothing, nothing, neutral_vz_advect, r, z, vzeta, vr,
-                vz, composition, geometry, scratch_dummy)
+                vz, composition, geometry, scratch_dummy, advance.r_diffusion,
+                advance.vz_diffusion)
             begin_sn_r_z_region()
             @loop_sn_r_z isn ir iz begin
                 @views hard_force_moment_constraints_neutral!(
@@ -379,8 +380,8 @@ indicate which parts of the equations are to be advanced concurrently.
 if no splitting of operators, all terms advanced concurrently;
 else, will advance one term at a time.
 """
-function setup_advance_flags(moments, composition, t_input, collisions, rk_coefs, r,
-                             vperp, vpa, vzeta, vr, vz)
+function setup_advance_flags(moments, composition, t_input, collisions, num_diss_params,
+                             rk_coefs, r, vperp, vpa, vzeta, vr, vz)
     # default is not to concurrently advance different operators
     advance_vpa_advection = false
     advance_z_advection = false
@@ -404,6 +405,7 @@ function setup_advance_flags(moments, composition, t_input, collisions, rk_coefs
     advance_neutral_energy = false
     r_diffusion = false
     vpa_diffusion = false
+    vz_diffusion = false
     # all advance flags remain false if using operator-splitting
     # otherwise, check to see if the flags need to be set to true
     if !t_input.split_operators
@@ -445,6 +447,12 @@ function setup_advance_flags(moments, composition, t_input, collisions, rk_coefs
                     advance_ionization_1V = true
                 elseif vperp.n > 1 && vr.n > 1 && vzeta.n > 1
                     advance_ionization = true
+                else
+                    error("If any perpendicular velocity has length>1 they all must. "
+                          * "If all perpendicular velocities have length=1, then vpa and "
+                          * "vz should be the same.\n"
+                          * "vperp.n=$(vperp.n), vr.n=$(vr.n), vzeta.n=$(vzeta.n), "
+                          * "vpa.n=$(vpa.n), vz.n=$(vz.n)")
                 end
             end
         end
@@ -486,6 +494,12 @@ function setup_advance_flags(moments, composition, t_input, collisions, rk_coefs
                 advance_neutral_energy = true
             end
         end
+
+        # flag to determine if a d^2/dr^2 operator is present
+        r_diffusion = (advance_numerical_dissipation && num_diss_params.r_dissipation_coefficient > 0.0)
+        # flag to determine if a d^2/dvpa^2 operator is present
+        vpa_diffusion = (advance_numerical_dissipation && num_diss_params.vpa_dissipation_coefficient > 0.0)
+        vz_diffusion = (advance_numerical_dissipation && num_diss_params.vz_dissipation_coefficient > 0.0)
     end
 
     manufactured_solns_test = t_input.use_manufactured_solns_for_advance
@@ -499,7 +513,7 @@ function setup_advance_flags(moments, composition, t_input, collisions, rk_coefs
                         advance_energy, advance_neutral_sources,
                         advance_neutral_continuity, advance_neutral_force_balance,
                         advance_neutral_energy, rk_coefs, manufactured_solns_test,
-                        r_diffusion, vpa_diffusion)
+                        r_diffusion, vpa_diffusion, vz_diffusion)
 end
 
 function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies_ion,nspecies_neutral)
@@ -1082,7 +1096,8 @@ or update them by taking the appropriate velocity moment of the evolved pdf
 """
 function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, vr, vzeta,
                     vpa, vperp, z, r, advect_objects, rk_coefs, istage, composition,
-                    geometry, num_diss_params, z_spectral, r_spectral, scratch_dummy)
+                    geometry, num_diss_params, z_spectral, r_spectral, advance,
+                    scratch_dummy)
     begin_s_r_z_region()
 
     new_scratch = scratch[istage+1]
@@ -1116,7 +1131,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     enforce_boundary_conditions!(new_scratch, moments,
         boundary_distributions.pdf_rboundary_charged, vpa.bc, z.bc, r.bc, vpa, vperp, z,
         r, vpa_advect, z_advect, r_advect, composition, scratch_dummy,
-        num_diss_params.r_dissipation_coefficient>0.0)
+        advance.r_diffusion, advance.vpa_diffusion)
 
     if moments.evolve_density && moments.enforce_conservation
         begin_s_r_z_region()
@@ -1188,7 +1203,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
             boundary_distributions, new_scratch.density_neutral, new_scratch.uz_neutral,
             new_scratch.pz_neutral, moments, new_scratch.density, new_scratch.upar,
             fields.Er, neutral_r_advect, neutral_z_advect, nothing, nothing,
-            neutral_vz_advect, r, z, vzeta, vr, vz, composition, geometry, scratch_dummy)
+            neutral_vz_advect, r, z, vzeta, vr, vz, composition, geometry, scratch_dummy,
+            advance.r_diffusion, advance.vz_diffusion)
 
         if moments.evolve_density && moments.enforce_conservation
             begin_sn_r_z_region()
@@ -1435,7 +1451,7 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
                           vzeta, vpa, vperp, z, r, advect_objects,
                           advance.rk_coefs[:,istage], istage, composition, geometry,
                           num_diss_params, spectral_objects.z_spectral,
-                          spectral_objects.r_spectral, scratch_dummy)
+                          spectral_objects.r_spectral, advance, scratch_dummy)
     end
 
     istage = n_rk_stages+1

From 8af54d01da4ef074798c5a7a6fbfaa856f4ab8ed Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 13:32:10 +0100
Subject: [PATCH 011/428] Remove 'unnorm' version of pdf

It is not needed anywhere, so removing it reduces the amount of code and
gets rid of some unnecessary array copies and computations.
---
 src/initial_conditions.jl | 63 ++++++++++++---------------
 src/time_advance.jl       | 90 +--------------------------------------
 2 files changed, 29 insertions(+), 124 deletions(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 322266242..f33c92952 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -39,7 +39,6 @@ using ..manufactured_solns: manufactured_solutions
 """
 struct pdf_substruct{n_distribution}
     norm::MPISharedArray{mk_float,n_distribution}
-    unnorm::MPISharedArray{mk_float,n_distribution}
     buffer::MPISharedArray{mk_float,n_distribution} # for collision operator terms when pdfs must be interpolated onto different velocity space grids
 end
 
@@ -123,14 +122,12 @@ Allocate arrays for pdfs
 function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     # allocate pdf arrays
     pdf_charged_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_ion_species)
-    pdf_charged_unnorm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_ion_species)
     pdf_charged_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
-    pdf_neutral_unnorm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
 
-    return pdf_struct(pdf_substruct(pdf_charged_norm, pdf_charged_unnorm, pdf_charged_buffer),
-                    pdf_substruct(pdf_neutral_norm, pdf_neutral_unnorm, pdf_neutral_buffer))
+    return pdf_struct(pdf_substruct(pdf_charged_norm, pdf_charged_buffer),
+                      pdf_substruct(pdf_neutral_norm, pdf_neutral_buffer))
 
 end
 
@@ -177,7 +174,7 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         # create and initialise the normalised particle distribution function (pdf)
         # such that ∫dwpa pdf.norm = 1, ∫dwpa wpa * pdf.norm = 0, and ∫dwpa wpa^2 * pdf.norm = 1/2
         # note that wpa = vpa - upar, unless moments.evolve_ppar = true, in which case wpa = (vpa - upar)/vth
-        # the definition of pdf.norm changes accordingly from pdf.unnorm / density to pdf.unnorm * vth / density
+        # the definition of pdf.norm changes accordingly from pdf_unnorm / density to pdf_unnorm * vth / density
         # when evolve_ppar = true.
         initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
                         vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
@@ -191,17 +188,17 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         if n_neutral_species > 0
             update_neutral_qz!(moments.neutral.qz, moments.neutral.qz_updated,
                                moments.neutral.dens, moments.neutral.uz,
-                               moments.neutral.vth, pdf.neutral.unnorm, vz, vr, vzeta, z,
+                               moments.neutral.vth, pdf.neutral.norm, vz, vr, vzeta, z,
                                r, composition, moments.evolve_density,
                                moments.evolve_upar, moments.evolve_ppar)
             update_neutral_pz!(moments.neutral.pz, moments.neutral.pz_updated,
                                moments.neutral.dens, moments.neutral.uz,
-                               pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition,
+                               pdf.neutral.norm, vz, vr, vzeta, z, r, composition,
                                moments.evolve_density, moments.evolve_upar)
             update_neutral_pr!(moments.neutral.pr, moments.neutral.pr_updated,
-                               pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
+                               pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
             update_neutral_pzeta!(moments.neutral.pzeta, moments.neutral.pzeta_updated,
-                                  pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
+                                  pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
         end
     end
 
@@ -241,13 +238,11 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
             @views wall_flux_L[ir,is] = moments.charged.dens[end,ir,is] *
                                         moments.charged.upar[end,ir,is]
 
-            pdf.charged.unnorm[:,:,:,ir,is] .= pdf.charged.norm[:,:,:,ir,is]
             @loop_z iz begin
-                @. pdf.charged.unnorm[:,:,iz,ir,is] *= moments.charged.dens[iz,ir,is]
                 if moments.evolve_ppar
                     @. pdf.charged.norm[:,:,iz,ir,is] *= moments.charged.vth[iz,ir,is]
                 elseif moments.evolve_density == false
-                    @. pdf.charged.norm[:,:,iz,ir,is] = pdf.charged.unnorm[:,:,iz,ir,is]
+                    @. pdf.charged.norm[:,:,iz,ir,is] *= moments.charged.dens[iz,ir,is]
                 end
             end
         end
@@ -268,13 +263,11 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
                 wall_flux_0[ir,min(isn,composition.n_ion_species)],
                 wall_flux_L[ir,min(isn,composition.n_ion_species)])
 
-            pdf.neutral.unnorm[:,:,:,:,ir,isn] .= pdf.neutral.norm[:,:,:,:,ir,isn]
             @loop_z iz begin
-                @. pdf.neutral.unnorm[:,:,:,iz,ir,isn] *= moments.neutral.dens[iz,ir,isn]
                 if moments.evolve_ppar
                     @. pdf.neutral.norm[:,:,:,iz,ir,isn] *= moments.neutral.vth[iz,ir,isn]
                 elseif moments.evolve_density == false
-                    @. pdf.neutral.norm[:,:,:,iz,ir,isn] = pdf.neutral.unnorm[:,:,:,iz,ir,isn]
+                    @. pdf.neutral.norm[:,:,:,iz,ir,isn] *= moments.neutral.dens[iz,ir,isn]
                 end
             end
         end
@@ -759,16 +752,15 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     @loop_s_r_z is ir iz begin
         moments.charged.dens[iz,ir,is] = densi_func(z.grid[iz],r.grid[ir],0.0)
         @loop_vperp_vpa ivperp ivpa begin
-            pdf.charged.unnorm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
-            pdf.charged.norm[ivpa,ivperp,iz,ir,is] = pdf.charged.unnorm[ivpa,ivperp,iz,ir,is]
+            pdf.charged.norm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
         end
     end
     # update upar, ppar, qpar, vth consistent with manufactured solns
-    update_density!(moments.charged.dens, pdf.charged.unnorm, vpa, vperp, z, r, composition)
-    update_qpar!(moments.charged.qpar, pdf.charged.unnorm, vpa, vperp, z, r, composition)
-    update_ppar!(moments.charged.ppar, pdf.charged.unnorm, vpa, vperp, z, r, composition)
+    update_density!(moments.charged.dens, pdf.charged.norm, vpa, vperp, z, r, composition)
+    update_qpar!(moments.charged.qpar, pdf.charged.norm, vpa, vperp, z, r, composition)
+    update_ppar!(moments.charged.ppar, pdf.charged.norm, vpa, vperp, z, r, composition)
     # get particle flux
-    update_upar!(moments.charged.upar, pdf.charged.unnorm, vpa, vperp, z, r, composition)
+    update_upar!(moments.charged.upar, pdf.charged.norm, vpa, vperp, z, r, composition)
     # convert from particle particle flux to parallel flow
     begin_s_r_z_region()
     @loop_s_r_z is ir iz begin
@@ -782,16 +774,15 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
         @loop_sn_r_z isn ir iz begin
             moments.neutral.dens[iz,ir,isn] = densn_func(z.grid[iz],r.grid[ir],0.0)
             @loop_vzeta_vr_vz ivzeta ivr ivz begin
-                pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn] = dfnn_func(vz.grid[ivz],vr.grid[ivr],vzeta.grid[ivzeta],z.grid[iz],r.grid[ir],0.0)
-                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn]
+                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = dfnn_func(vz.grid[ivz],vr.grid[ivr],vzeta.grid[ivzeta],z.grid[iz],r.grid[ir],0.0)
             end
         end
         # get consistent moments with manufactured solutions
-        update_neutral_density!(moments.neutral.dens, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_qz!(moments.neutral.qz, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_pz!(moments.neutral.pz, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_pr!(moments.neutral.pr, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_pzeta!(moments.neutral.pzeta, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
+        update_neutral_density!(moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_qz!(moments.neutral.qz, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_pz!(moments.neutral.pz, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_pr!(moments.neutral.pr, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_pzeta!(moments.neutral.pzeta, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
         #update ptot (isotropic pressure)
         if r.n > 1 #if 2D geometry
             begin_sn_r_z_region()
@@ -802,9 +793,9 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
             moments.neutral.ptot .= moments.neutral.pz
         end
         # nb bad naming convention uz -> n uz below
-        update_neutral_uz!(moments.neutral.uz, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_ur!(moments.neutral.ur, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
-        update_neutral_uzeta!(moments.neutral.uzeta, pdf.neutral.unnorm, vz, vr, vzeta, z, r, composition)
+        update_neutral_uz!(moments.neutral.uz, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_ur!(moments.neutral.ur, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+        update_neutral_uzeta!(moments.neutral.uzeta, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
         # now convert from particle particle flux to parallel flow
         begin_sn_r_z_region()
         @loop_sn_r_z isn ir iz begin
@@ -890,14 +881,14 @@ function init_rboundary_pdfs!(rboundary_charged, rboundary_neutral, pdf::pdf_str
 
     begin_s_z_region() #do not parallelise r here
     @loop_s_z_vperp_vpa is iz ivperp ivpa begin
-        rboundary_charged[ivpa,ivperp,iz,1,is] = pdf.charged.unnorm[ivpa,ivperp,iz,1,is]
-        rboundary_charged[ivpa,ivperp,iz,end,is] = pdf.charged.unnorm[ivpa,ivperp,iz,end,is]
+        rboundary_charged[ivpa,ivperp,iz,1,is] = pdf.charged.norm[ivpa,ivperp,iz,1,is]
+        rboundary_charged[ivpa,ivperp,iz,end,is] = pdf.charged.norm[ivpa,ivperp,iz,end,is]
     end
     if n_neutral_species > 0
         begin_sn_z_region() #do not parallelise r here
         @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
-            rboundary_neutral[ivz,ivr,ivzeta,iz,1,isn] = pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,1,isn]
-            rboundary_neutral[ivz,ivr,ivzeta,iz,end,isn] = pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,end,isn]
+            rboundary_neutral[ivz,ivr,ivzeta,iz,1,isn] = pdf.neutral.norm[ivz,ivr,ivzeta,iz,1,isn]
+            rboundary_neutral[ivz,ivr,ivzeta,iz,end,isn] = pdf.neutral.norm[ivz,ivr,ivzeta,iz,end,isn]
         end
     end
     return rboundary_charged, rboundary_neutral
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 06f37388d..68351101a 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -357,9 +357,6 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             scratch[1].pz_neutral[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
         end
     end
-    # update unnormalised pdf and phi in case they were affected by applying boundary
-    # conditions or constraints to the pdf
-    update_pdf_unnorm!(pdf, moments, scratch[1].temp_z_s, composition)
 
     update_phi!(fields, scratch[1], z, r, composition, z_spectral, r_spectral,
                 scratch_dummy)
@@ -1066,19 +1063,16 @@ function time_advance_no_splitting!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa
            composition, collisions, geometry, boundary_distributions,
            num_diss_params, advance, scratch_dummy, manufactured_source_list, istep)
 
-    #pdf_in = pdf.unnorm #get input pdf values in case we wish to impose a constant-in-time boundary condition in r
-    #use pdf.norm for this fn for now.
-
     if t_input.n_rk_stages > 1
         ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase, z, r,
             moments, fields, spectral_objects, advect_objects,
             composition, collisions, geometry, boundary_distributions, num_diss_params,
-            advance,  scratch_dummy, manufactured_source_list, istep)#pdf_in,
+            advance,  scratch_dummy, manufactured_source_list, istep)
     else
         euler_time_advance!(scratch, scratch, pdf, fields, moments,
             advect_objects, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, t,
             t_input, spectral_objects, composition,
-            collisions, geometry, boundary_distributions, scratch_dummy, manufactured_source_list, advance, 1)#pdf_in,
+            collisions, geometry, boundary_distributions, scratch_dummy, manufactured_source_list, advance, 1)
         # NB: this must be broken -- scratch is updated in euler_time_advance!,
         # but not the pdf or moments. need to add update to these quantities here. Also
         # need to apply boundary conditions, possibly other things that are taken care
@@ -1169,9 +1163,6 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                   num_diss_params)
 
-    # update the 'true', un-normalized pdf
-    update_unnormalized_pdf!(pdf.charged.unnorm, new_scratch, moments)
-
     ##
     # update the neutral particle distribution and moments
     ##
@@ -1232,9 +1223,6 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 
         calculate_moment_derivatives_neutral!(moments, new_scratch, scratch_dummy, z,
                                               z_spectral, num_diss_params)
-
-        # update the 'true', un-normalized pdf
-        update_unnormalized_pdf_neutral!(pdf.neutral.unnorm, new_scratch, moments)
     end
 
     # update the electrostatic potential phi
@@ -1344,58 +1332,6 @@ function update_derived_moments_neutral!(new_scratch, moments, vz, vr, vzeta, z,
     end
 end
 
-"""
-update the 'true', un-normalized, charged pdf
-"""
-function update_unnormalized_pdf!(pdf_unnorm, new_scratch, moments)
-    # if no moments are evolved separately from the pdf, then the
-    # evolved pdf is the 'true', un-normalized pdf;
-    # initialize to this value and modify below if necessary
-    @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        pdf_unnorm[ivpa,ivperp,iz,ir,is] = new_scratch.pdf[ivpa,ivperp,iz,ir,is]
-    end
-    # if separately evolving the particle density, the evolved
-    # pdf is the 'true' pdf divided by the particle density
-    if moments.evolve_density
-        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-            pdf_unnorm[ivpa,ivperp,iz,ir,is] = new_scratch.pdf[ivpa,ivperp,iz,ir,is] * new_scratch.density[iz,ir,is]
-        end
-    end
-    # if separately evolving the parallel pressure, the evolved
-    # pdf is the 'true' pdf multiplied by the thermal speed
-    if moments.evolve_ppar
-        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-            pdf_unnorm[ivpa,ivperp,iz,ir,is] /= moments.charged.vth[iz,ir,is]
-        end
-    end
-end
-
-"""
-update the 'true', un-normalized, neutral pdf
-"""
-function update_unnormalized_pdf_neutral!(pdf_unnorm, new_scratch, moments)
-    # if no moments are evolved separately from the pdf, then the
-    # evolved pdf is the 'true', un-normalized pdf;
-    # initialize to this value and modify below if necessary
-    @loop_sn_r_z_vzeta_vr_vz is ir iz ivzeta ivr ivz begin
-        pdf_unnorm[ivz,ivr,ivzeta,iz,ir,is] = new_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,is]
-    end
-    # if separately evolving the particle density, the evolved
-    # pdf is the 'true' pdf divided by the particle density
-    if moments.evolve_density
-        @loop_sn_r_z_vzeta_vr_vz is ir iz ivzeta ivr ivz begin
-            pdf_unnorm[ivz,ivr,ivzeta,iz,ir,is] = new_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,is] * new_scratch.density_neutral[iz,ir,is]
-        end
-    end
-    # if separately evolving the parallel pressure, the evolved
-    # pdf is the 'true' pdf multiplied by the thermal speed
-    if moments.evolve_ppar
-        @loop_sn_r_z_vzeta_vr_vz is ir iz ivzeta ivr ivz begin
-            pdf_unnorm[ivz,ivr,ivzeta,iz,ir,is] /= moments.neutral.vth[iz,ir,is]
-        end
-    end
-end
-
 """
 """
 function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase, z, r,
@@ -1421,7 +1357,6 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         begin_sn_r_z_region()
         @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
             first_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn] = pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn]
-            # change norm -> unnorm if remove moment-based evolution?
         end
         @loop_sn_r_z isn ir iz begin
             first_scratch.density_neutral[iz,ir,isn] = moments.neutral.dens[iz,ir,isn]
@@ -1461,7 +1396,6 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
     final_scratch = scratch[istage]
     @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
         pdf.charged.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
-        # change norm -> unnorm if remove moment-based evolution?
     end
     @loop_s_r_z is ir iz begin
         moments.charged.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
@@ -1474,7 +1408,6 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         begin_sn_r_z_region(no_synchronize=true)
         @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
             pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = final_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn]
-            # change norm -> unnorm if remove moment-based evolution?
         end
         @loop_sn_r_z isn ir iz begin
             moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
@@ -1522,7 +1455,6 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         end
     end
 
-    update_pdf_unnorm!(pdf, moments, scratch[istage].temp_z_s, composition)
     return nothing
 end
 
@@ -1729,22 +1661,4 @@ function update_solution_vector!(evolved, moments, istage, composition, vpa, vpe
     return nothing
 end
 
-"""
-if separately evolving the density via the continuity equation,
-the evolved pdf has been normalised by the particle density
-undo this normalisation to get the true particle distribution function
-"""
-function update_pdf_unnorm!(pdf, moments, scratch, composition)
-    begin_s_r_z_vperp_vpa_region()
-    @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        pdf.charged.unnorm[ivpa,ivperp,iz,ir,is] = pdf.charged.norm[ivpa,ivperp,iz,ir,is]
-    end
-    if composition.n_neutral_species > 0
-        begin_sn_r_z_vzeta_vr_vz_region()
-        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
-            pdf.neutral.unnorm[ivz,ivr,ivzeta,iz,ir,isn] = pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn]
-        end
-    end
-end
-
 end

From 9b099f893815c145939038428ff8e8c8922be3cf Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 13:34:49 +0100
Subject: [PATCH 012/428] Make spectral_object_struct a templated struct

Should allow it to be properly concretely typed, which might be a slight
optimization.
---
 src/time_advance.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/time_advance.jl b/src/time_advance.jl
index 68351101a..8dc7f1cd2 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -139,14 +139,14 @@ end
 
 # consider changing code structure so that
 # we can avoid arbitrary types below?
-struct spectral_object_struct
-    vz_spectral::T where T
-    vr_spectral::T where T
-    vzeta_spectral::T where T
-    vpa_spectral::T where T
-    vperp_spectral::T where T
-    z_spectral::T where T
-    r_spectral::T where T
+struct spectral_object_struct{Tvz,Tvr,Tvzeta,Tvpa,Tvperp,Tz,Tr}
+    vz_spectral::Tvz
+    vr_spectral::Tvr
+    vzeta_spectral::Tvzeta
+    vpa_spectral::Tvpa
+    vperp_spectral::Tvperp
+    z_spectral::Tz
+    r_spectral::Tr
 end
 
 """

From dd8def50dd81bcf7b39f72c1717c7f40368e82b3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 13:38:20 +0100
Subject: [PATCH 013/428] Remove some unnecessary `mutable` keywords from
 structs

Might be a slight optimization.
---
 src/time_advance.jl     | 2 +-
 src/velocity_moments.jl | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/time_advance.jl b/src/time_advance.jl
index 8dc7f1cd2..b460b608e 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -58,7 +58,7 @@ using Dates
 using Plots
 using ..post_processing: draw_v_parallel_zero!
 
-mutable struct scratch_dummy_arrays
+struct scratch_dummy_arrays
     dummy_sr::Array{mk_float,2}
     dummy_vpavperp::Array{mk_float,2}
     dummy_zrs::MPISharedArray{mk_float,3}
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index ea06fbe4c..ae3fe8c29 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -36,7 +36,7 @@ using ..looping
 
 """
 """
-mutable struct moments_charged_substruct
+struct moments_charged_substruct
     # this is the particle density
     dens::MPISharedArray{mk_float,3}
     # flag that keeps track of if the density needs updating before use
@@ -96,7 +96,7 @@ end
 
 """
 """
-mutable struct moments_neutral_substruct
+struct moments_neutral_substruct
     # this is the particle density
     dens::MPISharedArray{mk_float,3}
     # flag that keeps track of if the density needs updating before use

From 463351ca30d1d6b169ee3a2bec5cc415847a33a5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 May 2023 14:06:28 +0100
Subject: [PATCH 014/428] Write moments to 'dfns' output files

Allows these files to be used as checkpoints for restarting, and
simplifies post-processing as there is no worry about synchronizing
moments and dfns from separate files.
---
 src/file_io.jl         |  52 ++++++------
 src/file_io_hdf5.jl    |  11 +++
 src/file_io_netcdf.jl  |  11 +++
 src/load_data.jl       |  27 +-----
 src/moment_kinetics.jl |   6 +-
 src/post_processing.jl | 189 +++++++++++++++++++++++------------------
 src/time_advance.jl    |   4 +-
 7 files changed, 163 insertions(+), 137 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index 718f44bae..808e6c1a4 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -74,11 +74,9 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn}
 structure containing the data/metadata needed for binary file i/o
 distribution function data only
 """
-struct io_dfns_info{Tfile, Ttime, Tfi, Tfn}
-     # file identifier for the binary file to which data is written
+struct io_dfns_info{Tfile, Tfi, Tfn, Tmoments}
+    # file identifier for the binary file to which data is written
     fid::Tfile
-    # handle for the time variable
-    time::Ttime
     # handle for the charged species distribution function variable
     f::Tfi
     # handle for the neutral species distribution function variable
@@ -86,6 +84,9 @@ struct io_dfns_info{Tfile, Ttime, Tfi, Tfn}
 
     # Use parallel I/O?
     parallel_io::Bool
+
+    # Handles for moment variables
+    io_moments::Tmoments
 end
 
 """
@@ -140,6 +141,11 @@ function setup_file_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition, c
     return nothing, nothing, nothing
 end
 
+"""
+Get a (sub-)group from a file or group
+"""
+function get_group() end
+
 """
     write_single_value!(file_or_group, name, value; description=nothing)
 
@@ -438,18 +444,11 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz,
                                        n_ion_species, n_neutral_species, parallel_io)
 
     @serial_region begin
-        if !haskey(fid, "dynamic_data")
-            dynamic = create_io_group(fid, "dynamic_data",
-                                      description="time evolving variables")
-        else
-            dynamic = fid["dynamic_data"]
-        end
+        io_moments = define_dynamic_moment_variables!(fid, n_ion_species,
+                                                      n_neutral_species, r, z,
+                                                      parallel_io)
 
-        if !haskey(dynamic, "time")
-            io_time = create_dynamic_variable!(dynamic, "time", mk_float;
-                                               parallel_io=parallel_io,
-                                               description="simulation time")
-        end
+        dynamic = get_group(fid, "dynamic_data")
 
         # io_f is the handle for the ion pdf
         io_f = create_dynamic_variable!(dynamic, "f", mk_float, vpa, vperp, z, r;
@@ -463,7 +462,7 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz,
                                                 parallel_io=parallel_io,
                                                 description="neutral species distribution function")
 
-        return io_dfns_info(fid, io_time, io_f, io_f_neutral, parallel_io)
+        return io_dfns_info(fid, io_f, io_f_neutral, parallel_io, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -611,13 +610,16 @@ end
 """
 write time-dependent distribution function data to the binary output file
 """
-function write_dfns_data_to_binary(ff, ff_neutral, t, n_ion_species, n_neutral_species,
-                                   io_dfns, t_idx, r, z, vperp, vpa, vzeta, vr, vz)
+function write_dfns_data_to_binary(ff, ff_neutral, moments, fields, t, n_ion_species,
+                                   n_neutral_species, io_dfns, t_idx, r, z, vperp, vpa,
+                                   vzeta, vr, vz)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        # add the time for this time slice to the hdf5 file
-        append_to_dynamic_var(io_dfns.time, t, t_idx)
+        # Write the moments for this time slice to the output file.
+        # This also updates the time.
+        write_moments_data_to_binary(moments, fields, t, n_ion_species, n_neutral_species,
+                                     io_dfns.io_moments, t_idx, r, z)
 
         # add the distribution function data at this time slice to the output file
         append_to_dynamic_var(io_dfns.f, ff, t_idx, vpa, vperp, z, r, n_ion_species)
@@ -725,13 +727,15 @@ end
     # Special versions when using DebugMPISharedArray to avoid implicit conversion to
     # Array, which is forbidden.
     function write_dfns_data_to_binary(ff::DebugMPISharedArray,
-            ff_neutral::DebugMPISharedArray, t, n_ion_species, n_neutral_species, io_dfns,
-            t_idx, r, z, vperp, vpa, vzeta, vr, vz)
+            ff_neutral::DebugMPISharedArray, moments, t, n_ion_species, n_neutral_species,
+            io_dfns, t_idx, r, z, vperp, vpa, vzeta, vr, vz)
         @serial_region begin
             # Only read/write from first process in each 'block'
 
-            # add the time for this time slice to the hdf5 file
-            append_to_dynamic_var(io_dfns.time, t, t_idx)
+            # Write the moments for this time slice to the output file
+            # This also updates the time.
+            write_moments_data_to_binary(moments, fields, t, n_ion_species, n_neutral_species,
+                                         io_dfns.io_moments, t_idx, r, z)
 
             # add the distribution function data at this time slice to the output file
             append_to_dynamic_var(io_dfns.f, ff.data, t_idx, vpa, vperp, z, r, n_ion_species)
diff --git a/src/file_io_hdf5.jl b/src/file_io_hdf5.jl
index 7ac31fc5a..b79ce6dc4 100644
--- a/src/file_io_hdf5.jl
+++ b/src/file_io_hdf5.jl
@@ -49,6 +49,17 @@ function add_attribute!(var::HDF5.Dataset, name, value)
     attributes(var)[name] = value
 end
 
+function get_group(file_or_group::HDF5.H5DataStore, name::String)
+    # This overload deals with cases where fid is an HDF5 `File` or `Group` (`H5DataStore`
+    # is the abstract super-type for both
+    try
+        return file_or_group[name]
+    catch
+        println("An error occured while opening the $name group")
+        rethrow()
+    end
+end
+
 # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group
 function write_single_value!(file_or_group::HDF5.H5DataStore, name,
                              data::Union{Number,AbstractArray{T,N}},
diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 634d3dfe0..9571d11f6 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -30,6 +30,17 @@ function create_io_group(parent::NCDataset, name; description=nothing)
     return defGroup(parent, name, attrib=attributes)
 end
 
+function get_group(file_or_group::NCDataset, name::String)
+    # This overload deals with cases where fid is a NetCDF `Dataset` (which could be a
+    # file or a group).
+    try
+        return file_or_group.group[name]
+    catch
+        println("An error occured while opening the $name group")
+        rethrow()
+    end
+end
+
 function add_attribute!(file_or_group::NCDataset, name, value)
     file_or_group.attrib[name] = value
 end
diff --git a/src/load_data.jl b/src/load_data.jl
index 4bdbc11ac..8ebff4364 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -14,6 +14,8 @@ export load_block_data
 export load_rank_data
 export load_species_data
 
+using ..file_io: get_group
+
 using HDF5
 using NCDatasets
 
@@ -100,31 +102,6 @@ function load_variable(file_or_group::NCDataset, name::String)
     end
 end
 
-"""
-Get a (sub-)group from a file or group
-"""
-function get_group() end
-function get_group(file_or_group::HDF5.H5DataStore, name::String)
-    # This overload deals with cases where fid is an HDF5 `File` or `Group` (`H5DataStore`
-    # is the abstract super-type for both
-    try
-        return file_or_group[name]
-    catch
-        println("An error occured while opening the $name group")
-        rethrow()
-    end
-end
-function get_group(file_or_group::NCDataset, name::String)
-    # This overload deals with cases where fid is a NetCDF `Dataset` (which could be a
-    # file or a group).
-    try
-        return file_or_group.group[name]
-    catch
-        println("An error occured while opening the $name group")
-        rethrow()
-    end
-end
-
 """
 Load data for a coordinate
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 1572b2799..f61d0d821 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -342,9 +342,9 @@ function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
 
     write_moments_data_to_binary(moments, fields, code_time, composition.n_ion_species,
         composition.n_neutral_species, io_moments, 1, r, z)
-    write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, code_time,
-        composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, r, z, vperp,
-        vpa, vzeta, vr, vz)
+    write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, moments, fields,
+         code_time, composition.n_ion_species, composition.n_neutral_species, io_dfns, 1,
+         r, z, vperp, vpa, vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 395813350..e3a0bd68c 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -429,6 +429,18 @@ function analyze_and_plot_data(prefix...)
             get_tuple_of_return_values(allocate_global_zr_neutral_moments, nz_global,
                                        nr_global, n_neutral_species, ntime)
     end
+    phi_at_pdf_times, Ez_at_pdf_times, Er_at_pdf_times =
+        get_tuple_of_return_values(allocate_global_zr_fields, nz_global, nr_global, ntime)
+    density_at_pdf_times, parallel_flow_at_pdf_times, parallel_pressure_at_pdf_times,
+    parallel_heat_flux_at_pdf_times, thermal_speed_at_pdf_times =
+        get_tuple_of_return_values(allocate_global_zr_charged_moments, nz_global,
+                                   nr_global, n_ion_species, ntime)
+    if any(n_neutral_species .> 0)
+        neutral_density_at_pdf_times, neutral_uz_at_pdf_times, neutral_pz_at_pdf_times,
+        neutral_qz_at_pdf_times, neutral_thermal_speed_at_pdf_times =
+            get_tuple_of_return_values(allocate_global_zr_neutral_moments, nz_global,
+                                       nr_global, n_neutral_species, ntime)
+    end
     # read in the data from different block files
     # grids
     z, z_wgts, r, r_wgts =
@@ -470,6 +482,47 @@ function analyze_and_plot_data(prefix...)
                                    "thermal_speed_neutral", run_names, "moments", nblocks,
                                    nz_local, nr_local)
     end
+    # fields
+    get_tuple_of_return_values(read_distributed_zr_data!, phi_at_pdf_times, "phi",
+                               run_names, "dfns", nblocks, nz_local, nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, Ez_at_pdf_times, "Ez",
+                               run_names, "dfns", nblocks, nz_local, nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, Er_at_pdf_times, "Er",
+                               run_names, "dfns", nblocks, nz_local, nr_local)
+    # charged particle moments
+    get_tuple_of_return_values(read_distributed_zr_data!, density_at_pdf_times, "density",
+                               run_names, "dfns", nblocks, nz_local, nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, parallel_flow_at_pdf_times,
+                               "parallel_flow", run_names, "dfns", nblocks, nz_local,
+                               nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, parallel_pressure_at_pdf_times,
+                               "parallel_pressure", run_names, "dfns", nblocks,
+                               nz_local, nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, parallel_heat_flux_at_pdf_times,
+                               "parallel_heat_flux", run_names, "dfns", nblocks,
+                               nz_local, nr_local)
+    get_tuple_of_return_values(read_distributed_zr_data!, thermal_speed_at_pdf_times,
+                               "thermal_speed", run_names, "dfns", nblocks, nz_local,
+                               nr_local)
+    # neutral particle moments
+    if any(n_neutral_species .> 0)
+        get_tuple_of_return_values(read_distributed_zr_data!,
+                                   neutral_density_at_pdf_times, "density_neutral",
+                                   run_names, "dfns", nblocks, nz_local, nr_local)
+        get_tuple_of_return_values(read_distributed_zr_data!, neutral_uz_at_pdf_times,
+                                   "uz_neutral", run_names, "dfns", nblocks, nz_local,
+                                   nr_local)
+        get_tuple_of_return_values(read_distributed_zr_data!, neutral_pz_at_pdf_times,
+                                   "pz_neutral", run_names, "dfns", nblocks, nz_local,
+                                   nr_local)
+        get_tuple_of_return_values(read_distributed_zr_data!, neutral_qz_at_pdf_times,
+                                   "qz_neutral", run_names, "dfns", nblocks, nz_local,
+                                   nr_local)
+        get_tuple_of_return_values(read_distributed_zr_data!,
+                                   neutral_thermal_speed_at_pdf_times,
+                                   "thermal_speed_neutral", run_names, "dfns", nblocks,
+                                   nz_local, nr_local)
+    end
     # load time data from `dfns' cdf
     dfns_files0 = get_tuple_of_return_values(open_readonly_output_file, run_names, "dfns")
     # note that ntime may differ in these output files
@@ -529,6 +582,12 @@ function analyze_and_plot_data(prefix...)
             Tuple(ppar[:,ir0,:,:] for ppar ∈ parallel_pressure),
             Tuple(qpar[:,ir0,:,:] for qpar ∈ parallel_heat_flux),
             Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed),
+            Tuple(p[:,ir0,:] for p ∈ phi_at_pdf_times),
+            Tuple(n[:,ir0,:,:] for n ∈ density_at_pdf_times),
+            Tuple(upar[:,ir0,:,:] for upar ∈ parallel_flow_at_pdf_times),
+            Tuple(ppar[:,ir0,:,:] for ppar ∈ parallel_pressure_at_pdf_times),
+            Tuple(qpar[:,ir0,:,:] for qpar ∈ parallel_heat_flux_at_pdf_times),
+            Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed_at_pdf_times),
             Tuple(f[:,ivperp0,:,ir0,:,:] for f ∈ ff),
             n_ion_species, n_neutral_species, evolve_density, evolve_upar, evolve_ppar,
             nvpa, vpa, vpa_wgts, nz_global, z, z_wgts, Lz, ntime, time, ntime_pdfs,
@@ -549,10 +608,6 @@ function analyze_and_plot_data(prefix...)
 
     diagnostics_chodura = false
     if diagnostics_chodura
-        density_at_pdf_times = get_tuple_of_return_values(moment_at_pdf_times, density,
-                                                          ntime, ntime_pdfs)
-        Er_at_pdf_tames = get_tuple_of_return_values(moment_at_pdf_times, Er, ntime,
-                                                     ntime_pdfs)
         Chodura_ratio_lower, Chodura_ratio_upper =
             get_tuple_of_return_values(check_Chodura_condition, run_name, vpa_local,
                                        vpa_local_wgts, vperp_local, vperp_local_wgts,
@@ -809,9 +864,12 @@ end
 """
 function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         nwrite_movie_pdfs, itime_min_pdfs, itime_max_pdfs, ivpa0, iz0, ir0, r, phi,
-        density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed, ff,
-        n_ion_species, n_neutral_species, evolve_density, evolve_upar, evolve_ppar, nvpa,
-        vpa, vpa_wgts, nz, z, z_wgts, Lz, ntime, time, ntime_pdfs, time_pdfs)
+        density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed,
+        phi_at_pdf_times, density_at_pdf_times, parallel_flow_at_pdf_times,
+        parallel_pressure_at_pdf_times, parallel_heat_flux_at_pdf_times,
+        thermal_speed_at_pdf_times, ff, n_ion_species, n_neutral_species, evolve_density,
+        evolve_upar, evolve_ppar, nvpa, vpa, vpa_wgts, nz, z, z_wgts, Lz, ntime, time,
+        ntime_pdfs, time_pdfs)
 
     n_runs = length(run_names)
 
@@ -868,16 +926,14 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         else
             spec_string = ""
         end
-        if ntime == ntime_pdfs
-            # plot difference between evolved density and ∫dvpa f; only possibly different if density removed from
-            # normalised distribution function at run-time
-            plot(legend=legend)
-            for (t, n, n_int, run_label) ∈ zip(time, density, dens_moment, run_names)
-                @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)
+        # plot difference between evolved density and ∫dvpa f; only possibly different if density removed from
+        # normalised distribution function at run-time
+        plot(legend=legend)
+        for (t, n, n_int, run_label) ∈ zip(time_pdfs, density_at_pdf_times, dens_moment, run_names)
+            @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
         plot(legend=legend)
@@ -891,17 +947,15 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         end
         outfile = string(prefix, "_intwf0_vs_t", spec_string, ".pdf")
         savefig(outfile)
-        if ntime == ntime_pdfs
-            # 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
-            plot(legend=legend)
-            for (t, ppar, ppar_int, run_label) ∈ zip(time_pdfs, parallel_pressure, ppar_moment, run_names)
-                @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)
+        # 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
+        plot(legend=legend)
+        for (t, ppar, ppar_int, run_label) ∈ zip(time_pdfs, parallel_pressure_at_pdf_times, ppar_moment, run_names)
+            @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,:])
         if pp.animate_f_vs_vpa_z
@@ -963,9 +1017,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             it = itime_max_pdfs
             # 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_names)
+                 run_label) ∈ zip(1:n_runs, ff, density_at_pdf_times,
+                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
+                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
+                                  run_names)
 
                 PyPlot.subplot(1, n_runs, run_ind)
                 @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
@@ -981,9 +1036,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             it = itime_max_pdfs
             # 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_names)
+                 run_label) ∈ zip(1:n_runs, ff, density_at_pdf_times,
+                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
+                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
+                                  run_names)
 
                 @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
                     f[:,1,is,it], this_vpa, n[1,is,it], upar[1,is,it], vth[1,is,it],
@@ -997,9 +1053,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             it = itime_max_pdfs
             # 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_names)
+                 run_label) ∈ zip(1:n_runs, ff, density_at_pdf_times,
+                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
+                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
+                                  run_names)
 
                 @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
                     f[:,end,is,it], this_vpa, n[end,is,it], upar[end,is,it], vth[end,is,it],
@@ -1051,9 +1108,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                 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_names)
+                     run_label) ∈ zip(1:n_runs, ff, density_at_pdf_times,
+                                      parallel_flow_at_pdf_times,
+                                      thermal_speed_at_pdf_times, evolve_density,
+                                      evolve_upar, evolve_ppar, z, vpa, run_names)
 
                     PyPlot.subplot(1, n_runs, run_ind)
                     @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
@@ -1074,9 +1132,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                 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_names)
+                     run_label) ∈ zip(1:n_runs, ff, density_at_pdf_times,
+                                      parallel_flow_at_pdf_times,
+                                      thermal_speed_at_pdf_times, evolve_density,
+                                      evolve_upar, evolve_ppar, z, vpa, run_names)
 
                     PyPlot.subplot(1, n_runs, run_ind)
                     @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
@@ -1097,8 +1156,9 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 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_names)
+                    zip(ff, density_at_pdf_times, parallel_flow_at_pdf_times,
+                        thermal_speed_at_pdf_times, evolve_density, evolve_upar,
+                        evolve_ppar, vpa, run_names)
                     @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)
@@ -1112,8 +1172,9 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 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_names)
+                    zip(ff, density_at_pdf_times, parallel_flow_at_pdf_times,
+                        thermal_speed_at_pdf_times, evolve_density, evolve_upar,
+                        evolve_ppar, vpa, run_names)
                     @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)
@@ -1935,44 +1996,6 @@ function plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; plot_log=false, kwargs...)
     return p
 end
 
-"""
-Get values of 'moment' variable at time points where distribution functions were saved
-"""
-function moment_at_pdf_times(moment, ntime, ntime_pdfs)
-    if ntime_pdfs == 1
-        # Distribution functions only written at initial time
-        step = ntime + 1
-    else
-        if (ntime-1) % (ntime_pdfs-1) != 0
-            #error("ntime is not a multiple of ntime_pdfs => nwrite_pdfs was not a multiple "
-            #      * "of nwrite so cannot get moment at time points where pdf was saved.")
-            ntime = ((ntime-1)÷(ntime_pdfs-1)) * (ntime_pdfs-1) + 1
-        end
-
-        step = (ntime-1) ÷ (ntime_pdfs-1)
-    end
-    n = ndims(moment)
-    if n == 1
-        return moment[begin:step:end]
-    elseif n == 2
-        return moment[:,begin:step:end]
-    elseif n == 3
-        return moment[:,:,begin:step:end]
-    elseif n == 4
-        return moment[:,:,:,begin:step:end]
-    elseif n == 5
-        return moment[:,:,:,:,begin:step:end]
-    elseif n == 6
-        return moment[:,:,:,:,:,begin:step:end]
-    elseif n == 7
-        return moment[:,:,:,:,:,:,begin:step:end]
-    elseif n == 8
-        return moment[:,:,:,:,:,:,:,begin:step:end]
-    else
-        error("ndims=$n is not supported yet")
-    end
-end
-
 """
 """
 function compare_fields_symbolic_test(run_name,field,field_sym,z,r,time,nz,nr,ntime,field_label,field_sym_label,norm_label,file_string)
diff --git a/src/time_advance.jl b/src/time_advance.jl
index b460b608e..c2e84cfea 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -889,8 +889,8 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                                    Dates.format(now(), dateformat"H:MM:SS"))
                 end
             end
-            write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, t,
-                                      composition.n_ion_species,
+            write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, moments, fields,
+                                      t, composition.n_ion_species,
                                       composition.n_neutral_species, io_dfns, iwrite_dfns,
                                       r, z, vperp, vpa, vzeta, vr, vz)
             iwrite_dfns += 1

From f0afe3829969becf9d6bd894e955db522b0c690b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 12 Mar 2023 23:56:04 +0000
Subject: [PATCH 015/428] Make File I/O support String variables

---
 src/file_io_hdf5.jl   | 4 ++--
 src/file_io_netcdf.jl | 4 ++--
 src/load_data.jl      | 6 +++++-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/file_io_hdf5.jl b/src/file_io_hdf5.jl
index b79ce6dc4..aee3a2cdc 100644
--- a/src/file_io_hdf5.jl
+++ b/src/file_io_hdf5.jl
@@ -62,9 +62,9 @@ end
 
 # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group
 function write_single_value!(file_or_group::HDF5.H5DataStore, name,
-                             data::Union{Number,AbstractArray{T,N}},
+                             data::Union{Number, AbstractString, AbstractArray{T,N}},
                              coords::coordinate...; parallel_io, description=nothing) where {T,N}
-    if isa(data, Number)
+    if isa(data, Union{Number, AbstractString})
         file_or_group[name] = data
         return nothing
     end
diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 9571d11f6..3648781d5 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -59,7 +59,7 @@ function maybe_create_netcdf_dim(file_or_group::NCDataset, coord::coordinate)
 end
 
 function write_single_value!(file_or_group::NCDataset, name,
-                             value::Union{Number, AbstractArray{T,N}},
+                             value::Union{Number, AbstractString, AbstractArray{T,N}},
                              coords::coordinate...; parallel_io, description=nothing) where {T,N}
     if description !== nothing
         attributes = Dict("description" => description)
@@ -67,7 +67,7 @@ function write_single_value!(file_or_group::NCDataset, name,
         attributes = ()
     end
 
-    if isa(value, Number)
+    if isa(value, Number) || isa(value, String)
         coords !== () && error("cannot pass coordinates with a scalar")
         type = typeof(value)
         dims = ()
diff --git a/src/load_data.jl b/src/load_data.jl
index 8ebff4364..caa8c137e 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -91,7 +91,11 @@ function load_variable(file_or_group::NCDataset, name::String)
     # This overload deals with cases where fid is a NetCDF `Dataset` (which could be a
     # file or a group).
     try
-        var = file_or_group[name].var[:]
+        if size(file_or_group[name].var) == ()
+            var = file_or_group[name].var[]
+        else
+            var = file_or_group[name].var[:]
+        end
         if isa(var, Char)
             var = (var == Char(true))
         end

From 0cae0643b0d2016000b704a2f530452e833425f9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 12 Mar 2023 22:42:04 +0000
Subject: [PATCH 016/428] Return coordinate struct from load_coordinate_data!()

Makes it easier to use moment_kinetics functions in post-processing.
---
 src/analysis.jl                    |  67 ++--
 src/file_io.jl                     |  18 ++
 src/load_data.jl                   |  30 +-
 src/plot_MMS_sequence.jl           |  84 ++---
 src/post_processing.jl             | 498 ++++++++++++++++-------------
 test/harrisonthompson.jl           |   9 +-
 test/nonlinear_sound_wave_tests.jl |  15 +-
 test/sound_wave_tests.jl           |  16 +-
 test/wall_bc_tests.jl              |   6 +-
 9 files changed, 416 insertions(+), 327 deletions(-)

diff --git a/src/analysis.jl b/src/analysis.jl
index f3b1a9367..40029248c 100644
--- a/src/analysis.jl
+++ b/src/analysis.jl
@@ -13,15 +13,15 @@ using ..velocity_moments: integrate_over_vspace
 
 """
 """
-function analyze_fields_data(phi, ntime, nz, z_wgts, Lz)
+function analyze_fields_data(phi, ntime, z)
     print("Analyzing fields data...")
     phi_fldline_avg = allocate_float(ntime)
     for i ∈ 1:ntime
-        phi_fldline_avg[i] = field_line_average(view(phi,:,i), z_wgts, Lz)
+        phi_fldline_avg[i] = field_line_average(view(phi,:,i), z.wgts, z.L)
     end
     # delta_phi = phi - <phi> is the fluctuating phi
-    delta_phi = allocate_float(nz,ntime)
-    for iz ∈ 1:nz
+    delta_phi = allocate_float(z.n,ntime)
+    for iz ∈ 1:z.n
         delta_phi[iz,:] .= phi[iz,:] - phi_fldline_avg
     end
     println("done.")
@@ -168,70 +168,70 @@ end
 """
 """
 function analyze_moments_data(density, parallel_flow, parallel_pressure, thermal_speed,
-                              parallel_heat_flux, ntime, n_species, nz, z_wgts, Lz)
+                              parallel_heat_flux, ntime, n_species, z)
     print("Analyzing velocity moments data...")
     density_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
         for i ∈ 1:ntime
-            density_fldline_avg[is,i] = field_line_average(view(density,:,is,i), z_wgts, Lz)
+            density_fldline_avg[is,i] = field_line_average(view(density,:,is,i), z.wgts, z.L)
         end
     end
     upar_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
         for i ∈ 1:ntime
-            upar_fldline_avg[is,i] = field_line_average(view(parallel_flow,:,is,i), z_wgts, Lz)
+            upar_fldline_avg[is,i] = field_line_average(view(parallel_flow,:,is,i), z.wgts, z.L)
         end
     end
     ppar_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
         for i ∈ 1:ntime
-            ppar_fldline_avg[is,i] = field_line_average(view(parallel_pressure,:,is,i), z_wgts, Lz)
+            ppar_fldline_avg[is,i] = field_line_average(view(parallel_pressure,:,is,i), z.wgts, z.L)
         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)
+            vth_fldline_avg[is,i] = field_line_average(view(thermal_speed,:,is,i), z.wgts, z.L)
         end
     end
     qpar_fldline_avg = allocate_float(n_species, ntime)
     for is ∈ 1:n_species
         for i ∈ 1:ntime
-            qpar_fldline_avg[is,i] = field_line_average(view(parallel_heat_flux,:,is,i), z_wgts, Lz)
+            qpar_fldline_avg[is,i] = field_line_average(view(parallel_heat_flux,:,is,i), z.wgts, z.L)
         end
     end
     # delta_density = n_s - <n_s> is the fluctuating density
-    delta_density = allocate_float(nz,n_species,ntime)
+    delta_density = allocate_float(z.n,n_species,ntime)
     for is ∈ 1:n_species
-        for iz ∈ 1:nz
+        for iz ∈ 1:z.n
             @. delta_density[iz,is,:] = density[iz,is,:] - density_fldline_avg[is,:]
         end
     end
     # delta_upar = upar_s - <upar_s> is the fluctuating parallel flow
-    delta_upar = allocate_float(nz,n_species,ntime)
+    delta_upar = allocate_float(z.n,n_species,ntime)
     for is ∈ 1:n_species
-        for iz ∈ 1:nz
+        for iz ∈ 1:z.n
             @. delta_upar[iz,is,:] = parallel_flow[iz,is,:] - upar_fldline_avg[is,:]
         end
     end
     # delta_ppar = ppar_s - <ppar_s> is the fluctuating parallel pressure
-    delta_ppar = allocate_float(nz,n_species,ntime)
+    delta_ppar = allocate_float(z.n,n_species,ntime)
     for is ∈ 1:n_species
-        for iz ∈ 1:nz
+        for iz ∈ 1:z.n
             @. 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)
+    delta_vth = allocate_float(z.n,n_species,ntime)
     for is ∈ 1:n_species
-        for iz ∈ 1:nz
+        for iz ∈ 1:z.n
             @. 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)
+    delta_qpar = allocate_float(z.n,n_species,ntime)
     for is ∈ 1:n_species
-        for iz ∈ 1:nz
+        for iz ∈ 1:z.n
             @. delta_qpar[iz,is,:] = parallel_heat_flux[iz,is,:] - qpar_fldline_avg[is,:]
         end
     end
@@ -242,31 +242,30 @@ end
 
 """
 """
-function analyze_pdf_data(ff, vpa, nvpa, nz, n_species, ntime, vpa_wgts, z_wgts, Lz,
-                          vth, evolve_ppar)
+function analyze_pdf_data(ff, n_species, ntime, z, vpa, vth, evolve_ppar)
     print("Analyzing distribution function data...")
-    f_fldline_avg = allocate_float(nvpa,n_species,ntime)
+    f_fldline_avg = allocate_float(vpa.n,n_species,ntime)
     for i ∈ 1:ntime
         for is ∈ 1:n_species
-            for ivpa ∈ 1:nvpa
-                f_fldline_avg[ivpa,is,i] = field_line_average(view(ff,ivpa,:,is,i), z_wgts, Lz)
+            for ivpa ∈ 1:vpa.n
+                f_fldline_avg[ivpa,is,i] = field_line_average(view(ff,ivpa,:,is,i), z.wgts, z.L)
             end
         end
     end
     # delta_f = f - <f> is the fluctuating distribution function
-    delta_f = allocate_float(nvpa,nz,n_species,ntime)
-    for iz ∈ 1:nz
+    delta_f = allocate_float(vpa.n,z.n,n_species,ntime)
+    for iz ∈ 1:z.n
         @. delta_f[:,iz,:,:] = ff[:,iz,:,:] - f_fldline_avg
     end
-    dens_moment = allocate_float(nz,n_species,ntime)
-    upar_moment = allocate_float(nz,n_species,ntime)
-    ppar_moment = allocate_float(nz,n_species,ntime)
+    dens_moment = allocate_float(z.n,n_species,ntime)
+    upar_moment = allocate_float(z.n,n_species,ntime)
+    ppar_moment = allocate_float(z.n,n_species,ntime)
     for i ∈ 1:ntime
         for is ∈ 1:n_species
-            for iz ∈ 1:nz
-                @views dens_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa_wgts)
-                @views upar_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa, vpa_wgts)
-                @views ppar_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa, 2, vpa_wgts)
+            for iz ∈ 1:z.n
+                @views dens_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa.wgts)
+                @views upar_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa.grid, vpa.wgts)
+                @views ppar_moment[iz,is,i] = integrate_over_vspace(ff[:,iz,is,i], vpa.grid, 2, vpa.wgts)
             end
         end
     end
diff --git a/src/file_io.jl b/src/file_io.jl
index 808e6c1a4..795ea6899 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -288,6 +288,11 @@ function define_io_coordinate!(parent, coord, coord_name, description, parallel_
                                 description="number of local $coord_name grid points")
         end
 
+        # write the number of points within each element for this coordinate to variable
+        # "ngrid" within "coords/coord_name" group
+        write_single_value!(group, "ngrid", coord.ngrid; parallel_io=parallel_io,
+                            description="number of points in each element in $coord_name")
+
         # write the number of global grid points for this coordinate to variable "n_local"
         # within "coords/coord_name" group
         write_single_value!(group, "n_global", coord.n_global; parallel_io=parallel_io,
@@ -310,6 +315,19 @@ function define_io_coordinate!(parent, coord, coord_name, description, parallel_
         write_single_value!(group, "wgts", coord.wgts, coord; parallel_io=parallel_io,
                             description="integration weights associated with the $coord_name grid points")
 
+        # write the discretization option for the coordinate
+        write_single_value!(group, "discretization", coord.discretization;
+                            parallel_io=parallel_io,
+                            description="discretization used for $coord_name")
+
+        # write the finite-difference option for the coordinate
+        write_single_value!(group, "fd_option", coord.fd_option; parallel_io=parallel_io,
+                            description="type of finite difference for $coord_name, if used")
+
+        # write the boundary condition for the coordinate
+        write_single_value!(group, "bc", coord.bc; parallel_io=parallel_io,
+                            description="boundary condition for $coord_name")
+
         return group
     end
 
diff --git a/src/load_data.jl b/src/load_data.jl
index caa8c137e..3d7b88e81 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -14,9 +14,12 @@ export load_block_data
 export load_rank_data
 export load_species_data
 
+using ..coordinates: define_coordinate
 using ..file_io: get_group
+using ..input_structs: advection_input, grid_input
 
 using HDF5
+using MPI
 using NCDatasets
 
 """
@@ -116,14 +119,39 @@ function load_coordinate_data(fid, name; printout=false)
 
     coord_group = get_group(get_group(fid, "coords"), name)
 
+    ngrid = load_variable(coord_group, "ngrid")
     n_local = load_variable(coord_group, "n_local")
     n_global = load_variable(coord_group, "n_global")
     grid = load_variable(coord_group, "grid")
     wgts = load_variable(coord_group, "wgts")
+    irank = load_variable(coord_group, "irank")
     # L = global box length
     L = load_variable(coord_group, "L")
+    discretization = load_variable(coord_group, "discretization")
+    fd_option = load_variable(coord_group, "fd_option")
+    bc = load_variable(coord_group, "bc")
 
-    return n_local, n_global, grid, wgts, L
+    nelement_local = nothing
+    if n_local == 1 && ngrid == 1
+        nelement_local = 1
+    else
+        nelement_local = (n_local-1) ÷ (ngrid-1)
+    end
+    if n_global == 1 && ngrid == 1
+        nelement_global = 1
+    else
+        nelement_global = (n_global-1) ÷ (ngrid-1)
+    end
+
+    # Define input to create coordinate struct
+    # Some dummy inputs, at least for now: nrank=0
+    input = grid_input(name, ngrid, nelement_global, nelement_local, 0, irank, L,
+                       discretization, fd_option, bc, advection_input("", 0.0, 0.0, 0.0),
+                       MPI.COMM_NULL)
+
+    coord, spectral = define_coordinate(input)
+
+    return coord, spectral
 end
 
 """
diff --git a/src/plot_MMS_sequence.jl b/src/plot_MMS_sequence.jl
index 103a48652..d006ea10c 100644
--- a/src/plot_MMS_sequence.jl
+++ b/src/plot_MMS_sequence.jl
@@ -139,8 +139,8 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         # load local sizes of grids stored on each netCDF file 
         # z z_wgts r r_wgts may take different values on different blocks
         # we need to construct the global grid below
-        nz_local, nz_global, z_local, z_wgts, Lz = load_coordinate_data(fid, "z")
-        nr_local, nr_global, r_local, r_wgts, Lr = load_coordinate_data(fid, "r")
+        z = load_coordinate_data(fid, "z")
+        r = load_coordinate_data(fid, "r")
         # load time data 
         ntime, time = load_time_data(fid)
         # load species data 
@@ -150,72 +150,72 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         # load local velocity coordinate data from `moments' cdf
         # these values are currently the same for all blocks 
         fid = open_readonly_output_file(run_name,"dfns")
-        nvpa, nvpa_global, vpa, vpa_wgts, Lvpa = load_coordinate_data(fid, "vpa")
-        nvperp, nvperp_global, vperp, vperp_wgts, Lvperp = load_coordinate_data(fid, "vperp")
+        vpa = load_coordinate_data(fid, "vpa")
+        vperp = load_coordinate_data(fid, "vperp")
         if n_neutral_species > 0
-            nvz, nvz_global, vz, vz_wgts, Lvz = load_coordinate_data(fid, "vz")
-            nvr, nvr_global, vr, vr_wgts, Lvr = load_coordinate_data(fid, "vr")
-            nvzeta, nvzeta_global, vzeta, vzeta_wgts, Lvzeta = load_coordinate_data(fid, "vzeta")
+            vz = load_coordinate_data(fid, "vz")
+            vr = load_coordinate_data(fid, "vr")
+            vzeta = load_coordinate_data(fid, "vzeta")
         else # define nvz nvr nvzeta to avoid errors below
-            nvz = 1
-            nvr = 1
-            nvzeta = 1
+            vz.n = 1
+            vr.n = 1
+            vzeta.n = 1
         end
         close(fid)
         
         
         # allocate arrays to contain the global fields as a function of (z,r,t)
-        phi, Ez, Er = allocate_global_zr_fields(nz_global,nr_global,ntime)
+        phi, Ez, Er = allocate_global_zr_fields(z.n_global,r.n_global,ntime)
         density, parallel_flow, parallel_pressure, parallel_heat_flux,
-            thermal_speed = allocate_global_zr_charged_moments(nz_global,nr_global,n_ion_species,ntime)
+            thermal_speed = allocate_global_zr_charged_moments(z.n_global,r.n_global,n_ion_species,ntime)
         if n_neutral_species > 0
             neutral_density, neutral_uz, neutral_pz, 
-             neutral_qz, neutral_thermal_speed = allocate_global_zr_neutral_moments(nz_global,nr_global,n_neutral_species,ntime)
+             neutral_qz, neutral_thermal_speed = allocate_global_zr_neutral_moments(z.n_global,r.n_global,n_neutral_species,ntime)
         end 
         # read in the data from different block netcdf files
         # grids 
         z, z_wgts, r, r_wgts = construct_global_zr_grids(run_name,
-           "moments",nz_global,nr_global,nblocks)
+           "moments",z.n_global,r.n_global,nblocks)
         #println("z: ",z)
         #println("r: ",r)
         
         # fields 
-        read_distributed_zr_data!(phi,"phi",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(Ez,"Ez",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(Er,"Er",run_name,"moments",nblocks,nz_local,nr_local) 
+        read_distributed_zr_data!(phi,"phi",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(Ez,"Ez",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(Er,"Er",run_name,"moments",nblocks,z.n,r.n) 
         # charged particle moments
-        read_distributed_zr_data!(density,"density",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(parallel_flow,"parallel_flow",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(parallel_pressure,"parallel_pressure",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(parallel_heat_flux,"parallel_heat_flux",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(thermal_speed,"thermal_speed",run_name,"moments",nblocks,nz_local,nr_local) 
+        read_distributed_zr_data!(density,"density",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(parallel_flow,"parallel_flow",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(parallel_pressure,"parallel_pressure",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(parallel_heat_flux,"parallel_heat_flux",run_name,"moments",nblocks,z.n,r.n) 
+        read_distributed_zr_data!(thermal_speed,"thermal_speed",run_name,"moments",nblocks,z.n,r.n) 
         # neutral particle moments 
         if n_neutral_species > 0
-            read_distributed_zr_data!(neutral_density,"density_neutral",run_name,"moments",nblocks,nz_local,nr_local) 
-            read_distributed_zr_data!(neutral_uz,"uz_neutral",run_name,"moments",nblocks,nz_local,nr_local) 
-            read_distributed_zr_data!(neutral_pz,"pz_neutral",run_name,"moments",nblocks,nz_local,nr_local) 
-            read_distributed_zr_data!(neutral_qz,"qz_neutral",run_name,"moments",nblocks,nz_local,nr_local) 
-            read_distributed_zr_data!(neutral_thermal_speed,"thermal_speed_neutral",run_name,"moments",nblocks,nz_local,nr_local) 
+            read_distributed_zr_data!(neutral_density,"density_neutral",run_name,"moments",nblocks,z.n,r.n) 
+            read_distributed_zr_data!(neutral_uz,"uz_neutral",run_name,"moments",nblocks,z.n,r.n) 
+            read_distributed_zr_data!(neutral_pz,"pz_neutral",run_name,"moments",nblocks,z.n,r.n) 
+            read_distributed_zr_data!(neutral_qz,"qz_neutral",run_name,"moments",nblocks,z.n,r.n) 
+            read_distributed_zr_data!(neutral_thermal_speed,"thermal_speed_neutral",run_name,"moments",nblocks,z.n,r.n) 
         end
         
         
         r_bc = get(scan_input, "r_bc", "periodic")
         z_bc = get(scan_input, "z_bc", "periodic")
         # avoid passing Lr = 0 into manufactured_solns functions 
-        if nr_local > 1
-            Lr_in = Lr 
+        if r.n > 1
+            Lr_in = r.L 
         else 
             Lr_in = 1.0
         end
         geometry, composition = get_geometry_and_composition(scan_input,n_ion_species,n_neutral_species)
         
-        manufactured_solns_list = manufactured_solutions(Lr_in,Lz,r_bc,z_bc,geometry,composition,nr_local) 
+        manufactured_solns_list = manufactured_solutions(Lr_in,z.L,r_bc,z_bc,geometry,composition,r.n) 
         dfni_func = manufactured_solns_list.dfni_func
         densi_func = manufactured_solns_list.densi_func
         dfnn_func = manufactured_solns_list.dfnn_func
         densn_func = manufactured_solns_list.densn_func
         
-        manufactured_E_fields = manufactured_electric_fields(Lr_in,Lz,r_bc,z_bc,composition,nr_local)
+        manufactured_E_fields = manufactured_electric_fields(Lr_in,z.L,r_bc,z_bc,composition,r.n)
         Er_func = manufactured_E_fields.Er_func
         Ez_func = manufactured_E_fields.Ez_func
         phi_func = manufactured_E_fields.phi_func
@@ -225,21 +225,21 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         Er_sym = copy(phi[:,:,:])
         Ez_sym = copy(phi[:,:,:])
         for it in 1:ntime
-            for ir in 1:nr_global
-                for iz in 1:nz_global
+            for ir in 1:r.n_global
+                for iz in 1:z.n_global
                     phi_sym[iz,ir,it] = phi_func(z[iz],r[ir],time[it])
                     Ez_sym[iz,ir,it] = Ez_func(z[iz],r[ir],time[it])
                     Er_sym[iz,ir,it] = Er_func(z[iz],r[ir],time[it])
                 end
             end
         end
-        phi_error_t = compare_fields_symbolic_test(run_name,phi,phi_sym,z,r,time,nz_global,nr_global,ntime,
+        phi_error_t = compare_fields_symbolic_test(run_name,phi,phi_sym,z,r,time,z.n_global,r.n_global,ntime,
          L"\widetilde{\phi}",L"\widetilde{\phi}^{sym}",L"\sqrt{\sum || \widetilde{\phi} - \widetilde{\phi}^{sym} ||^2 / N} ","phi")
         phi_error_sequence[isim] = phi_error_t[end]
-        Er_error_t = compare_fields_symbolic_test(run_name,Er,Er_sym,z,r,time,nz_global,nr_global,ntime,
+        Er_error_t = compare_fields_symbolic_test(run_name,Er,Er_sym,z,r,time,z.n_global,r.n_global,ntime,
          L"\widetilde{E_r}",L"\widetilde{E_r}^{sym}",L"\sqrt{\sum || \widetilde{E_r} - \widetilde{E_r}^{sym} ||^2 /N} ","Er")
         Er_error_sequence[isim] = Er_error_t[end]
-        Ez_error_t = compare_fields_symbolic_test(run_name,Ez,Ez_sym,z,r,time,nz_global,nr_global,ntime,
+        Ez_error_t = compare_fields_symbolic_test(run_name,Ez,Ez_sym,z,r,time,z.n_global,r.n_global,ntime,
          L"\widetilde{E_z}",L"\widetilde{E_z}^{sym}",L"\sqrt{\sum || \widetilde{E_z} - \widetilde{E_z}^{sym} ||^2 /N} ","Ez")
         Ez_error_sequence[isim] = Ez_error_t[end]
         
@@ -247,13 +247,13 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         density_sym = copy(density[:,:,:,:])
         is = 1
         for it in 1:ntime
-            for ir in 1:nr_global
-                for iz in 1:nz_global
+            for ir in 1:r.n_global
+                for iz in 1:z.n_global
                     density_sym[iz,ir,is,it] = densi_func(z[iz],r[ir],time[it])
                 end
             end
         end
-        ion_density_error_t = compare_moments_symbolic_test(run_name,density,density_sym,"ion",z,r,time,nz_global,nr_global,ntime,
+        ion_density_error_t = compare_moments_symbolic_test(run_name,density,density_sym,"ion",z,r,time,z.n_global,r.n_global,ntime,
          L"\widetilde{n}_i",L"\widetilde{n}_i^{sym}",L"\sum || \widetilde{n}_i - \widetilde{n}_i^{sym} ||^2 ","dens")
         # use final time point for analysis
         ion_density_error_sequence[isim] = ion_density_error_t[end] 
@@ -267,13 +267,13 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
             neutral_density_sym = copy(density[:,:,:,:])
             is = 1
             for it in 1:ntime
-                for ir in 1:nr_global
-                    for iz in 1:nz_global
+                for ir in 1:r.n_global
+                    for iz in 1:z.n_global
                         neutral_density_sym[iz,ir,is,it] = densn_func(z[iz],r[ir],time[it])
                     end
                 end
             end
-            neutral_density_error_t = compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z,r,time,nz_global,nr_global,ntime,
+            neutral_density_error_t = compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z,r,time,z.n_global,r.n_global,ntime,
              L"\widetilde{n}_n",L"\widetilde{n}_n^{sym}",L"\sum || \widetilde{n}_n - \widetilde{n}_n^{sym} ||^2 ","dens")
             neutral_density_error_sequence[isim] = neutral_density_error_t[end]
         
diff --git a/src/post_processing.jl b/src/post_processing.jl
index e3a0bd68c..f08027df4 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -201,30 +201,30 @@ function construct_global_zr_grids(run_name::String,file_key::String,nz_global::
     # whereas the optimal routine would loop over a single z/r group
     for iblock in 0:nblocks-1
         fid = open_readonly_output_file(run_name,file_key,iblock=iblock,printout=false)
-        nz_local, nz_global, z_local, z_local_wgts, Lz = load_coordinate_data(fid, "z")
-        nr_local, nr_global, r_local, r_local_wgts, Lr = load_coordinate_data(fid, "r")
+        z, z_spectral = load_coordinate_data(fid, "z")
+        r, r_spectral = load_coordinate_data(fid, "r")
 
         z_irank, r_irank = load_rank_data(fid)
         # MRH should wgts at element boundaries be doubled
         # MRH in the main code duplicated points have half the integration wgt
         if z_irank == 0
-            zgrid[1] = z_local[1]
-            zwgts[1] = z_local_wgts[1]
+            zgrid[1] = z.grid[1]
+            zwgts[1] = z.wgts[1]
         end
-        for iz_local in 2:nz_local
-            iz_global = iglobal_func(iz_local,z_irank,nz_local)
-            zgrid[iz_global] = z_local[iz_local]
-            zwgts[iz_global] = z_local_wgts[iz_local]
+        for iz_local in 2:z.n
+            iz_global = iglobal_func(iz_local,z_irank,z.n)
+            zgrid[iz_global] = z.grid[iz_local]
+            zwgts[iz_global] = z.wgts[iz_local]
         end
 
         if r_irank == 0
-            rgrid[1] = r_local[1]
-            rwgts[1] = r_local_wgts[1]
+            rgrid[1] = r.grid[1]
+            rwgts[1] = r.wgts[1]
         end
-        for ir_local in 2:nr_local
-            ir_global = iglobal_func(ir_local,r_irank,nr_local)
-            rgrid[ir_global] = r_local[ir_local]
-            rwgts[ir_global] = r_local_wgts[ir_local]
+        for ir_local in 2:r.n
+            ir_global = iglobal_func(ir_local,r_irank,r.n)
+            rgrid[ir_global] = r.grid[ir_local]
+            rwgts[ir_global] = r.wgts[ir_local]
         end
     end
     return zgrid, zwgts, rgrid, rwgts
@@ -404,10 +404,8 @@ function analyze_and_plot_data(prefix...)
     # load global and local sizes of grids stored on each output file
     # z z_wgts r r_wgts may take different values on different blocks
     # we need to construct the global grid below
-    nz_local, nz_global, z_local, z_local_wgts, Lz =
-        get_tuple_of_return_values(load_coordinate_data, moments_files0, "z")
-    nr_local, nr_global, r_local, r_local_wgts, Lr =
-        get_tuple_of_return_values(load_coordinate_data, moments_files0, "r")
+    z, z_spectral = get_tuple_of_return_values(load_coordinate_data, moments_files0, "z")
+    r, r_spectral = get_tuple_of_return_values(load_coordinate_data, moments_files0, "r")
     # load time data
     ntime, time = get_tuple_of_return_values(load_time_data, moments_files0)
     # load species data
@@ -419,158 +417,218 @@ function analyze_and_plot_data(prefix...)
     end
 
     # allocate arrays to contain the global fields as a function of (z,r,t)
-    phi, Ez, Er = get_tuple_of_return_values(allocate_global_zr_fields, nz_global,
-                                             nr_global, ntime)
+    phi, Ez, Er = get_tuple_of_return_values(allocate_global_zr_fields,
+                                             Tuple(this_z.n_global for this_z ∈ z),
+                                             Tuple(this_r.n_global for this_r ∈ r),
+                                             ntime)
     density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments, nz_global,
-                                   nr_global, n_ion_species, ntime)
+        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+                                   Tuple(this_z.n_global for this_z ∈ z),
+                                   Tuple(this_r.n_global for this_r ∈ r),
+                                   n_ion_species, ntime)
     if any(n_neutral_species .> 0)
         neutral_density, neutral_uz, neutral_pz, neutral_qz, neutral_thermal_speed =
-            get_tuple_of_return_values(allocate_global_zr_neutral_moments, nz_global,
-                                       nr_global, n_neutral_species, ntime)
-    end
-    phi_at_pdf_times, Ez_at_pdf_times, Er_at_pdf_times =
-        get_tuple_of_return_values(allocate_global_zr_fields, nz_global, nr_global, ntime)
-    density_at_pdf_times, parallel_flow_at_pdf_times, parallel_pressure_at_pdf_times,
-    parallel_heat_flux_at_pdf_times, thermal_speed_at_pdf_times =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments, nz_global,
-                                   nr_global, n_ion_species, ntime)
-    if any(n_neutral_species .> 0)
-        neutral_density_at_pdf_times, neutral_uz_at_pdf_times, neutral_pz_at_pdf_times,
-        neutral_qz_at_pdf_times, neutral_thermal_speed_at_pdf_times =
-            get_tuple_of_return_values(allocate_global_zr_neutral_moments, nz_global,
-                                       nr_global, n_neutral_species, ntime)
+            get_tuple_of_return_values(allocate_global_zr_neutral_moments,
+                                       Tuple(this_z.n_global for this_z ∈ z),
+                                       Tuple(this_r.n_global for this_r ∈ r),
+                                       n_neutral_species, ntime)
     end
     # read in the data from different block files
     # grids
-    z, z_wgts, r, r_wgts =
+    z_global, z_wgts, r_global, r_wgts =
         get_tuple_of_return_values(construct_global_zr_grids, run_names, "moments",
-                                   nz_global, nr_global, nblocks)
+                                   Tuple(this_z.n_global for this_z ∈ z),
+                                   Tuple(this_r.n_global for this_r ∈ r),
+                                   nblocks)
 
     # fields
     get_tuple_of_return_values(read_distributed_zr_data!, phi, "phi", run_names, "moments",
-                               nblocks, nz_local, nr_local)
+                               nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, Ez, "Ez", run_names, "moments",
-                               nblocks, nz_local, nr_local)
+                               nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, Er, "Er", run_names, "moments",
-                               nblocks, nz_local, nr_local)
+                               nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     # charged particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density, "density", run_names,
-                               "moments", nblocks, nz_local, nr_local)
+                               "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_flow, "parallel_flow",
-                               run_names, "moments", nblocks, nz_local, nr_local)
+                               run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_pressure,
                                "parallel_pressure", run_names, "moments", nblocks,
-                               nz_local, nr_local)
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_heat_flux,
                                "parallel_heat_flux", run_names, "moments", nblocks,
-                               nz_local, nr_local)
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, thermal_speed, "thermal_speed",
-                               run_names, "moments", nblocks, nz_local, nr_local)
+                               run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     # neutral particle moments
     if any(n_neutral_species .> 0)
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_density,
                                    "density_neutral", run_names, "moments", nblocks,
-                                   nz_local, nr_local)
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_uz, "uz_neutral",
-                                   run_names, "moments", nblocks, nz_local, nr_local)
+                                   run_names, "moments", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_pz, "pz_neutral",
-                                   run_names, "moments", nblocks, nz_local, nr_local)
+                                   run_names, "moments", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_qz, "qz_neutral",
-                                   run_names, "moments", nblocks, nz_local, nr_local)
+                                   run_names, "moments", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_thermal_speed,
                                    "thermal_speed_neutral", run_names, "moments", nblocks,
-                                   nz_local, nr_local)
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
+    end
+
+    # load time data from `dfns' cdf
+    dfns_files0 = get_tuple_of_return_values(open_readonly_output_file, run_names, "dfns")
+    # note that ntime may differ in these output files
+
+    ntime_pdfs, time_pdfs = get_tuple_of_return_values(load_time_data, dfns_files0)
+
+    # load local velocity coordinate data from `dfns' cdf
+    # these values are currently the same for all blocks
+    vpa, vpa_spectral = get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vpa")
+    vperp, vperp_spectral = get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vperp")
+    if any(n_neutral_species .> 0)
+        vzeta, vzeta_spectral = get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vzeta")
+        vr, vr_spectral = get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vr")
+        vz, vz_spectral = get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vz")
+    else # define nvz nvr nvzeta to avoid errors below
+        vzeta = vzeta_spectral = nothing
+        vr = vr_spectral = nothing
+        vz = vz_spectral = nothing
+    end
+
+    phi_at_pdf_times, Ez_at_pdf_times, Er_at_pdf_times =
+        get_tuple_of_return_values(allocate_global_zr_fields,
+                                   Tuple(this_z.n_global for this_z ∈ z),
+                                   Tuple(this_r.n_global for this_r ∈ r), ntime_pdfs)
+    density_at_pdf_times, parallel_flow_at_pdf_times, parallel_pressure_at_pdf_times,
+    parallel_heat_flux_at_pdf_times, thermal_speed_at_pdf_times =
+        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+                                   Tuple(this_z.n_global for this_z ∈ z),
+                                   Tuple(this_r.n_global for this_r ∈ r), n_ion_species,
+                                   ntime_pdfs)
+    if any(n_neutral_species .> 0)
+        neutral_density_at_pdf_times, neutral_uz_at_pdf_times, neutral_pz_at_pdf_times,
+        neutral_qz_at_pdf_times, neutral_thermal_speed_at_pdf_times =
+            get_tuple_of_return_values(allocate_global_zr_neutral_moments,
+                                       Tuple(this_z.n_global for this_z ∈ z),
+                                       Tuple(this_r.n_global for this_r ∈ r),
+                                       n_neutral_species, ntime_pdfs)
     end
     # fields
     get_tuple_of_return_values(read_distributed_zr_data!, phi_at_pdf_times, "phi",
-                               run_names, "dfns", nblocks, nz_local, nr_local)
+                               run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, Ez_at_pdf_times, "Ez",
-                               run_names, "dfns", nblocks, nz_local, nr_local)
+                               run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, Er_at_pdf_times, "Er",
-                               run_names, "dfns", nblocks, nz_local, nr_local)
+                               run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     # charged particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density_at_pdf_times, "density",
-                               run_names, "dfns", nblocks, nz_local, nr_local)
+                               run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_flow_at_pdf_times,
-                               "parallel_flow", run_names, "dfns", nblocks, nz_local,
-                               nr_local)
+                               "parallel_flow", run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_pressure_at_pdf_times,
                                "parallel_pressure", run_names, "dfns", nblocks,
-                               nz_local, nr_local)
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, parallel_heat_flux_at_pdf_times,
                                "parallel_heat_flux", run_names, "dfns", nblocks,
-                               nz_local, nr_local)
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     get_tuple_of_return_values(read_distributed_zr_data!, thermal_speed_at_pdf_times,
-                               "thermal_speed", run_names, "dfns", nblocks, nz_local,
-                               nr_local)
+                               "thermal_speed", run_names, "dfns", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     # neutral particle moments
     if any(n_neutral_species .> 0)
         get_tuple_of_return_values(read_distributed_zr_data!,
                                    neutral_density_at_pdf_times, "density_neutral",
-                                   run_names, "dfns", nblocks, nz_local, nr_local)
+                                   run_names, "dfns", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_uz_at_pdf_times,
-                                   "uz_neutral", run_names, "dfns", nblocks, nz_local,
-                                   nr_local)
+                                   "uz_neutral", run_names, "dfns", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_pz_at_pdf_times,
-                                   "pz_neutral", run_names, "dfns", nblocks, nz_local,
-                                   nr_local)
+                                   "pz_neutral", run_names, "dfns", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_qz_at_pdf_times,
-                                   "qz_neutral", run_names, "dfns", nblocks, nz_local,
-                                   nr_local)
+                                   "qz_neutral", run_names, "dfns", nblocks,
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
         get_tuple_of_return_values(read_distributed_zr_data!,
                                    neutral_thermal_speed_at_pdf_times,
                                    "thermal_speed_neutral", run_names, "dfns", nblocks,
-                                   nz_local, nr_local)
-    end
-    # load time data from `dfns' cdf
-    dfns_files0 = get_tuple_of_return_values(open_readonly_output_file, run_names, "dfns")
-    # note that ntime may differ in these output files
-
-    ntime_pdfs, time_pdfs = get_tuple_of_return_values(load_time_data, dfns_files0)
-
-    # load local velocity coordinate data from `dfns' cdf
-    # these values are currently the same for all blocks
-    nvpa, nvpa_global, vpa, vpa_wgts, Lvpa =
-        get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vpa")
-    nvperp, nvperp_global, vperp, vperp_wgts, Lvperp =
-        get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vperp")
-    if any(n_neutral_species .> 0)
-        nvzeta, nvzeta_global, vzeta, vzeta_wgts, Lvzeta =
-            get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vzeta")
-        nvr, nvr_global, vr, vr_wgts, Lvr =
-            get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vr")
-        nvz, nvz_global, vz, vz_wgts, Lvz =
-            get_tuple_of_return_values(load_coordinate_data, dfns_files0, "vz")
-    else # define nvz nvr nvzeta to avoid errors below
-        nvz = 0
-        nvr = 0
-        nvzeta = 0
+                                   Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
     end
 
     for f in dfns_files0
         close(f)
     end
 
-    geometry, composition =
-        get_tuple_of_return_values(get_geometry_and_composition, scan_input,
-                                   n_ion_species, n_neutral_species)
+    #geometry, composition =
+    #    get_tuple_of_return_values(get_geometry_and_composition, scan_input,
+    #                               n_ion_species, n_neutral_species)
 
     # initialise the post-processing input options
     nwrite_movie, itime_min, itime_max, nwrite_movie_pdfs, itime_min_pdfs, itime_max_pdfs,
     ivpa0, ivperp0, iz0, ir0, ivz0, ivr0, ivzeta0 =
-        init_postprocessing_options(pp, minimum(nvpa), minimum(nvperp),
-            minimum(nz_global), minimum(nr_global), minimum(nvz), minimum(nvr),
-            minimum(nvzeta), minimum(ntime), minimum(ntime_pdfs))
-
-    diagnostics_1d = false
+        init_postprocessing_options(pp, minimum(this_vpa.n_global for this_vpa ∈ vpa),
+            minimum(this_vperp.n_global for this_vperp ∈ vperp),
+            minimum(this_z.n_global for this_z ∈ z),
+            minimum(this_r.n_global for this_r ∈ r),
+            minimum(this_vz.n_global for this_vz ∈ vz),
+            minimum(this_vr.n_global for this_vr ∈ vr),
+            minimum(this_vzeta.n_global for this_vzeta ∈ vzeta),
+            minimum(ntime), minimum(ntime_pdfs))
+
+    diagnostics_1d = true
     if diagnostics_1d
         # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-        ff = get_tuple_of_return_values(allocate_global_zr_charged_dfns, nvpa_global,
-                                        nvperp_global, nz_global, nr_global,
+        ff = get_tuple_of_return_values(allocate_global_zr_charged_dfns,
+                                        Tuple(this_vpa.n_global for this_vpa ∈ vpa),
+                                        Tuple(this_vperp.n_global for this_vperp ∈ vperp),
+                                        Tuple(this_z.n_global for this_z ∈ z),
+                                        Tuple(this_r.n_global for this_r ∈ r),
                                         n_ion_species, ntime_pdfs)
         get_tuple_of_return_values(read_distributed_zr_data!, ff, "f", run_names, "dfns",
-                                   nblocks, nz_local, nr_local)
+                                   nblocks, Tuple(this_z.n for this_z ∈ z),
+                                   Tuple(this_r.n for this_r ∈ r))
 
         #evaluate 1D-1V diagnostics at fixed ir0
         plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
@@ -590,8 +648,7 @@ function analyze_and_plot_data(prefix...)
             Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed_at_pdf_times),
             Tuple(f[:,ivperp0,:,ir0,:,:] for f ∈ ff),
             n_ion_species, n_neutral_species, evolve_density, evolve_upar, evolve_ppar,
-            nvpa, vpa, vpa_wgts, nz_global, z, z_wgts, Lz, ntime, time, ntime_pdfs,
-            time_pdfs)
+            vpa, z, ntime, time, ntime_pdfs, time_pdfs)
     end
 
     diagnostics_2d = false
@@ -599,8 +656,7 @@ function analyze_and_plot_data(prefix...)
         # analyze the fields data
         phi_iz0 = Tuple(p[iz0,:,:] for p ∈ phi)
         phi_fldline_avg, delta_phi =
-            get_tuple_of_return_values(analyze_fields_data, phi_iz0, ntime, nr_global,
-                                       r_wgts, Lr)
+            get_tuple_of_return_values(analyze_fields_data, phi_iz0, ntime, r)
         get_tuple_of_return_values(plot_fields_rt, phi_iz0, delta_phi, time, itime_min,
                                    itime_max, nwrite_movie, r, ir0, run_name, delta_phi,
                                    pp)
@@ -609,8 +665,7 @@ function analyze_and_plot_data(prefix...)
     diagnostics_chodura = false
     if diagnostics_chodura
         Chodura_ratio_lower, Chodura_ratio_upper =
-            get_tuple_of_return_values(check_Chodura_condition, run_name, vpa_local,
-                                       vpa_local_wgts, vperp_local, vperp_local_wgts,
+            get_tuple_of_return_values(check_Chodura_condition, run_name, vpa, vperp,
                                        density_at_pdf_times, composition.T_e,
                                        Er_at_pdf_times, geometry, "wall", nblocks)
 
@@ -643,40 +698,35 @@ function analyze_and_plot_data(prefix...)
     phi = phi[1]
     Ez = Ez[1]
     Er = Er[1]
-    nvpa_global = nvpa_global[1]
-    nvperp_global = nvperp_global[1]
-    nz_global = nz_global[1]
-    nr_global = nr_global[1]
-    ntime = ntime[1]
-    nz_local = nz_local[1]
-    nr_local = nr_local[1]
+    vpa = vpa[1]
+    vperp = vperp[1]
     geometry = geometry[1]
     composition = composition[1]
 
     # make plots and animations of the phi, Ez and Er
-    plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time, z, r, iz0,
+    plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time, z.grid, r.grid, iz0,
                             ir0, n_ion_species, itime_min, itime_max, nwrite_movie,
                             run_name, pp)
     # make plots and animations of the phi, Ez and Er
-    plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0, itime_min, itime_max, nwrite_movie,
+    plot_fields_2D(phi, Ez, Er, time, z.grid, r.grid, iz0, ir0, itime_min, itime_max, nwrite_movie,
                    run_name, pp, "")
     # make plots and animations of the ion pdf
     # only if ntime == ntime_pdfs & data on one shared memory process
-    if ntime == ntime_pdfs && nr_global == nr_local && nz_global == nz_local
+    if ntime == ntime_pdfs && r.n_global == r.n && z.n_global == z.n
         # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-        ff = allocate_global_zr_charged_dfns(nvpa_global, nvperp_global, nz_global,
-                                             nr_global, n_ion_species, ntime)
-        read_distributed_zr_data!(ff, "f", run_names, "dfns", nblocks, nz_local, nr_local)
+        ff = allocate_global_zr_charged_dfns(vpa.n_global, vperp.n_global, z.n_global,
+                                             r.n_global, n_ion_species, ntime)
+        read_distributed_zr_data!(ff, "f", run_names, "dfns", nblocks, z.n, r.n)
 
         spec_type = "ion"
-        plot_charged_pdf(ff, vpa_local, vperp_local, z_local, r_local, ivpa0, ivperp0,
+        plot_charged_pdf(ff, vpa_local, vperp_local, z.grid, r.grid, ivpa0, ivperp0,
                          iz0, ir0, spec_type, n_ion_species, itime_min_pdfs,
                          itime_max_pdfs, nwrite_movie_pdfs, run_name, pp)
         # make plots and animations of the neutral pdf
         if n_neutral_species > 0
             spec_type = "neutral"
-            plot_neutral_pdf(neutral_ff, vz_local, vr_local, vzeta_local, z_local,
-                             r_local, ivz0, ivr0, ivzeta0, iz0, ir0, spec_type,
+            plot_neutral_pdf(neutral_ff, vz.grid, vr.grid, vzeta.grid, z.grid,
+                             r.grid, ivz0, ivr0, ivzeta0, iz0, ir0, spec_type,
                              n_neutral_species, itime_min_pdfs, itime_max_pdfs,
                              nwrite_movie_pdfs, run_name, pp)
         end
@@ -694,20 +744,20 @@ function analyze_and_plot_data(prefix...)
         r_bc = get(scan_input, "r_bc", "periodic")
         z_bc = get(scan_input, "z_bc", "periodic")
         # avoid passing Lr = 0 into manufactured_solns functions
-        if nr_local > 1
-            Lr_in = Lr
+        if r.n > 1
+            Lr_in = r.L
         else
             Lr_in = 1.0
         end
 
         #geometry, composition = get_geometry_and_composition(scan_input,n_ion_species,n_neutral_species)
 
-        manufactured_solns_list = manufactured_solutions(Lr_in,Lz,r_bc,z_bc,geometry,composition,nr_local)
+        manufactured_solns_list = manufactured_solutions(Lr_in,z.L,r_bc,z_bc,geometry,composition,r.n)
         dfni_func = manufactured_solns_list.dfni_func
         densi_func = manufactured_solns_list.densi_func
         dfnn_func = manufactured_solns_list.dfnn_func
         densn_func = manufactured_solns_list.densn_func
-        manufactured_E_fields = manufactured_electric_fields(Lr_in,Lz,r_bc,z_bc,composition,nr_local)
+        manufactured_E_fields = manufactured_electric_fields(Lr_in,z.L,r_bc,z_bc,composition,r.n)
         Er_func = manufactured_E_fields.Er_func
         Ez_func = manufactured_E_fields.Ez_func
         phi_func = manufactured_E_fields.phi_func
@@ -719,21 +769,21 @@ function analyze_and_plot_data(prefix...)
         for it in 1:ntime
             for ir in 1:nr_global
                 for iz in 1:nz_global
-                    phi_sym[iz,ir,it] = phi_func(z[iz],r[ir],time[it])
-                    Ez_sym[iz,ir,it] = Ez_func(z[iz],r[ir],time[it])
-                    Er_sym[iz,ir,it] = Er_func(z[iz],r[ir],time[it])
+                    phi_sym[iz,ir,it] = phi_func(z_grid[iz],r_grid[ir],time[it])
+                    Ez_sym[iz,ir,it] = Ez_func(z_grid[iz],r_grid[ir],time[it])
+                    Er_sym[iz,ir,it] = Er_func(z_grid[iz],r_grid[ir],time[it])
                 end
             end
         end
         # make plots and animations of the phi, Ez and Er
-        #plot_fields_2D(phi_sym, Ez_sym, Er_sym, time, z, r, iz0, ir0,
+        #plot_fields_2D(phi_sym, Ez_sym, Er_sym, time, z_grid, r_grid, iz0, ir0,
         #    itime_min, itime_max, nwrite_movie, run_name, pp, "_sym")
         println("time/ (Lref/cref): ", time)
-        compare_fields_symbolic_test(run_name,phi,phi_sym,z,r,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,phi,phi_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
          L"\widetilde{\phi}",L"\widetilde{\phi}^{sym}",L"\varepsilon(\widetilde{\phi})","phi")
-        compare_fields_symbolic_test(run_name,Er,Er_sym,z,r,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,Er,Er_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
          L"\widetilde{E}_r",L"\widetilde{E}^{sym}_r",L"\varepsilon(\widetilde{E}_r)","Er")
-        compare_fields_symbolic_test(run_name,Ez,Ez_sym,z,r,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,Ez,Ez_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
          L"\widetilde{E}_z",L"\widetilde{E}^{sym}_z",L"\varepsilon(\widetilde{E}_z)","Ez")
 
         # ion test
@@ -742,11 +792,11 @@ function analyze_and_plot_data(prefix...)
         for it in 1:ntime
             for ir in 1:nr_global
                 for iz in 1:nz_global
-                    density_sym[iz,ir,is,it] = densi_func(z[iz],r[ir],time[it])
+                    density_sym[iz,ir,is,it] = densi_func(z_grid[iz],r_grid[ir],time[it])
                 end
             end
         end
-        compare_moments_symbolic_test(run_name,density,density_sym,"ion",z,r,time,nz_global,nr_global,ntime,
+        compare_moments_symbolic_test(run_name,density,density_sym,"ion",z_grid,r_grid,time,nz_global,nr_global,ntime,
          L"\widetilde{n}_i",L"\widetilde{n}_i^{sym}",L"\varepsilon(\widetilde{n}_i)","dens")
 
         compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
@@ -758,11 +808,11 @@ function analyze_and_plot_data(prefix...)
             for it in 1:ntime
                 for ir in 1:nr_global
                     for iz in 1:nz_global
-                        neutral_density_sym[iz,ir,is,it] = densn_func(z[iz],r[ir],time[it])
+                        neutral_density_sym[iz,ir,is,it] = densn_func(z_grid[iz],r_grid[ir],time[it])
                     end
                 end
             end
-            compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z,r,time,nz_global,nr_global,ntime,
+            compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z_grid,r_grid,time,nz_global,nr_global,ntime,
              L"\widetilde{n}_n",L"\widetilde{n}_n^{sym}",L"\varepsilon(\widetilde{n}_n)","dens")
 
             compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,"neutral",
@@ -868,8 +918,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         phi_at_pdf_times, density_at_pdf_times, parallel_flow_at_pdf_times,
         parallel_pressure_at_pdf_times, parallel_heat_flux_at_pdf_times,
         thermal_speed_at_pdf_times, ff, n_ion_species, n_neutral_species, evolve_density,
-        evolve_upar, evolve_ppar, nvpa, vpa, vpa_wgts, nz, z, z_wgts, Lz, ntime, time,
-        ntime_pdfs, time_pdfs)
+        evolve_upar, evolve_ppar, vpa, z, ntime, time, ntime_pdfs, time_pdfs)
 
     n_runs = length(run_names)
 
@@ -879,7 +928,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
 
     # analyze the fields data
     phi_fldline_avg, delta_phi = get_tuple_of_return_values(analyze_fields_data, phi,
-                                                            ntime, nz, z_wgts, Lz)
+                                                            ntime, z)
 
     # use a fit to calculate and write to file the damping rate and growth rate of the
     # perturbed electrostatic potential
@@ -894,8 +943,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
     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 =
         get_tuple_of_return_values(analyze_moments_data, density, parallel_flow,
-            parallel_pressure, thermal_speed, parallel_heat_flux, ntime, n_ion_species,
-            nz, z_wgts, Lz)
+            parallel_pressure, thermal_speed, parallel_heat_flux, ntime, n_ion_species, z)
     # create the requested plots of the moments
     plot_moments(density, delta_density, density_fldline_avg,
         parallel_flow, delta_upar, upar_fldline_avg,
@@ -906,8 +954,8 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
     # load particle distribution function (pdf) data
     # analyze the pdf data
     f_fldline_avg, delta_f, dens_moment, upar_moment, ppar_moment =
-        get_tuple_of_return_values(analyze_pdf_data, ff, vpa, nvpa, nz, n_ion_species,
-            ntime_pdfs, vpa_wgts, z_wgts, Lz, thermal_speed, evolve_ppar)
+        get_tuple_of_return_values(analyze_pdf_data, ff, n_ion_species, ntime_pdfs, z,
+                                   vpa, thermal_speed_at_pdf_times, evolve_ppar)
 
     println("Plotting distribution function data...")
     cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
@@ -962,7 +1010,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             # make a gif animation of ln f(vpa,z,t)
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                subplots = (@views heatmap(this_z, this_vpa, log.(abs.(f[:,:,is,i])),
+                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
@@ -973,7 +1021,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             # make a gif animation of f(vpa,z,t)
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                subplots = (@views heatmap(this_z, this_vpa, f[:,:,is,i], xlabel="z",
+                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
@@ -983,9 +1031,9 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             gif(anim, outfile, fps=5)
             # make pdf of f(vpa,z,t_final) for each species
             str = string("spec ", string(is), " pdf")
-            subplots = (@views heatmap(this_z, this_vpa, f[:,:,is,end], xlabel="vpa", ylabel="z",
-                                       c = :deep, interpolation = :cubic,
-                                       title=string(run_label, str))
+            subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
             plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
             outfile = string(prefix, "_f_vs_z_vpa_final", spec_string, ".pdf")
@@ -994,7 +1042,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
                 for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa, f[:,1,is,i], xlabel="vpa", ylabel="f(z=0)",
+                    @views plot!(this_vpa.grid, f[:,1,is,i], xlabel="vpa", ylabel="f(z=0)",
                                  label=run_label)
                 end
             end
@@ -1004,7 +1052,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
                 for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa, f[:,end,is,i], xlabel="vpa", ylabel="f(z=L)",
+                    @views plot!(this_vpa.grid, f[:,end,is,i], xlabel="vpa", ylabel="f(z=L)",
                                  label=run_label)
                 end
             end
@@ -1024,7 +1072,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
 
                 PyPlot.subplot(1, n_runs, run_ind)
                 @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                    f[:,:,is,it], this_z, this_vpa, n[:,is,it],
+                    f[:,:,is,it], this_z.grid, this_vpa.grid, n[:,is,it],
                     upar[:,is,it], vth[:,is,it], ev_n, ev_u, ev_p)
                 plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; title=run_label,
                                       plot_log=false)
@@ -1042,7 +1090,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                                   run_names)
 
                 @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                    f[:,1,is,it], this_vpa, n[1,is,it], upar[1,is,it], vth[1,is,it],
+                    f[:,1,is,it], this_vpa.grid, n[1,is,it], upar[1,is,it], vth[1,is,it],
                     ev_n, ev_u, ev_p)
                 @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=0)",
                              label=run_label)
@@ -1059,7 +1107,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                                   run_names)
 
                 @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                    f[:,end,is,it], this_vpa, n[end,is,it], upar[end,is,it], vth[end,is,it],
+                    f[:,end,is,it], this_vpa.grid, n[end,is,it], upar[end,is,it], vth[end,is,it],
                     ev_n, ev_u, ev_p)
                 @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=L)",
                              label=run_label)
@@ -1071,7 +1119,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             ## use Plots.jl...
             ## make a gif animation of f_unnorm(v_parallel_unnorm,z,t)
             #anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa,
+            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid,
             #                           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,
@@ -1085,7 +1133,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             #gif(anim, outfile, fps=5)
             ## make a gif animation of log(f_unnorm)(v_parallel_unnorm,z,t)
             #anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z, this_vpa, n[:,is,i],
+            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid, 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,
@@ -1115,7 +1163,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
 
                     PyPlot.subplot(1, n_runs, run_ind)
                     @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                        f[:,:,is,iframe], this_z, this_vpa, n[:,is,iframe],
+                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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)
@@ -1139,7 +1187,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
 
                     PyPlot.subplot(1, n_runs, run_ind)
                     @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                        f[:,:,is,iframe], this_z, this_vpa, n[:,is,iframe],
+                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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)
@@ -1160,7 +1208,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                         thermal_speed_at_pdf_times, evolve_density, evolve_upar,
                         evolve_ppar, vpa, run_names)
                     @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],
+                        f[:,1,is,i], this_vpa.grid, 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)
@@ -1176,7 +1224,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                         thermal_speed_at_pdf_times, evolve_density, evolve_upar,
                         evolve_ppar, vpa, run_names)
                     @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                        f[:,end,is,i], this_vpa, n[end,is,i], upar[end,is,i],
+                        f[:,end,is,i], this_vpa.grid, 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)
@@ -1188,7 +1236,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         if pp.animate_deltaf_vs_vpa_z
             # make a gif animation of δf(vpa,z,t)
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                subplots = (@views heatmap(this_z, this_vpa, delta_f[:,:,is,i],
+                subplots = (@views heatmap(this_z, this_vpa.grid, delta_f[:,:,is,i],
                                        xlabel="z", ylabel="vpa", c = :deep,
                                        interpolation = :cubic, title=run_label)
                             for (df, this_z, this_vpa, run_label) ∈
@@ -1257,7 +1305,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                 #@views plot(vpa, ff[iz0,:,is,i], ylims = (fmin,fmax))
                 plot(legend=legend)
                 for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa, f[:,iz0,is,i], label=run_label)
+                    @views plot!(this_vpa.grid, f[:,iz0,is,i], label=run_label)
                 end
             end
             outfile = string(prefix, "_f_vs_vpa", spec_string, ".gif")
@@ -1271,7 +1319,7 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
                 plot(legend=legend)
                 for (df, this_vpa, fn, fx, run_label) ∈
                         zip(delta_f, vpa, fmin, fmax, run_names)
-                    @views plot!(this_vpa, delta_f[:,iz0,is,i], ylims = (fn,fx),
+                    @views plot!(this_vpa.grid, delta_f[:,iz0,is,i], ylims = (fn,fx),
                                 label=run_label)
                 end
             end
@@ -1372,8 +1420,8 @@ function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
     end
     if pp.plot_phi_vs_z_t
         # make a heatmap plot of ϕ(z,t)
-        subplots = (heatmap(t, this_z, p, xlabel="time", ylabel="z", title=run_label, c =
-                            :deep)
+        subplots = (heatmap(t, this_z.grid, p, xlabel="time", ylabel="z", title=run_label,
+                            c = :deep)
                     for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names))
         plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
         outfile = string(prefix, "_phi_vs_z_t.pdf")
@@ -1384,7 +1432,7 @@ function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             plot(legend=legend)
             for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
-                @views plot!(this_z, p[:,i], xlabel="z", ylabel="ϕ",
+                @views plot!(this_z.grid, p[:,i], xlabel="z", ylabel="ϕ",
                              ylims=(phimin, phimax), label=run_label)
             end
         end
@@ -1399,8 +1447,8 @@ function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
     # savefig(outfile)
     plot(legend=legend)
     for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
-        plot!(this_z, p[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label=run_label,
-             linewidth=2)
+        plot!(this_z.grid, p[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label=run_label,
+              linewidth=2)
     end
     outfile = string(prefix, "_phi_final.pdf")
     savefig(outfile)
@@ -1586,7 +1634,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
         end
         if pp.plot_dens_vs_z_t
             # make a heatmap plot of n_s(z,t)
-            subplots = (heatmap(t, this_z, n[:,is,:], xlabel="time", ylabel="z",
+            subplots = (heatmap(t, this_z.grid, n[:,is,:], xlabel="time", ylabel="z",
                                 title=run_label, c = :deep)
                         for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names))
             plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
@@ -1595,7 +1643,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
         end
         if pp.plot_upar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z, upar[:,is,:], xlabel="time", ylabel="z",
+            subplots = (heatmap(t, this_z.grid, upar[:,is,:], xlabel="time", ylabel="z",
                                 title=run_label, c = :deep)
                         for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow,
                                                               run_names))
@@ -1605,7 +1653,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
         end
         if pp.plot_ppar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z, ppar[:,is,:], xlabel="time", ylabel="z",
+            subplots = (heatmap(t, this_z.grid, ppar[:,is,:], xlabel="time", ylabel="z",
                                 title=run_label, c = :deep)
                         for (t, this_z, ppar, run_label) ∈
                             zip(time, z, parallel_pressure, run_names))
@@ -1615,7 +1663,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
         end
         if pp.plot_qpar_vs_z_t
             # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z, qpar[:,is,:], xlabel="time", ylabel="z",
+            subplots = (heatmap(t, this_z.grid, 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_names))
@@ -1628,7 +1676,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 plot(legend=legend)
                 for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names)
-                    @views plot!(this_z, n[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ",
+                    @views plot!(this_z.grid, n[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ",
                                  ylims=(dens_min, dens_max), label=run_label)
                 end
             end
@@ -1640,7 +1688,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 plot(legend=legend)
                 for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow, run_names)
-                    @views plot!(this_z, upar[:,is,i], xlabel="z", ylabel="upars/vt",
+                    @views plot!(this_z.grid, upar[:,is,i], xlabel="z", ylabel="upars/vt",
                                  ylims=(upar_min, upar_max), label=run_label)
                 end
             end
@@ -1652,7 +1700,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 plot(legend=legend)
                 for (t, this_z, ppar, run_label) ∈ zip(time, z, parallel_pressure, run_names)
-                    @views plot!(this_z, ppar[:,is,i], xlabel="z", ylabel="ppars",
+                    @views plot!(this_z.grid, ppar[:,is,i], xlabel="z", ylabel="ppars",
                                  ylims=(ppar_min, ppar_max), label=run_label)
                 end
             end
@@ -1664,7 +1712,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
             anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                 plot(legend=legend)
                 for (t, this_z, vth, run_label) ∈ zip(time, z, thermal_speed, run_names)
-                    @views plot!(this_z, vth[:,is,i], xlabel="z", ylabel="vths",
+                    @views plot!(this_z.grid, vth[:,is,i], xlabel="z", ylabel="vths",
                                  ylims=(vth_min, vth_max), label=run_label)
                 end
             end
@@ -1677,7 +1725,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                 plot(legend=legend)
                 for (t, this_z, qpar, run_label) ∈ zip(time, z, parallel_heat_flux,
                                                        run_names)
-                    @views plot!(this_z, qpar[:,is,i], xlabel="z", ylabel="qpars",
+                    @views plot!(this_z.grid, qpar[:,is,i], xlabel="z", ylabel="qpars",
                                  ylims=(qpar_min, qpar_max), label=run_label)
                 end
             end
@@ -2117,11 +2165,11 @@ function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec
     # load block data on iblock=0
     nblocks, iblock = load_block_data(fid, printout=false)
     # load global sizes of grids that are distributed in memory
-    nz_local, nz_global, z_local, z_wgts_local, Lz = load_coordinate_data(fid, "z", printout=false)
-    nr_local, nr_global, r_local, r_wgts_local, Lr = load_coordinate_data(fid, "r", printout=false)
+    z0, _ = load_coordinate_data(fid, "z", printout=false)
+    r0, _ = load_coordinate_data(fid, "r", printout=false)
     # velocity grid data on iblock=0 (same for all blocks)
-    nvpa, _, vpa, vpa_wgts, Lvpa = load_coordinate_data(fid, "vpa", printout=false)
-    nvperp, _, vperp, vperp_wgts, Lvperp = load_coordinate_data(fid, "vperp", printout=false)
+    vpa, _ = load_coordinate_data(fid, "vpa", printout=false)
+    vperp, _ = load_coordinate_data(fid, "vperp", printout=false)
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time = load_time_data(fid, printout=false)
     close(fid)
@@ -2135,20 +2183,20 @@ function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec
         z_irank, r_irank = load_rank_data(fid_pdfs,printout=false)
         pdf = load_pdf_data(fid_pdfs, printout=false)
         # local local grid data on iblock=0
-        nz_local, _, z_local, z_wgts_local, Lz = load_coordinate_data(fid_pdfs, "z")
-        nr_local, _, r_local, r_wgts_local, Lr = load_coordinate_data(fid_pdfs, "r")
+        z_local, _ = load_coordinate_data(fid_pdfs, "z")
+        r_local, _ = load_coordinate_data(fid_pdfs, "r")
         close(fid_pdfs)
         imin_r = min(1,r_irank) + 1
         imin_z = min(1,z_irank) + 1
-        for it in 1:ntime, ir in imin_r:nr_local, iz in imin_z:nz_local,
-                ivperp in 1:nvperp, ivpa in 1:nvpa
+        for it in 1:ntime, ir in imin_r:r_local.n, iz in imin_z:z_local.n,
+                ivperp in 1:vperp.n, ivpa in 1:vpa.n
 
-            pdf_sym = dfni_func(vpa[ivpa],vperp[ivperp],z_local[iz],r_local[ir],time[it])
+            pdf_sym = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z_local.grid[iz],r_local.grid[ir],time[it])
             pdf_norm[it] += (pdf[ivpa,ivperp,iz,ir,is,it] - pdf_sym)^2
         end
     end
     for it in 1:ntime
-        pdf_norm[it] = sqrt(pdf_norm[it]/(nr_global*nz_global*nvpa*nvperp))
+        pdf_norm[it] = sqrt(pdf_norm[it]/(r0.n_global*z0.n_global*vpa.n*vperp.n))
     end
     println("test: ",file_string,": ",spec_string," ",pdf_norm)
     @views plot(time, pdf_norm[:], xlabel=L"t L_z/v_{ti}", ylabel=norm_label) #, yaxis=:log)
@@ -2164,11 +2212,11 @@ function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec
         if (z_irank == 0 || z_irank == z_nrank - 1) && r_irank == 0
             pdf = load_pdf_data(fid_pdfs, printout=false)
             # local local grid data on iblock=0
-            nz_local, _, z_local, z_wgts_local, Lz = load_coordinate_data(fid_pdfs, "z")
-            nr_local, _, r_local, r_wgts_local, Lr = load_coordinate_data(fid_pdfs, "r")
-            pdf_sym_array = copy(vpa)
+            z_local, _ = load_coordinate_data(fid_pdfs, "z")
+            r_local = load_coordinate_data(fid_pdfs, "r")
+            pdf_sym_array = copy(vpa.grid)
             # plot a thermal vperp on line plots
-            ivperp0 = max(floor(mk_int,nvperp/3),1)
+            ivperp0 = max(floor(mk_int,vperp.n/3),1)
             # plot a typical r on line plots
             ir0 = 1
             # plot at the wall boundary
@@ -2176,14 +2224,14 @@ function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec
                 iz0 = 1
                 zlabel="wall-"
             elseif z_irank == z_nrank - 1
-                iz0 = nz_local
+                iz0 = z_local.n
                 zlabel="wall+"
             end
-            for ivpa in 1:nvpa
-                pdf_sym_array[ivpa] = dfni_func(vpa[ivpa],vperp[ivperp0],z_local[iz0],r_local[ir0],time[ntime])
+            for ivpa in 1:vpa.n
+                pdf_sym_array[ivpa] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp0],z_local.grid[iz0],r_local.grid[ir0],time[ntime])
             end
             # plot f(vpa,ivperp0,iz_wall,ir0,is,itime) at the wall
-            @views plot(vpa, [pdf[:,ivperp0,iz0,ir0,is,ntime], pdf_sym_array], xlabel=L"v_{\|\|}/L_{v_{\|\|}}", ylabel=L"f_i", label=["num" "sym"])
+            @views plot(vpa.grid, [pdf[:,ivperp0,iz0,ir0,is,ntime], pdf_sym_array], xlabel=L"v_{\|\|}/L_{v_{\|\|}}", ylabel=L"f_i", label=["num" "sym"])
             outfile = string(run_name, "_pdf(vpa,vperp0,iz_"*zlabel*",ir0)_sym_vs_vpa.pdf")
             savefig(outfile)
         end
@@ -2198,12 +2246,12 @@ function compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,spec
     # load block data on iblock=0
     nblocks, iblock = load_block_data(fid, printout=false)
     # load global sizes of grids that are distributed in memory
-    nz_local, nz_global, z_local, z_wgts_local, Lz = load_coordinate_data(fid, "z", printout=false)
-    nr_local, nr_global, r_local, r_wgts_local, Lr = load_coordinate_data(fid, "r", printout=false)
+    z0, _ = load_coordinate_data(fid, "z", printout=false)
+    r0, _ = load_coordinate_data(fid, "r", printout=false)
     # velocity grid data on iblock=0 (same for all blocks)
-    nvz, _, vz, vz_wgts, Lvz = load_coordinate_data(fid, "vz", printout=false)
-    nvr, _, vr, vr_wgts, Lvr = load_coordinate_data(fid, "vr", printout=false)
-    nvzeta, _, vzeta, vzeta_wgts, Lvzeta = load_coordinate_data(fid, "vzeta", printout=false)
+    vz, _ = load_coordinate_data(fid, "vz", printout=false)
+    vr, _ = load_coordinate_data(fid, "vr", printout=false)
+    vzeta, _ = load_coordinate_data(fid, "vzeta", printout=false)
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time = load_time_data(fid, printout=false)
     close(fid)
@@ -2217,20 +2265,20 @@ function compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,spec
         z_irank, r_irank = load_rank_data(fid_pdfs,printout=false)
         pdf = load_neutral_pdf_data(fid_pdfs, printout=false)
         # load local grid data
-        nz_local, _, z_local, z_wgts_local, Lz = load_coordinate_data(fid_pdfs, "z", printout=false)
-        nr_local, _, r_local, r_wgts_local, Lr = load_coordinate_data(fid_pdfs, "r", printout=false)
+        z_local, _ = load_coordinate_data(fid_pdfs, "z", printout=false)
+        r_local, _ = load_coordinate_data(fid_pdfs, "r", printout=false)
         close(fid_pdfs)
         imin_r = min(1,r_irank) + 1
         imin_z = min(1,z_irank) + 1
-        for it in 1:ntime, ir in imin_r:nr_local, iz in imin_z:nz_local,
-                ivzeta in 1:nvzeta, ivr in 1:nvr, ivz in 1:nvz
+        for it in 1:ntime, ir in imin_r:r_local.n, iz in imin_z:z_local.n,
+                ivzeta in 1:vzeta.n, ivr in 1:vr.n, ivz in 1:vz.n
 
-            pdf_sym = dfnn_func(vz[ivz],vr[ivr],vzeta[ivzeta],z_local[iz],r_local[ir],time[it])
+            pdf_sym = dfnn_func(vz.grid[ivz],vr.grid[ivr],vzeta.grid[ivzeta],z_local.grid[iz],r_local.grid[ir],time[it])
             pdf_norm[it] += (pdf[ivz,ivr,ivzeta,iz,ir,is,it] - pdf_sym)^2
         end
     end
     for it in 1:ntime
-        pdf_norm[it] = sqrt(pdf_norm[it]/(nr_global*nz_global*nvz*nvr*nvzeta))
+        pdf_norm[it] = sqrt(pdf_norm[it]/(r0.n_global*z0.n_global*vz.n*vr.n*vzeta.n))
     end
     println("test: ",file_string,": ",spec_string," ",pdf_norm)
     @views plot(time, pdf_norm[:], xlabel=L"t L_z/v_{ti}", ylabel=norm_label) #, yaxis=:log)
@@ -2608,8 +2656,8 @@ function plot_charged_pdf_2D_at_wall(run_name)
     # load block data on iblock=0
     nblocks, iblock = load_block_data(fid, printout=false)
     # velocity grid data on iblock=0 (same for all blocks)
-    nvpa, _, vpa, vpa_wgts, Lvpa = load_coordinate_data(fid, "vpa", printout=false)
-    nvperp, _, vperp, vperp_wgts, Lvperp = load_coordinate_data(fid, "vperp", printout=false)
+    vpa, _ = load_coordinate_data(fid, "vpa", printout=false)
+    vperp, _ = load_coordinate_data(fid, "vperp", printout=false)
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time = load_time_data(fid, printout=false)
     # load species data
@@ -2619,9 +2667,9 @@ function plot_charged_pdf_2D_at_wall(run_name)
     # plot only data at last timestep
     itime0 = ntime
     # plot a thermal vpa on line plots
-    ivpa0 = floor(mk_int,nvpa/3)
+    ivpa0 = floor(mk_int,vpa.n/3)
     # plot a thermal vperp on line plots
-    ivperp0 = max(1,floor(mk_int,nvperp/3))
+    ivperp0 = max(1,floor(mk_int,vperp.n/3))
     # plot a typical r on line plots
     ir0 = 1
 
@@ -2635,16 +2683,16 @@ function plot_charged_pdf_2D_at_wall(run_name)
         if (z_irank == 0 || z_irank == z_nrank-1) && r_irank == 0
             # plot data from lower wall boundary near z = -L/2
             # load local grid data
-            nz_local, _, z_local, z_wgts_local, Lz = load_coordinate_data(fid_pdfs, "z", printout=false)
-            nr_local, _, r_local, r_wgts_local, Lr = load_coordinate_data(fid_pdfs, "r", printout=false)
+            z, _ = load_coordinate_data(fid_pdfs, "z", printout=false)
+            r, _ = load_coordinate_data(fid_pdfs, "r", printout=false)
 
             if z_irank == 0
                 iz_wall = 1
-                #print("z_local[iz_wall-]: ",z_local[iz_wall])
+                #print("z.grid[iz_wall-]: ",z.grid[iz_wall])
                 zlabel = "wall-"
             elseif z_irank == z_nrank-1
-                iz_wall = nz_local
-                #print("z_local[iz_wall+]: ",z_local[iz_wall])
+                iz_wall = z.n
+                #print("z.grid[iz_wall+]: ",z.grid[iz_wall])
                 zlabel = "wall+"
             end
             # load local pdf data
@@ -2653,29 +2701,29 @@ function plot_charged_pdf_2D_at_wall(run_name)
                 description = "_ion_spec"*string(is)*"_"
 
                 # plot f(vpa,ivperp0,iz_wall,ir0,is,itime) at the wall
-                @views plot(vpa, pdf[:,ivperp0,iz_wall,ir0,is,itime0], xlabel=L"v_{\|\|}/L_{v_{\|\|}}", ylabel=L"f_i")
+                @views plot(vpa.grid, pdf[:,ivperp0,iz_wall,ir0,is,itime0], xlabel=L"v_{\|\|}/L_{v_{\|\|}}", ylabel=L"f_i")
                 outfile = string(run_name, "_pdf(vpa,vperp0,iz_"*zlabel*",ir0)"*description*"_vs_vpa.pdf")
                 savefig(outfile)
 
                 # plot f(vpa,vperp,iz_wall,ir0,is,itime) at the wall
-                @views heatmap(vperp, vpa, pdf[:,:,iz_wall,ir0,is,itime0], xlabel=L"v_{\perp}", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
+                @views heatmap(vperp.grid, vpa.grid, pdf[:,:,iz_wall,ir0,is,itime0], xlabel=L"v_{\perp}", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
                 windowsize = (360,240), margin = 15pt)
                 outfile = string(run_name, "_pdf(vpa,vperp,iz_"*zlabel*",ir0)"*description*"_vs_vperp_vpa.pdf")
                 savefig(outfile)
 
                 # plot f(vpa,ivperp0,z,ir0,is,itime) near the wall
-                @views heatmap(z_local, vpa, pdf[:,ivperp0,:,ir0,is,itime0], xlabel=L"z", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
+                @views heatmap(z.grid, vpa.grid, pdf[:,ivperp0,:,ir0,is,itime0], xlabel=L"z", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
                 windowsize = (360,240), margin = 15pt)
                 outfile = string(run_name, "_pdf(vpa,ivperp0,z_"*zlabel*",ir0)"*description*"_vs_z_vpa.pdf")
                 savefig(outfile)
 
                 # plot f(ivpa0,ivperp0,z,r,is,itime) near the wall
-                if nr_local > 1
-                    @views heatmap(r_local, z_local, pdf[ivpa0,ivperp0,:,:,is,itime0], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
+                if r.n > 1
+                    @views heatmap(r.grid, z.grid, pdf[ivpa0,ivperp0,:,:,is,itime0], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
                     windowsize = (360,240), margin = 15pt)
                     outfile = string(run_name, "_pdf(ivpa0,ivperp0,z_"*zlabel*",r)"*description*"_vs_r_z.pdf")
                     savefig(outfile)
-                    @views heatmap(r_local, vpa, pdf[:,ivperp0,iz_wall,:,is,itime0], xlabel=L"r", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
+                    @views heatmap(r_local, vpa.grid, pdf[:,ivperp0,iz_wall,:,is,itime0], xlabel=L"r", ylabel=L"v_{||}", c = :deep, interpolation = :cubic,
                     windowsize = (360,240), margin = 15pt)
                     outfile = string(run_name, "_pdf(vpa,ivperp0,z_"*zlabel*",r)"*description*"_vs_r_vpa.pdf")
                     savefig(outfile)
diff --git a/test/harrisonthompson.jl b/test/harrisonthompson.jl
index e80a68b93..fcb0f0825 100644
--- a/test/harrisonthompson.jl
+++ b/test/harrisonthompson.jl
@@ -184,8 +184,8 @@ function run_test(test_input, analytic_rtol, analytic_atol, expected_phi,
             fid = open_readonly_output_file(path,"moments")
 
             # load space-time coordinate data
-            nz, nz_global, z, z_wgts, Lz = load_coordinate_data(fid, "z")
-            nr, nr_global, r, r_wgts, Lr = load_coordinate_data(fid, "r")
+            z, z_spectral = load_coordinate_data(fid, "z")
+            r, r_spectral = load_coordinate_data(fid, "r")
             ntime, time = load_time_data(fid)
             n_ion_species, n_neutral_species = load_species_data(fid)
             
@@ -196,12 +196,11 @@ function run_test(test_input, analytic_rtol, analytic_atol, expected_phi,
 
             phi = phi_zrt[:,1,:]
             
-            analytic_phi = [findphi(zval, input["ionization_frequency"]) for zval ∈ z]
+            analytic_phi = [findphi(zval, input["ionization_frequency"]) for zval ∈ z.grid]
         end
 
-        nz = length(z)
         # Analytic solution defines phi=0 at mid-point, so need to offset the code solution
-        offset = phi[(nz+1)÷2, end]
+        offset = phi[(z.n+1)÷2, end]
         # Error is large on the boundary points, so test those separately
         @test isapprox(phi[2:end-1, end] .- offset, analytic_phi[2:end-1],
                        rtol=analytic_rtol, atol=analytic_atol)
diff --git a/test/nonlinear_sound_wave_tests.jl b/test/nonlinear_sound_wave_tests.jl
index e52ce71bc..ca0ff8105 100644
--- a/test/nonlinear_sound_wave_tests.jl
+++ b/test/nonlinear_sound_wave_tests.jl
@@ -284,9 +284,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # open the netcdf file containing moments data and give it the handle 'fid'
             fid = open_readonly_output_file(path, "moments")
 
-            # load space-time coordinate data
-            nz, nz_global, z, z_wgts, Lz = load_coordinate_data(fid, "z")
-            nr, nr_global, r, r_wgts, Lr = load_coordinate_data(fid, "r")
+            # load species, time coordinate data
             n_ion_species, n_neutral_species = load_species_data(fid)
             ntime, time = load_time_data(fid)
             n_ion_species, n_neutral_species = load_species_data(fid)
@@ -295,8 +293,9 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, qpar_charged_zrst, v_t_charged_zrst, evolve_ppar = load_charged_particle_moments_data(fid)
+            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, qpar_charged_zrst, v_t_charged_zrst = load_charged_particle_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
+            z, z_spectral = load_coordinate_data(fid, "z")
 
             close(fid)
             
@@ -325,18 +324,18 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
 
             # Unnormalize f
             if input["evolve_moments_density"]
-                for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:nz
+                for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
                     f_charged[:,iz,is,it] .*= n_charged[iz,is,it]
                 end
-                for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:nz
+                for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] .*= n_neutral[iz,isn,it]
                 end
             end
             if input["evolve_moments_parallel_pressure"]
-                for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:nz
+                for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
                     f_charged[:,iz,is,it] ./= v_t_charged[iz,is,it]
                 end
-                for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:nz
+                for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] ./= v_t_neutral[iz,isn,it]
                 end
             end
diff --git a/test/sound_wave_tests.jl b/test/sound_wave_tests.jl
index 23f9e3732..d6f12bc81 100644
--- a/test/sound_wave_tests.jl
+++ b/test/sound_wave_tests.jl
@@ -181,8 +181,8 @@ function run_test(test_input, analytic_frequency, analytic_growth_rate,
             fid = open_readonly_output_file(path,"moments")
 
             # load space-time coordinate data
-            nz, nz_global, z, z_wgts, Lz = load_coordinate_data(fid, "z")
-            nr, nr_global, r, r_wgts, Lr = load_coordinate_data(fid, "r")
+            z, z_spectral = load_coordinate_data(fid, "z")
+            r, r_spectral = load_coordinate_data(fid, "r")
             n_ion_species, n_neutral_species = load_species_data(fid)
             ntime, time = load_time_data(fid)
             
@@ -196,25 +196,25 @@ function run_test(test_input, analytic_frequency, analytic_growth_rate,
             phi = phi_zrt[:,ir0,:]
 
             # analyze the fields data
-            phi_fldline_avg, delta_phi = analyze_fields_data(phi, ntime, nz, z_wgts, Lz)
+            phi_fldline_avg, delta_phi = analyze_fields_data(phi, ntime, z)
 
             # use a fit to calculate the damping rate and growth rate of the perturbed
             # electrostatic potential
             itime_max = ntime
-            iz0 = cld(nz, 3)
+            iz0 = cld(z.n, 3)
             shifted_time = allocate_float(ntime)
             @. shifted_time = time - time[itime_min]
-            @views phi_fit = fit_delta_phi_mode(shifted_time[itime_min:itime_max], z,
+            @views phi_fit = fit_delta_phi_mode(shifted_time[itime_min:itime_max], z.grid,
                                                 delta_phi[:, itime_min:itime_max])
             ## The following plot code (copied from post_processing.jl) may be helpful for
             ## debugging tests. Uncomment to use, and also uncomment
             ## `using Plots: plot, plot!, gui at the top of the file.
-            #L = z[end] - z[begin]
+            #L = z.grid[end] - z.grid[begin]
             #fitted_delta_phi =
-            #    @. (phi_fit.amplitude0 * cos(2.0 * π * (z[iz0] + phi_fit.offset0) / L)
+            #    @. (phi_fit.amplitude0 * cos(2.0 * π * (z.grid[iz0] + phi_fit.offset0) / L)
             #        * exp(phi_fit.growth_rate * shifted_time)
             #        * cos(phi_fit.frequency * shifted_time + phi_fit.phase))
-            #@views plot(time, abs.(delta_phi[iz0,:]), xlabel="t*Lz/vti", ylabel="δϕ", yaxis=:log)
+            #@views plot(time, abs.(delta_phi[iz0,:]), xlabel="t*z.L/vti", ylabel="δϕ", yaxis=:log)
             #plot!(time, abs.(fitted_delta_phi))
             #gui()
         end
diff --git a/test/wall_bc_tests.jl b/test/wall_bc_tests.jl
index 721fd431d..9a5708921 100644
--- a/test/wall_bc_tests.jl
+++ b/test/wall_bc_tests.jl
@@ -15,7 +15,7 @@ using moment_kinetics.interpolation: interpolate_to_grid_z
 using moment_kinetics.load_data: open_readonly_output_file
 using moment_kinetics.load_data: load_fields_data,
                                  load_pdf_data, load_time_data,
-                                 load_species_data, load_coordinate_data
+                                 load_species_data
 
 # Create a temporary directory for test output
 test_output_directory = tempname()
@@ -161,9 +161,7 @@ function run_test(test_input, expected_phi, tolerance; args...)
             # open the netcdf file and give it the handle 'fid'
             fid = open_readonly_output_file(path,"moments")
 
-            # load space-time coordinate data
-            nz, nz_global, z, z_wgts, Lz = load_coordinate_data(fid, "z")
-            nr, nr_global, r, r_wgts, Lr = load_coordinate_data(fid, "r")
+            # load species, time coordinate data
             n_ion_species, n_neutral_species = load_species_data(fid)
             ntime, time = load_time_data(fid)
             n_ion_species, n_neutral_species = load_species_data(fid)

From 04dcb430fe9dd9a170b26ad7dd9600000a79a2f6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 May 2023 20:01:44 +0100
Subject: [PATCH 017/428] Update filename for distributed, multiple-file mode

Swap order of the ".dfns" or ".moments" suffix and the ".<iblock>" part
(which is present when multiple files are being written, but not when
parallel I/O is being used). This makes it easier to extract the
'iblock' component from the filename, as we know that it is always
preceded by either ".dfns" or ".moments", whereas the user defined
prefix might end in something like ".2" by chance, which could be
mistaken for an 'iblock' suffix in the naming scheme before this update.
---
 src/file_io.jl   | 20 +++++++++++---------
 src/load_data.jl |  4 ++--
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index 795ea6899..32b9f121f 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -110,11 +110,7 @@ function setup_file_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition, c
         # check to see if output_dir exists in the current directory
         # if not, create it
         isdir(io_input.output_dir) || mkdir(io_input.output_dir)
-        if io_input.parallel_io
-            out_prefix = string(io_input.output_dir, "/", io_input.run_name)
-        else
-            out_prefix = string(io_input.output_dir, "/", io_input.run_name, ".", iblock_index[])
-        end
+        out_prefix = string(io_input.output_dir, "/", io_input.run_name)
 
         if io_input.ascii_output
             #ff_io = open_ascii_output_file(out_prefix, "f_vs_t")
@@ -511,8 +507,11 @@ setup file i/o for moment variables
 function setup_moments_io(prefix, binary_format, r, z, composition, collisions,
                           evolve_density, evolve_upar, evolve_ppar, parallel_io, io_comm)
     @serial_region begin
-        fid = open_output_file(string(prefix, ".moments"), binary_format, parallel_io,
-                               io_comm)
+        moments_prefix = string(prefix, ".moments")
+        if !parallel_io
+            moments_prefix *= ".$(iblock_index[])"
+        end
+        fid = open_output_file(moments_prefix, binary_format, parallel_io, io_comm)
 
         # write a header to the output file
         add_attribute!(fid, "file_info", "Output moments data from the moment_kinetics code")
@@ -544,8 +543,11 @@ function setup_dfns_io(prefix, binary_format, r, z, vperp, vpa, vzeta, vr, vz, c
                        io_comm)
 
     @serial_region begin
-        fid = open_output_file(string(prefix, ".dfns"), binary_format, parallel_io,
-                               io_comm)
+        dfns_prefix = string(prefix, ".dfns")
+        if !parallel_io
+            dfns_prefix *= ".$(iblock_index[])"
+        end
+        fid = open_output_file(dfns_prefix, binary_format, parallel_io, io_comm)
 
         # write a header to the output file
         add_attribute!(fid, "file_info",
diff --git a/src/load_data.jl b/src/load_data.jl
index 3d7b88e81..77fa24757 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -27,9 +27,9 @@ using NCDatasets
 function open_readonly_output_file(run_name, ext; iblock=0, printout=false)
     possible_names = (
         string(run_name, ".", ext, ".h5"),
-        string(run_name, ".", iblock,".", ext, ".h5"),
+        string(run_name, ".", ext, ".", iblock, ".h5"),
         string(run_name, ".", ext, ".cdf"),
-        string(run_name, ".", iblock,".", ext, ".cdf"),
+        string(run_name, ".", ext, ".", iblock, ".cdf"),
     )
     existing_files = Tuple(f for f in possible_names if isfile(f))
     exists_count = length(existing_files)

From c5d9c1b31d0a528b637b5b504aa45868d1520f01 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 20 May 2023 13:48:34 +0100
Subject: [PATCH 018/428] Update restart functionality to work with new I/O
 implementation

---
 src/file_io.jl         | 102 +++++++++----------
 src/file_io_hdf5.jl    |  18 +++-
 src/file_io_netcdf.jl  |  16 ++-
 src/load_data.jl       | 225 +++++++++++++++++++++++++++++++++++++++++
 src/moment_kinetics.jl | 111 ++++++++++++++------
 5 files changed, 382 insertions(+), 90 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index 32b9f121f..bc97d064a 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -101,8 +101,8 @@ function io_has_parallel() end
 """
 open the necessary output files
 """
-function setup_file_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition, collisions,
-                       evolve_density, evolve_upar, evolve_ppar)
+function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vperp, z, r,
+                       composition, collisions, evolve_density, evolve_upar, evolve_ppar)
     begin_serial_region()
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -126,10 +126,10 @@ function setup_file_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition, c
                                       composition, collisions, evolve_density,
                                       evolve_upar, evolve_ppar, io_input.parallel_io,
                                       comm_inter_block[])
-        io_dfns = setup_dfns_io(out_prefix, io_input.binary_format, r, z, vperp, vpa,
-                                vzeta, vr, vz, composition, collisions, evolve_density,
-                                evolve_upar, evolve_ppar, io_input.parallel_io,
-                                comm_inter_block[])
+        io_dfns = setup_dfns_io(out_prefix, io_input.binary_format,
+                                boundary_distributions, r, z, vperp, vpa, vzeta, vr, vz,
+                                composition, collisions, evolve_density, evolve_upar,
+                                evolve_ppar, io_input.parallel_io, comm_inter_block[])
 
         return ascii, io_moments, io_dfns
     end
@@ -183,9 +183,40 @@ function write_overview!(fid, composition, collisions, parallel_io, evolve_densi
         write_single_value!(overview, "evolve_ppar", evolve_ppar,
                             parallel_io=parallel_io,
                             description="is parallel pressure evolved separately from the distribution function?")
+        write_single_value!(overview, "parallel_io", parallel_io,
+                            parallel_io=parallel_io,
+                            description="is parallel I/O being used?")
     end
 end
 
+"""
+Write the distributions that may be used for boundary conditions to the output file
+"""
+function write_boundary_distributions!(fid, boundary_distributions, parallel_io,
+                                       composition, z, vperp, vpa, vzeta, vr, vz)
+    @serial_region begin
+        boundary_distributions_io = create_io_group(fid, "boundary_distributions")
+
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_left",
+            boundary_distributions.pdf_rboundary_charged[:,:,:,1,:], vpa, vperp, z,
+            parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
+            description="Initial charged-particle pdf at left radial boundary")
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_right",
+            boundary_distributions.pdf_rboundary_charged[:,:,:,2,:], vpa, vperp, z,
+            parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
+            description="Initial charged-particle pdf at right radial boundary")
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_neutral_left",
+            boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:], vz, vr, vzeta, z,
+            parallel_io=parallel_io, n_neutral_species=composition.n_neutral_species,
+            description="Initial neutral-particle pdf at left radial boundary")
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_neutral_right",
+            boundary_distributions.pdf_rboundary_neutral[:,:,:,:,2,:], vz, vr, vzeta, z,
+            parallel_io=parallel_io, n_neutral_species=composition.n_neutral_species,
+            description="Initial neutral-particle pdf at right radial boundary")
+    end
+    return nothing
+end
+
 """
 Define coords group for coordinate information in the output file and write information
 about spatial coordinate grids
@@ -538,9 +569,9 @@ end
 """
 setup file i/o for distribution function variables
 """
-function setup_dfns_io(prefix, binary_format, r, z, vperp, vpa, vzeta, vr, vz, composition,
-                       collisions, evolve_density, evolve_upar, evolve_ppar, parallel_io,
-                       io_comm)
+function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vperp, vpa,
+                       vzeta, vr, vz, composition, collisions, evolve_density,
+                       evolve_upar, evolve_ppar, parallel_io, io_comm)
 
     @serial_region begin
         dfns_prefix = string(prefix, ".dfns")
@@ -553,10 +584,15 @@ function setup_dfns_io(prefix, binary_format, r, z, vperp, vpa, vzeta, vr, vz, c
         add_attribute!(fid, "file_info",
                        "Output distribution function data from the moment_kinetics code")
 
-        # write some overview information to the hdf5 file
+        # write some overview information to the output file
         write_overview!(fid, composition, collisions, parallel_io, evolve_density,
                         evolve_upar, evolve_ppar)
 
+        # write the distributions that may be used for boundary conditions to the output
+        # file
+        write_boundary_distributions!(fid, boundary_distributions, parallel_io,
+                                      composition, z, vperp, vpa, vzeta, vr, vz)
+
         ### define coordinate dimensions ###
         coords_group = define_spatial_coordinates!(fid, z, r, parallel_io)
         add_vspace_coordinates!(coords_group, vz, vr, vzeta, vpa, vperp, parallel_io)
@@ -651,52 +687,6 @@ function write_dfns_data_to_binary(ff, ff_neutral, moments, fields, t, n_ion_spe
     return nothing
 end
 
-"""
-Reload pdf and moments from an existing output file.
-"""
-function reload_evolving_fields!(pdf, moments, restart_filename, time_index,
-                                 composition, r, z, vpa)
-    code_time = 0.0
-    begin_serial_region()
-    @serial_region begin
-        fid = NCDataset(restart_filename,"r")
-        try
-            if time_index < 0
-                time_index = fid.dim["ntime"]
-            end
-            restart_n_species = fid.dim["n_species"]
-            restart_nr = fid.dim["nr"]
-            restart_nz = fid.dim["nz"]
-            restart_nvpa = fid.dim["nvpa"]
-            if restart_n_species != composition.n_species || restart_nr != r.n ||
-                restart_nz != z.n || restart_nvpa != vpa.n
-
-                error("Dimensions of restart file and input do not match.\n" *
-                      "Restart file was n_species=$restart_n_species, nr=$restart_nr, " *
-                      "nz=$restart_nz, nvpa=$restart_nvpa.\n" *
-                      "Input file gave  n_species=$(composition.n_species), nr=$(r.n), " *
-                      "nz=$(z.n), nvpa=$(vpa.n).")
-            end
-
-            code_time = fid["time"].var[time_index]
-            pdf.norm .= fid["f"].var[:,:,:,:,time_index]
-            moments.dens .= fid["density"].var[:,:,:,time_index]
-            moments.dens_updated .= true
-            moments.upar .= fid["parallel_flow"].var[:,:,:,time_index]
-            moments.upar_updated .= true
-            moments.ppar .= fid["parallel_pressure"].var[:,:,:,time_index]
-            moments.ppar_updated .= true
-            moments.qpar .= fid["parallel_heat_flux"].var[:,:,:,time_index]
-            moments.qpar_updated .= true
-            moments.vth .= fid["thermal_speed"].var[:,:,:,time_index]
-        finally
-            close(fid)
-        end
-    end
-
-    return code_time
-end
-
 @debug_shared_array begin
     # Special versions when using DebugMPISharedArray to avoid implicit conversion to
     # Array, which is forbidden.
diff --git a/src/file_io_hdf5.jl b/src/file_io_hdf5.jl
index aee3a2cdc..2d1d68a47 100644
--- a/src/file_io_hdf5.jl
+++ b/src/file_io_hdf5.jl
@@ -63,16 +63,28 @@ end
 # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group
 function write_single_value!(file_or_group::HDF5.H5DataStore, name,
                              data::Union{Number, AbstractString, AbstractArray{T,N}},
-                             coords::coordinate...; parallel_io, description=nothing) where {T,N}
+                             coords::Union{coordinate,mk_int}...; parallel_io,
+                             n_ion_species=nothing, n_neutral_species=nothing,
+                             description=nothing) where {T,N}
     if isa(data, Union{Number, AbstractString})
         file_or_group[name] = data
         return nothing
     end
 
+    if n_ion_species !== nothing && n_neutral_species != nothing
+        error("Cannot have both ion-species and neutral species dimensions." *
+              "Got n_ion_species=$n_ion_species, n_neutral_species=$n_neutral_species")
+    end
+
+    if n_ion_species !== nothing
+        coords = tuple(coords..., n_ion_species)
+    elseif n_neutral_species !== nothing
+        coords = tuple(coords..., n_neutral_species)
+    end
     dim_sizes, chunk_sizes = hdf5_get_fixed_dim_sizes(coords, parallel_io)
     io_var = create_dataset(file_or_group, name, T, dim_sizes, chunk=chunk_sizes)
-    local_ranges = Tuple(c.local_io_range for c ∈ coords)
-    global_ranges = Tuple(c.global_io_range for c ∈ coords)
+    local_ranges = Tuple(isa(c, coordinate) ? c.local_io_range : 1:c for c ∈ coords)
+    global_ranges = Tuple(isa(c, coordinate) ? c.global_io_range : 1:c for c ∈ coords)
 
     if N == 1
         io_var[global_ranges[1]] = @view data[local_ranges[1]]
diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 3648781d5..f0db2c2c7 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -60,13 +60,19 @@ end
 
 function write_single_value!(file_or_group::NCDataset, name,
                              value::Union{Number, AbstractString, AbstractArray{T,N}},
-                             coords::coordinate...; parallel_io, description=nothing) where {T,N}
+                             coords::coordinate...; parallel_io, n_ion_species=nothing,
+                             n_neutral_species=nothing, description=nothing) where {T,N}
     if description !== nothing
         attributes = Dict("description" => description)
     else
         attributes = ()
     end
 
+    if n_ion_species !== nothing && n_neutral_species != nothing
+        error("Cannot have both ion-species and neutral species dimensions." *
+              "Got n_ion_species=$n_ion_species, n_neutral_species=$n_neutral_species")
+    end
+
     if isa(value, Number) || isa(value, String)
         coords !== () && error("cannot pass coordinates with a scalar")
         type = typeof(value)
@@ -77,6 +83,14 @@ function write_single_value!(file_or_group::NCDataset, name,
             maybe_create_netcdf_dim(file_or_group, c)
         end
         dims = Tuple(c.name for c in coords)
+
+        if n_ion_species !== nothing
+            maybe_create_netcdf_dim(file_or_group, "ion_species", n_ion_species)
+            dims = tuple(dims..., "ion_species")
+        elseif n_neutral_species !== nothing
+            maybe_create_netcdf_dim(file_or_group, "neutral_species", n_neutral_species)
+            dims = tuple(dims..., "neutral_species")
+        end
     end
     if isa(value, Bool)
         # As a hack, write bools to NetCDF as Char, as NetCDF does not support bools (?),
diff --git a/src/load_data.jl b/src/load_data.jl
index 77fa24757..676bc32c7 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -17,6 +17,7 @@ export load_species_data
 using ..coordinates: define_coordinate
 using ..file_io: get_group
 using ..input_structs: advection_input, grid_input
+using ..looping
 
 using HDF5
 using MPI
@@ -109,6 +110,32 @@ function load_variable(file_or_group::NCDataset, name::String)
     end
 end
 
+"""
+Load a slice of a single variable from a file
+"""
+function load_slice() end
+function load_slice(file_or_group::HDF5.H5DataStore, name::String, slices_or_indices...)
+    # This overload deals with cases where fid is an HDF5 `File` or `Group` (`H5DataStore`
+    # is the abstract super-type for both
+    try
+        return file_or_group[name][slices_or_indices...]
+    catch
+        println("An error occured while loading $name")
+        rethrow()
+    end
+end
+function load_slice(file_or_group::NCDataset, name::String, slices_or_indices...)
+    # This overload deals with cases where fid is a NetCDF `Dataset` (which could be a
+    # file or a group).
+    try
+        var = file_or_group[name].var[slices_or_indices...]
+        return var
+    catch
+        println("An error occured while loading $name")
+        rethrow()
+    end
+end
+
 """
 Load data for a coordinate
 """
@@ -350,4 +377,202 @@ function load_neutral_pdf_data(fid; printout=false)
     return neutral_pdf
 end
 
+"""
+Reload pdf and moments from an existing output file.
+"""
+function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_prefix_iblock,
+                                 time_index, composition, r, z, vpa, vperp, vzeta, vr, vz)
+    code_time = 0.0
+    begin_serial_region()
+    @serial_region begin
+        fid = open_readonly_output_file(restart_prefix_iblock[1], "dfns";
+                                        iblock=restart_prefix_iblock[2])
+        try # finally to make sure to close file0
+            overview = get_group(fid, "overview")
+            dynamic = get_group(fid, "dynamic_data")
+            parallel_io = load_variable(overview, "parallel_io")
+            if time_index < 0
+                time_index, _ = load_time_data(fid)
+            end
+
+            if parallel_io
+                restart_n_ion_species, restart_n_neutral_species = load_species_data(fid)
+                restart_z, _ = load_coordinate_data(fid, "z")
+                restart_r, _ = load_coordinate_data(fid, "r")
+                restart_vperp, _ = load_coordinate_data(fid, "vperp")
+                restart_vpa, _ = load_coordinate_data(fid, "vpa")
+                restart_vzeta, _ = load_coordinate_data(fid, "vzeta")
+                restart_vr, _ = load_coordinate_data(fid, "vr")
+                restart_vz, _ = load_coordinate_data(fid, "vz")
+                if (restart_n_ion_species != composition.n_ion_species ||
+                    restart_n_neutral_species != composition.n_neutral_species ||
+                    restart_z.n != z.n_global || restart_r.n != r.n_global ||
+                    restart_vperp.n_global != vperp.n_global ||
+                    restart_vpa.n != vpa.n_global || restart_vzeta.n != vzeta.n_global ||
+                    restart_vr.n != vr.n_global || restart_vz.n != vz.n_global)
+
+                    error("Dimensions of restart file and input do not match.\n" *
+                          "Restart file was n_ion_species=$restart_n_ion_species, " *
+                          "n_neutral_species=$restart_n_neutral_species, nr=$(restart_r.n), " *
+                          "nz=$(restart_z.n), nvperp=$(restart_vperp.n), nvpa=$(restart_vpa.n).\n" *
+                          "nvzeta=$(restart_vzeta.n), nvr=$(restart_vr.n), nvz=$(restart_vz.n)." *
+                          "Input file gave n_ion_species=$(composition.n_ion_species), " *
+                          "n_neutral_species=$(composition.n_neutral_species), nr=$(r.n), " *
+                          "nz=$(z.n), nvperp=$(vperp.n), nvpa=$(vpa.n), nvzeta=$(vzeta.n), " *
+                          "nvr=$(vr.n), nvz=$(vz.n).")
+                end
+
+                code_time = load_slice(dynamic, "time", time_index)
+
+                function get_range(coord)
+                    if coord.irank == coord.nrank - 1
+                        return coord.global_io_range
+                    else
+                        # Need to modify the range to load the end-point that is duplicated on
+                        # the next process
+                        r = coord.global_io_range
+                        return r.start:(r.stop+1)
+                    end
+                end
+                r_range = get_range(r)
+                z_range = get_range(z)
+                vperp_range = get_range(vperp)
+                vpa_range = get_range(vpa)
+                vzeta_range = get_range(vzeta)
+                vr_range = get_range(vr)
+                vz_range = get_range(vz)
+
+                pdf.charged.norm .= load_slice(dynamic, "f", vpa_range, vperp_range,
+                                               z_range, r_range, :, time_index)
+                moments.charged.dens .= load_slice(dynamic, "density", z_range, r_range,
+                                                   :, time_index)
+                moments.charged.dens_updated .= true
+                moments.charged.upar .= load_slice(dynamic, "parallel_flow", z_range,
+                                                   r_range, :, time_index)
+                moments.charged.upar_updated .= true
+                moments.charged.ppar .= load_slice(dynamic, "parallel_pressure", z_range,
+                                                   r_range, :, time_index)
+                moments.charged.ppar_updated .= true
+                moments.charged.qpar .= load_slice(dynamic, "parallel_heat_flux", z_range,
+                                                   r_range, :, time_index)
+                moments.charged.qpar_updated .= true
+                moments.charged.vth .= load_slice(dynamic, "thermal_speed", z_range,
+                                                  r_range, :, time_index)
+
+                boundary_distributions_io = get_group(fid, "boundary_distributions")
+                boundary_distributions.pdf_rboundary_charged[:,:,:,1,:] .=
+                load_slice(boundary_distributions_io, "pdf_rboundary_charged_left",
+                           vpa_range, vperp_range, z_range, :)
+                boundary_distributions.pdf_rboundary_charged[:,:,:,2,:] .=
+                load_slice(boundary_distributions_io, "pdf_rboundary_charged_right",
+                           vpa_range, vperp_range, z_range, :)
+
+                if composition.n_neutral_species > 0
+                    pdf.neutral.norm .= load_slice(dynamic, "f_neutral", vz_range,
+                                                   vr_range, vzeta_range, z_range,
+                                                   r_range, :, time_index)
+                    moments.neutral.dens .= load_slice(dynamic, "density_neutral",
+                                                       z_range, r_range, :, time_index)
+                    moments.neutral.dens_updated .= true
+                    moments.neutral.uz .= load_slice(dynamic, "uz_neutral", z_range,
+                                                     r_range, :, time_index)
+                    moments.neutral.uz_updated .= true
+                    moments.neutral.pz .= load_slice(dynamic, "pz_neutral", z_range,
+                                                     r_range, :, time_index)
+                    moments.neutral.pz_updated .= true
+                    moments.neutral.qz .= load_slice(dynamic, "qz_neutral", z_range,
+                                                     r_range, :, time_index)
+                    moments.neutral.qz_updated .= true
+                    moments.neutral.vth .= load_slice(dynamic, "thermal_speed", z_range,
+                                                      r_range, :, time_index)
+
+                    boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:] .=
+                    load_slice(boundary_distributions_io, "pdf_rboundary_neutral_left",
+                               vz_range, vr_range, vzeta_range, z_range, :)
+                    boundary_distributions.pdf_rboundary_neutral[:,:,:,:,2,:] .=
+                    load_slice(boundary_distributions_io, "pdf_rboundary_neutral_right",
+                               vz_range, vr_range, vzeta_range, z_range, :)
+                end
+            else
+                restart_n_ion_species, restart_n_neutral_species = load_species_data(fid)
+                restart_z, _ = load_coordinate_data(fid, "z")
+                restart_r, _ = load_coordinate_data(fid, "r")
+                restart_vperp, _ = load_coordinate_data(fid, "vperp")
+                restart_vpa, _ = load_coordinate_data(fid, "vpa")
+                restart_vzeta, _ = load_coordinate_data(fid, "vzeta")
+                restart_vr, _ = load_coordinate_data(fid, "vr")
+                restart_vz, _ = load_coordinate_data(fid, "vz")
+                if (restart_n_ion_species != composition.n_ion_species ||
+                    restart_n_neutral_species != composition.n_neutral_species ||
+                    restart_z.n != z.n || restart_r.n != r.n || restart_vperp.n != vperp.n ||
+                    restart_vpa.n != vpa.n || restart_vzeta.n != vzeta.n ||
+                    restart_vr.n != vr.n || restart_vz.n != vz.n)
+
+                    error("Dimensions of restart file and input do not match.\n" *
+                          "Restart file was n_ion_species=$restart_n_ion_species, " *
+                          "n_neutral_species=$restart_n_neutral_species, nr=$(restart_r.n), " *
+                          "nz=$(restart_z.n), nvperp=$(restart_vperp.n), nvpa=$(restart_vpa.n).\n" *
+                          "nvzeta=$(restart_vzeta.n), nvr=$(restart_vr.n), nvz=$(restart_vz.n)." *
+                          "Input file gave n_ion_species=$(composition.n_ion_species), " *
+                          "n_neutral_species=$(composition.n_neutral_species), nr=$(r.n), " *
+                          "nz=$(z.n), nvperp=$(vperp.n), nvpa=$(vpa.n), nvzeta=$(vzeta.n), " *
+                          "nvr=$(vr.n), nvz=$(vz.n).")
+                end
+
+                code_time = load_slice(dynamic, "time", time_index)
+
+                pdf.charged.norm .= load_slice(dynamic, "f", :, :, :, :, :, time_index)
+                moments.charged.dens .= load_slice(dynamic, "density", :, :, :,
+                                                   time_index)
+                moments.charged.dens_updated .= true
+                moments.charged.upar .= load_slice(dynamic, "parallel_flow", :, :, :,
+                                                   time_index)
+                moments.charged.upar_updated .= true
+                moments.charged.ppar .= load_slice(dynamic, "parallel_pressure", :, :, :,
+                                                   time_index)
+                moments.charged.ppar_updated .= true
+                moments.charged.qpar .= load_slice(dynamic, "parallel_heat_flux", :, :, :,
+                                                   time_index)
+                moments.charged.qpar_updated .= true
+                moments.charged.vth .= load_slice(dynamic, "thermal_speed", :, :, :,
+                                                  time_index)
+
+                boundary_distributions_io = get_group(fid, "boundary_distributions")
+                boundary_distributions.pdf_rboundary_charged[:,:,:,1,:] .=
+                load_variable(boundary_distributions_io, "pdf_rboundary_charged_left")
+                boundary_distributions.pdf_rboundary_charged[:,:,:,2,:] .=
+                load_variable(boundary_distributions_io, "pdf_rboundary_charged_right")
+
+                if composition.n_neutral_species > 0
+                    pdf.neutral.norm .= load_slice(dynamic, "f_neutral", :, :, :, :, :, :,
+                                                   time_index)
+                    moments.neutral.dens .= load_slice(dynamic, "density_neutral", :, :,
+                                                       :, time_index)
+                    moments.neutral.dens_updated .= true
+                    moments.neutral.uz .= load_slice(dynamic, "uz_neutral", :, :, :,
+                                                     time_index)
+                    moments.neutral.uz_updated .= true
+                    moments.neutral.pz .= load_slice(dynamic, "pz_neutral", :, :, :,
+                                                     time_index)
+                    moments.neutral.pz_updated .= true
+                    moments.neutral.qz .= load_slice(dynamic, "qz_neutral", :, :, :,
+                                                     time_index)
+                    moments.neutral.qz_updated .= true
+                    moments.neutral.vth .= load_slice(dynamic, "thermal_speed", :, :, :,
+                                                      time_index)
+
+                    boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:] .=
+                    load_variable(boundary_distributions_io, "pdf_rboundary_neutral_left")
+                    boundary_distributions.pdf_rboundary_neutral[:,:,:,:,2,:] .=
+                    load_variable(boundary_distributions_io, "pdf_rboundary_neutral_right")
+                end
+            end
+        finally
+            close(fid)
+        end
+    end
+
+    return code_time
+end
+
 end
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index f61d0d821..211479f76 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -62,7 +62,7 @@ include("time_advance.jl")
 using TimerOutputs
 using Dates
 
-using .file_io: setup_file_io, finish_file_io, reload_evolving_fields!
+using .file_io: setup_file_io, finish_file_io
 using .file_io: write_data_to_ascii
 using .file_io: write_moments_data_to_binary, write_dfns_data_to_binary
 using .command_line_options: get_options
@@ -72,11 +72,13 @@ using .coordinates: define_coordinate
 using .debugging
 using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!,
                            enforce_boundary_conditions!
+using .load_data: reload_evolving_fields!
 using .looping
 using .moment_constraints: hard_force_moment_constraints!
 using .looping: debug_setup_loop_ranges_split_one_combination!
 using .moment_kinetics_input: mk_input, read_input_file, run_type, performance_test
 using .time_advance: setup_time_advance!, time_advance!
+using .type_definitions: mk_int
 
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
@@ -158,18 +160,67 @@ Append a number to the filename, to get a new, non-existing filename to backup t
 to.
 """
 function get_backup_filename(filename)
+    if !isfile(filename)
+        error("Requested to restart from $filename, but this file does not exist")
+    end
     counter = 1
-    basename, extension = splitext(filename)
-    backup_name = ""
-    while true
-        backup_name = "$(basename)_$(counter)$(extension)"
-        if !isfile(backup_name)
-            break
+    temp, extension = splitext(filename)
+    extension = extension[2:end]
+    temp, iblock_or_type = splitext(temp)
+    iblock_or_type = iblock_or_type[2:end]
+    iblock = nothing
+    basename = nothing
+    type = nothing
+    if iblock_or_type == "dfns"
+        iblock = nothing
+        type = iblock_or_type
+        basename = temp
+        parallel_io = true
+    else
+        # Filename had an iblock, so we are not using parallel I/O, but actually want to
+        # use the iblock for this block, not necessarily for the exact file that was
+        # passed.
+        iblock = iblock_index[]
+        basename, type = splitext(temp)
+        type = type[2:end]
+        parallel_io = false
+    end
+    if type != "dfns"
+        error("Must pass the '.dfns.h5' output file for restarting. Got $filename.")
+    end
+    backup_dfns_filename = ""
+    if parallel_io
+        # Using parallel I/O
+        while true
+            backup_dfns_filename = "$(basename)_$(counter).$(type).$(extension)"
+            if !isfile(backup_dfns_filename)
+                break
+            end
+            counter += 1
         end
-        counter += 1
+        # Create dfns_filename here even though it is the filename passed in, as
+        # parallel_io=false branch needs to get the right `iblock` for this block.
+        dfns_filename = "$(basename).dfns.$(extension)"
+        moments_filename = "$(basename).moments.$(extension)"
+        backup_moments_filename = "$(basename)_$(counter).moments.$(extension)"
+    else
+        while true
+            backup_dfns_filename = "$(basename)_$(counter).$(type).$(iblock).$(extension)"
+            if !isfile(backup_dfns_filename)
+                break
+            end
+            counter += 1
+        end
+        # Create dfns_filename here even though it is almost the filename passed in, in
+        # order to get the right `iblock` for this block.
+        dfns_filename = "$(basename).dfns.$(iblock).$(extension)"
+        moments_filename = "$(basename).moments.$(iblock).$(extension)"
+        backup_moments_filename = "$(basename)_$(counter).moments.$(iblock).$(extension)"
     end
-    backup_name == "" && error("Failed to find a name for backup file.")
-    return backup_name
+    backup_dfns_filename == "" && error("Failed to find a name for backup file.")
+    backup_prefix_iblock = ("$(basename)_$(counter)", iblock)
+    return dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
+           backup_moments_filename, backup_prefix_iblock
 end
 
 """
@@ -178,7 +229,7 @@ input must be the same as for the original run.
 """
 function restart_moment_kinetics(restart_filename::String, input_filename::String,
                                  time_index::Int=-1)
-    restart_moment_kinetics(restart_filename, input_from_TOML(input_filename),
+    restart_moment_kinetics(restart_filename, read_input_file(input_filename),
                             time_index)
     return nothing
 end
@@ -203,25 +254,25 @@ function restart_moment_kinetics(restart_filename::String, input_dict::Dict,
     try
         # Move the output file being restarted from to make sure it doesn't get
         # overwritten.
-        backup_filename = get_backup_filename(restart_filename)
-        global_rank[] == 0 && mv(restart_filename, backup_filename)
+        dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
+        backup_moments_filename, backup_prefix_iblock =
+            get_backup_filename(restart_filename)
+        if (parallel_io && global_rank[] == 0) || (!parallel_io && block_rank[] == 0)
+            mv(dfns_filename, backup_dfns_filename)
+            mv(moments_filename, backup_moments_filename)
+        end
 
         # Set up all the structs, etc. needed for a run.
-        pdf, scratch, code_time, t_input, vpa, z, r, vpa_spectral, z_spectral,
-        r_spectral, moments, fields, vpa_advect, z_advect, r_advect, composition,
-        collisions, num_diss_params, advance, scratch_dummy_sr, io, cdf =
-        setup_moment_kinetics(input_dict, backup_filename=backup_filename,
-                              restart_time_index=time_index)
+        mk_state = setup_moment_kinetics(input_dict,
+                                         restart_prefix_iblock=backup_prefix_iblock,
+                                         restart_time_index=time_index)
 
         try
-            time_advance!(pdf, scratch, code_time, t_input, vpa, z, r, vpa_spectral,
-                          z_spectral, r_spectral, moments, fields, vpa_advect, z_advect,
-                          r_advect, composition, collisions, num_diss_params, advance,
-                          scratch_dummy_sr, io, cdf)
+            time_advance!(mk_state...)
         finally
             # clean up i/o and communications
             # last 2 elements of mk_state are `io` and `cdf`
-            cleanup_moment_kinetics!(io, cdf)
+            cleanup_moment_kinetics!(mk_state[end-2:end]...)
         end
     catch e
         # Stop code from hanging when running on multiple processes if only one of them
@@ -247,7 +298,7 @@ reload data from time index given by `restart_time_index` for a restart.
 `debug_loop_type` and `debug_loop_parallel_dims` are used to force specific set ups for
 parallel loop ranges, and are only used by the tests in `debug_test/`.
 """
-function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
+function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
         restart_time_index=-1,
         debug_loop_type::Union{Nothing,NTuple{N,Symbol} where N}=nothing,
         debug_loop_parallel_dims::Union{Nothing,NTuple{N,Symbol} where N}=nothing)
@@ -305,7 +356,7 @@ function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
         allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
                                  evolve_moments, collisions, num_diss_params)
 
-    if backup_filename === nothing
+    if restart_prefix_iblock === nothing
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
@@ -319,8 +370,8 @@ function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
 
         # Reload pdf and moments from an existing output file
         code_time = reload_evolving_fields!(pdf, moments, boundary_distributions,
-                                            backup_filename, restart_time_index,
-                                            composition, r, z, vpa)
+                                            restart_prefix_iblock, restart_time_index,
+                                            composition, r, z, vpa, vperp, vzeta, vr, vz)
         _block_synchronize()
     end
     # create arrays and do other work needed to setup
@@ -333,9 +384,9 @@ function setup_moment_kinetics(input_dict::Dict; backup_filename=nothing,
             r_spectral, composition, drive_input, moments, t_input, collisions, species,
             geometry, boundary_distributions, num_diss_params, restarting)
     # setup i/o
-    ascii_io, io_moments, io_dfns = setup_file_io(io_input, vz, vr, vzeta, vpa, vperp, z,
-        r, composition, collisions, moments.evolve_density, moments.evolve_upar,
-        moments.evolve_ppar)
+    ascii_io, io_moments, io_dfns = setup_file_io(io_input, boundary_distributions, vz,
+        vr, vzeta, vpa, vperp, z, r, composition, collisions, moments.evolve_density,
+        moments.evolve_upar, moments.evolve_ppar)
     # write initial data to ascii files
     write_data_to_ascii(moments, fields, vpa, vperp, z, r, code_time, composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files

From f7530dd6f15e59bd94a30a798036aed4d1e0596a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 20 May 2023 16:35:39 +0100
Subject: [PATCH 019/428] Load evolve_* flags in post-processing

---
 src/load_data.jl       | 16 +++++++++++++---
 src/post_processing.jl |  5 ++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/load_data.jl b/src/load_data.jl
index 676bc32c7..6b62e629d 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -199,6 +199,18 @@ function load_species_data(fid; printout=false)
     return n_ion_species, n_neutral_species
 end
 
+"""
+"""
+function load_mk_options(fid)
+    overview = get_group(fid, "overview")
+
+    evolve_density = load_variable(overview, "evolve_density")
+    evolve_upar = load_variable(overview, "evolve_upar")
+    evolve_ppar = load_variable(overview, "evolve_ppar")
+
+    return evolve_density, evolve_upar, evolve_ppar
+end
+
 """
 """
 function load_time_data(fid; printout=false)
@@ -301,14 +313,12 @@ function load_charged_particle_moments_data(fid; printout=false)
 
     # Read charged species thermal speed
     thermal_speed = load_variable(group, "thermal_speed")
-    
-    evolve_ppar = false
 
     if printout
         println("done.")
     end
 
-    return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed, evolve_ppar
+    return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed
 end
 
 function load_neutral_particle_moments_data(fid; printout=false)
diff --git a/src/post_processing.jl b/src/post_processing.jl
index f08027df4..f83077ad3 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -39,7 +39,8 @@ using ..load_data: load_fields_data, load_pdf_data
 using ..load_data: load_charged_particle_moments_data, load_neutral_particle_moments_data
 using ..load_data: load_neutral_pdf_data
 using ..load_data: load_variable
-using ..load_data: load_coordinate_data, load_block_data, load_rank_data, load_species_data
+using ..load_data: load_coordinate_data, load_block_data, load_rank_data,
+                   load_species_data, load_mk_options
 using ..analysis: analyze_fields_data, analyze_moments_data, analyze_pdf_data,
                   check_Chodura_condition
 using ..velocity_moments: integrate_over_vspace
@@ -411,6 +412,8 @@ function analyze_and_plot_data(prefix...)
     # load species data
     n_ion_species, n_neutral_species =
         get_tuple_of_return_values(load_species_data, moments_files0)
+    evolve_density, evolve_upar, evolve_ppar =
+        get_tuple_of_return_values(load_mk_options, moments_files0)
 
     for f in moments_files0
         close(f)

From fc8a2a90933f949bf5a81074f47baa36a0e7717a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 May 2023 11:48:40 +0100
Subject: [PATCH 020/428] Fix error message in second_derivative!()

...needed brackets.
---
 src/calculus.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/calculus.jl b/src/calculus.jl
index eeb5f4087..b88fcb9f6 100644
--- a/src/calculus.jl
+++ b/src/calculus.jl
@@ -142,7 +142,7 @@ function second_derivative!(d2f, f, Q, coord, spectral)
         d2f[1] -= C * coord.scratch2_2d[end,end]
         d2f[end] += C * coord.scratch2_2d[1,1]
     else
-        error("Unsupported bc '$coord.bc'")
+        error("Unsupported bc '$(coord.bc)'")
     end
     return nothing
 end

From 1fdb162176f474eb9246a6354ff424194ab79566 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 May 2023 11:51:20 +0100
Subject: [PATCH 021/428] Add `vz_*` settings to wall-bc examples

---
 examples/wall-bc/wall+sheath-bc.toml      | 5 +++++
 examples/wall-bc/wall-bc_cheb.toml        | 5 +++++
 examples/wall-bc/wall-bc_cheb_split1.toml | 5 +++++
 examples/wall-bc/wall-bc_cheb_split2.toml | 5 +++++
 examples/wall-bc/wall-bc_cheb_split3.toml | 5 +++++
 5 files changed, 25 insertions(+)

diff --git a/examples/wall-bc/wall+sheath-bc.toml b/examples/wall-bc/wall+sheath-bc.toml
index 5be80021f..bff2fb441 100644
--- a/examples/wall-bc/wall+sheath-bc.toml
+++ b/examples/wall-bc/wall+sheath-bc.toml
@@ -61,3 +61,8 @@ vpa_nelement = 10
 vpa_L = 8.0
 vpa_bc = "periodic"
 vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
diff --git a/examples/wall-bc/wall-bc_cheb.toml b/examples/wall-bc/wall-bc_cheb.toml
index cbbfc0746..4f57a5084 100644
--- a/examples/wall-bc/wall-bc_cheb.toml
+++ b/examples/wall-bc/wall-bc_cheb.toml
@@ -60,6 +60,11 @@ vpa_nelement = 64
 vpa_L = 18.0
 vpa_bc = "both_zero"
 vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 9
+vz_nelement = 64
+vz_L = 18.0
+vz_bc = "both_zero"
+vz_discretization = "chebyshev_pseudospectral"
 
 [numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
diff --git a/examples/wall-bc/wall-bc_cheb_split1.toml b/examples/wall-bc/wall-bc_cheb_split1.toml
index 2f18c8fce..008b1ebf5 100644
--- a/examples/wall-bc/wall-bc_cheb_split1.toml
+++ b/examples/wall-bc/wall-bc_cheb_split1.toml
@@ -61,6 +61,11 @@ vpa_nelement = 64
 vpa_L = 18.0
 vpa_bc = "both_zero"
 vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 9
+vz_nelement = 64
+vz_L = 18.0
+vz_bc = "both_zero"
+vz_discretization = "chebyshev_pseudospectral"
 
 [numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
diff --git a/examples/wall-bc/wall-bc_cheb_split2.toml b/examples/wall-bc/wall-bc_cheb_split2.toml
index d06c4c27e..9ce9d7d74 100644
--- a/examples/wall-bc/wall-bc_cheb_split2.toml
+++ b/examples/wall-bc/wall-bc_cheb_split2.toml
@@ -61,6 +61,11 @@ vpa_nelement = 64
 vpa_L = 18.0
 vpa_bc = "both_zero"
 vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 9
+vz_nelement = 64
+vz_L = 18.0
+vz_bc = "both_zero"
+vz_discretization = "chebyshev_pseudospectral"
 
 [numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
diff --git a/examples/wall-bc/wall-bc_cheb_split3.toml b/examples/wall-bc/wall-bc_cheb_split3.toml
index e4ad530d0..b911f2898 100644
--- a/examples/wall-bc/wall-bc_cheb_split3.toml
+++ b/examples/wall-bc/wall-bc_cheb_split3.toml
@@ -61,6 +61,11 @@ vpa_nelement = 64
 vpa_L = 18.0
 vpa_bc = "both_zero"
 vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 9
+vz_nelement = 64
+vz_L = 18.0
+vz_bc = "both_zero"
+vz_discretization = "chebyshev_pseudospectral"
 
 [numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1

From 8ac8e944411835278820fd4c424ea6dc5eba2ef9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 3 May 2023 16:11:01 +0100
Subject: [PATCH 022/428] Fix initialisation of upar for distributed-MPI runs

---
 src/initial_conditions.jl | 28 ++++------------------------
 1 file changed, 4 insertions(+), 24 deletions(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index f33c92952..07c1d1cfd 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -339,17 +339,7 @@ function init_upar!(upar, z, r, spec, n_species)
                 # this is designed to give a nonzero J_{||i} at endpoints in z
                 # necessary for an electron sheath condition involving J_{||i}
                 # option "gaussian" to be consistent with usual init option for now
-                mid_ind = z.n ÷ 2
-                if z.n % 2 == 0
-                    z_midpoint = 0.5*(z.grid[mid_ind] + z.grid[mid_ind+1])
-                else
-                    # because ÷ does integer division (which floors the result), the
-                    # actual index of the mid-point is mid_ind+1
-                    z_midpoint = z.grid[mid_ind+1]
-                end
-                @. upar[:,ir,is] =
-                    (spec[is].z_IC.upar_amplitude * 2.0 *
-                           (z.grid[:] - z_midpoint)/z.L)
+                @. upar[:,ir,is] = spec[is].z_IC.upar_amplitude * 2.0 * z.grid / z.L
             else
                 @. upar[:,ir,is] = 0.0
             end
@@ -374,17 +364,7 @@ function init_uz!(uz, z, r, spec, n_species)
                 # this is designed to give a nonzero J_{||i} at endpoints in z
                 # necessary for an electron sheath condition involving J_{||i}
                 # option "gaussian" to be consistent with usual init option for now
-                mid_ind = z.n ÷ 2
-                if z.n % 2 == 0
-                    z_midpoint = 0.5*(z.grid[mid_ind] + z.grid[mid_ind+1])
-                else
-                    # because ÷ does integer division (which floors the result), the
-                    # actual index of the mid-point is mid_ind+1
-                    z_midpoint = z.grid[mid_ind+1]
-                end
-                @. uz[:,ir,is] =
-                    (spec[is].z_IC.upar_amplitude * 2.0 *
-                           (z.grid[:] - z_midpoint)/z.L)
+                @. uz[:,ir,is] = spec[is].z_IC.upar_amplitude * 2.0 * z.grid / z.L
             else
                 @. uz[:,ir,is] = 0.0
             end
@@ -502,10 +482,10 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
             zero = 1.e-14
             for ivpa ∈ 1:vpa.n
                 if vpa.grid[ivpa] > zero
-                    pdf[ivpa,:,1] .= 0.0
+                    lower_z_pdf_buffer[ivpa,:] .= 0.0
                 end
                 if vpa.grid[ivpa] < -zero
-                    pdf[ivpa,:,end] .= 0.0
+                    upper_z_pdf_buffer[ivpa,:] .= 0.0
                 end
             end
 

From b74ab358a7625816c013133a4e7495c92b62a63c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 May 2023 22:15:28 +0100
Subject: [PATCH 023/428] Fix initialisation with wall-bc for distributed-MPI

---
 src/initial_conditions.jl | 79 +++++++++++++++++++++++++++++++++------
 1 file changed, 67 insertions(+), 12 deletions(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 07c1d1cfd..171a41ba7 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -35,6 +35,8 @@ using ..velocity_moments: update_ppar!, update_upar!, update_density!
 
 using ..manufactured_solns: manufactured_solutions
 
+using MPI
+
 """
 """
 struct pdf_substruct{n_distribution}
@@ -454,7 +456,16 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
         else
             # First create distribution functions at the z-boundary points that obey the
             # boundary conditions.
-            for iz ∈ (1,z.n)
+            if z.irank == 0 && z.irank == z.nrank - 1
+                zrange = (1,z.n)
+            elseif z.irank == 0
+                zrange = (1,)
+            elseif z.irank == z.nrank - 1
+                zrange = (z.n,)
+            else
+                zrange = ()
+            end
+            for iz ∈ zrange
                 @loop_vperp ivperp begin
                     # Initialise as full-f distribution functions, then
                     # normalise/interpolate (if necessary). This makes it easier to
@@ -478,6 +489,18 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
                     @. pdf[:,ivperp,iz] *= 1.0 - exp(-vpa.grid^2*inverse_width)
                 end
             end
+            # Can use non-shared memory here because `init_charged_pdf_over_density!()` is
+            # called inside a `@serial_region`
+            lower_z_pdf_buffer = allocate_float(vpa.n, vperp.n)
+            upper_z_pdf_buffer = allocate_float(vpa.n, vperp.n)
+            if z.irank == 0
+                lower_z_pdf_buffer .= pdf[:,:,1]
+            end
+            if z.irank == z.nrank - 1
+                upper_z_pdf_buffer .= pdf[:,:,end]
+            end
+            @views MPI.Bcast!(lower_z_pdf_buffer, 0, z.comm)
+            @views MPI.Bcast!(upper_z_pdf_buffer, z.nrank - 1, z.comm)
 
             zero = 1.e-14
             for ivpa ∈ 1:vpa.n
@@ -495,15 +518,15 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
             @. z.scratch = 1.0 + 0.5 * (1.0 - (2.0 * z.grid / z.L)^2)
             for iz ∈ 1:z.n
                 # right_weight is 0 on left boundary and 1 on right boundary
-                right_weight = (z.grid[iz] - z.grid[1])/z.L
+                right_weight = z.grid[iz]/z.L + 0.5
                 #right_weight = min(max(0.5 + 0.5*(2.0*z.grid[iz]/z.L)^5, 0.0), 1.0)
                 #right_weight = min(max(0.5 +
                 #                       0.7*(2.0*z.grid[iz]/z.L) -
                 #                       0.2*(2.0*z.grid[iz]/z.L)^3, 0.0), 1.0)
                 # znorm is 1.0 at the boundary and 0.0 at the midplane
                 @views @. pdf[:,:,iz] = z.scratch[iz] * (
-                                            (1.0 - right_weight)*pdf[:,:,1] +
-                                            right_weight*pdf[:,:,end])
+                                            (1.0 - right_weight)*lower_z_pdf_buffer +
+                                            right_weight*upper_z_pdf_buffer)
             end
 
             # Get the unnormalised pdf and the moments of the constructed full-f
@@ -608,7 +631,16 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
         else
             # First create distribution functions at the z-boundary points that obey the
             # boundary conditions.
-            for iz ∈ (1,z.n)
+            if z.irank == 0 && z.irank == z.nrank - 1
+                zrange = (1,z.n)
+            elseif z.irank == 0
+                zrange = (1,)
+            elseif z.irank == z.nrank - 1
+                zrange = (z.n,)
+            else
+                zrange = ()
+            end
+            for iz ∈ zrange
                 @loop_vzeta_vr ivzeta ivr begin
                     # Initialise as full-f distribution functions, then
                     # normalise/interpolate (if necessary). This makes it easier to
@@ -633,6 +665,29 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
                     @. pdf[:,ivr,ivzeta,iz] *= 1.0 - exp(-vz.grid^2*inverse_width)
                 end
             end
+            # Can use non-shared memory here because `init_charged_pdf_over_density!()` is
+            # called inside a `@serial_region`
+            lower_z_pdf_buffer = allocate_float(vz.n, vr.n, vzeta.n)
+            upper_z_pdf_buffer = allocate_float(vz.n, vr.n, vzeta.n)
+            if z.irank == 0
+                lower_z_pdf_buffer .= pdf[:,:,:,1]
+            end
+            if z.irank == z.nrank - 1
+                upper_z_pdf_buffer .= pdf[:,:,:,end]
+            end
+
+            # Get the boundary pdfs from the processes that have the actual z-boundary
+            @views MPI.Bcast!(lower_z_pdf_buffer, 0, z.comm)
+            @views MPI.Bcast!(upper_z_pdf_buffer, z.nrank - 1, z.comm)
+
+            # Also need to get the (ion) wall fluxes from the processes that have the
+            # actual z-boundary
+            temp = Ref(wall_flux_0)
+            @views MPI.Bcast!(temp, 0, z.comm)
+            wall_flux_0 = temp[]
+            temp[] = wall_flux_L
+            @views MPI.Bcast!(temp, z.nrank - 1, z.comm)
+            wall_flux_L = temp[]
 
             knudsen_pdf = boundary_distributions.knudsen
 
@@ -641,28 +696,28 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
             # add this species' contribution to the combined ion/neutral particle flux
             # out of the domain at z=-Lz/2
             @views wall_flux_0 += integrate_over_negative_vz(
-                                      abs.(vz.grid) .* pdf[:,:,:,1], vz.grid, vz.wgts,
+                                      abs.(vz.grid) .* lower_z_pdf_buffer, vz.grid, vz.wgts,
                                       vz.scratch3, vr.grid, vr.wgts, vzeta.grid,
                                       vzeta.wgts)
             # for left boundary in zed (z = -Lz/2), want
             # f_n(z=-Lz/2, v_z > 0) = Γ_0 * f_KW(v_z) * pdf_norm_fac(-Lz/2)
             @loop_vz ivz begin
                 if vz.grid[ivz] > zero
-                    @. pdf[ivz,:,:,1] = wall_flux_0 * knudsen_pdf[ivz,:,:]
+                    @. lower_z_pdf_buffer[ivz,:,:] = wall_flux_0 * knudsen_pdf[ivz,:,:]
                 end
             end
 
             # add this species' contribution to the combined ion/neutral particle flux
             # out of the domain at z=-Lz/2
             @views wall_flux_L += integrate_over_positive_vz(
-                                      abs.(vz.grid) .* pdf[:,:,:,end], vz.grid, vz.wgts,
+                                      abs.(vz.grid) .* upper_z_pdf_buffer, vz.grid, vz.wgts,
                                       vz.scratch3, vr.grid, vr.wgts, vzeta.grid,
                                       vzeta.wgts)
             # for right boundary in zed (z = Lz/2), want
             # f_n(z=Lz/2, v_z < 0) = Γ_Lz * f_KW(v_z) * pdf_norm_fac(Lz/2)
             @loop_vz ivz begin
                 if vz.grid[ivz] < -zero
-                    @. pdf[ivz,:,:,end] = wall_flux_L * knudsen_pdf[ivz,:,:]
+                    @. upper_z_pdf_buffer[ivz,:,:] = wall_flux_L * knudsen_pdf[ivz,:,:]
                 end
             end
 
@@ -673,15 +728,15 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
 
             for iz ∈ 1:z.n
                 # right_weight is 0 on left boundary and 1 on right boundary
-                right_weight = (z.grid[iz] - z.grid[1])/z.L
+                right_weight = z.grid[iz]/z.L + 0.5
                 #right_weight = min(max(0.5 + 0.5*(2.0*z.grid[iz]/z.L)^5, 0.0), 1.0)
                 #right_weight = min(max(0.5 +
                 #                       0.7*(2.0*z.grid[iz]/z.L) -
                 #                       0.2*(2.0*z.grid[iz]/z.L)^3, 0.0), 1.0)
                 # znorm is 1.0 at the boundary and 0.0 at the midplane
                 @views @. pdf[:,:,:,iz] = z.scratch[iz] * (
-                                              (1.0 - right_weight)*pdf[:,:,:,1] +
-                                              right_weight*pdf[:,:,:,end])
+                                              (1.0 - right_weight)*lower_z_pdf_buffer +
+                                              right_weight*upper_z_pdf_buffer)
             end
 
             # Get the unnormalised pdf and the moments of the constructed full-f

From ee4850526f68eeba404824803f0db9a411c88420 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 May 2023 16:29:08 +0100
Subject: [PATCH 024/428] (Re-)implement plotting of neutral quantities for
 1D1V cases

---
 src/post_processing.jl | 1357 +++++++++++++++++++++-------------------
 1 file changed, 713 insertions(+), 644 deletions(-)

diff --git a/src/post_processing.jl b/src/post_processing.jl
index f83077ad3..8939b07ca 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -174,7 +174,7 @@ function read_distributed_zr_data!(var::Array{mk_float,7}, var_name::String,
             for iz_local in imin_z:nz_local
                 ir_global = iglobal_func(ir_local,r_irank,nr_local)
                 iz_global = iglobal_func(iz_local,z_irank,nz_local)
-                var[:,:,:,iz_global,ir_global,:,:] .= var_local[iz_local,ir_local,:,:]
+                var[:,:,:,iz_global,ir_global,:,:] .= var_local[:,:,:,iz_local,ir_local,:,:]
             end
         end
         close(fid)
@@ -632,6 +632,22 @@ function analyze_and_plot_data(prefix...)
         get_tuple_of_return_values(read_distributed_zr_data!, ff, "f", run_names, "dfns",
                                    nblocks, Tuple(this_z.n for this_z ∈ z),
                                    Tuple(this_r.n for this_r ∈ r))
+        if maximum(n_neutral_species) > 0
+            neutral_ff = get_tuple_of_return_values(
+                             allocate_global_zr_neutral_dfns,
+                             Tuple(this_vz.n_global for this_vz ∈ vz),
+                             Tuple(this_vr.n_global for this_vr ∈ vr),
+                             Tuple(this_vzeta.n_global for this_vzeta ∈ vzeta),
+                             Tuple(this_z.n_global for this_z ∈ z),
+                             Tuple(this_r.n_global for this_r ∈ r),
+                             n_neutral_species, ntime_pdfs)
+            get_tuple_of_return_values(read_distributed_zr_data!, neutral_ff, "f_neutral",
+                                       run_names, "dfns", nblocks,
+                                       Tuple(this_z.n for this_z ∈ z),
+                                       Tuple(this_r.n for this_r ∈ r))
+        else
+            neutral_ff = nothing
+        end
 
         #evaluate 1D-1V diagnostics at fixed ir0
         plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
@@ -650,8 +666,19 @@ function analyze_and_plot_data(prefix...)
             Tuple(qpar[:,ir0,:,:] for qpar ∈ parallel_heat_flux_at_pdf_times),
             Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed_at_pdf_times),
             Tuple(f[:,ivperp0,:,ir0,:,:] for f ∈ ff),
+            Tuple(neutral_n[:,ir0,:,:] for neutral_n ∈ neutral_density),
+            Tuple(uz[:,ir0,:,:] for uz ∈ neutral_uz),
+            Tuple(pz[:,ir0,:,:] for pz ∈ neutral_pz),
+            Tuple(qz[:,ir0,:,:] for qz ∈ neutral_qz),
+            Tuple(neutral_vth[:,ir0,:,:] for neutral_vth ∈ neutral_thermal_speed),
+            Tuple(neutral_n[:,ir0,:,:] for neutral_n ∈ neutral_density_at_pdf_times),
+            Tuple(uz[:,ir0,:,:] for uz ∈ neutral_uz_at_pdf_times),
+            Tuple(pz[:,ir0,:,:] for pz ∈ neutral_pz_at_pdf_times),
+            Tuple(qz[:,ir0,:,:] for qz ∈ neutral_qz_at_pdf_times),
+            Tuple(neutral_vth[:,ir0,:,:] for neutral_vth ∈ neutral_thermal_speed_at_pdf_times),
+            Tuple(neutral_f[:,ivr0,ivzeta0,:,ir0,:,:] for neutral_f ∈ neutral_ff),
             n_ion_species, n_neutral_species, evolve_density, evolve_upar, evolve_ppar,
-            vpa, z, ntime, time, ntime_pdfs, time_pdfs)
+            vz, vpa, z, ntime, time, ntime_pdfs, time_pdfs)
     end
 
     diagnostics_2d = false
@@ -920,8 +947,12 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed,
         phi_at_pdf_times, density_at_pdf_times, parallel_flow_at_pdf_times,
         parallel_pressure_at_pdf_times, parallel_heat_flux_at_pdf_times,
-        thermal_speed_at_pdf_times, ff, n_ion_species, n_neutral_species, evolve_density,
-        evolve_upar, evolve_ppar, vpa, z, ntime, time, ntime_pdfs, time_pdfs)
+        thermal_speed_at_pdf_times, ff, neutral_density, neutral_uz, neutral_pz,
+        neutral_qz, neutral_thermal_speed, neutral_density_at_pdf_times,
+        neutral_uz_at_pdf_times, neutral_pz_at_pdf_times, neutral_qz_at_pdf_times,
+        neutral_thermal_speed_at_pdf_times, neutral_ff, n_ion_species, n_neutral_species,
+        evolve_density, evolve_upar, evolve_ppar, vz, vpa, z, ntime, time, ntime_pdfs,
+        time_pdfs)
 
     n_runs = length(run_names)
 
@@ -941,7 +972,6 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
     # create the requested plots of the fields
     plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
                 z, iz0, run_names, 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 =
@@ -953,17 +983,108 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         parallel_pressure, delta_ppar, ppar_fldline_avg,
         thermal_speed, delta_vth, vth_fldline_avg,
         parallel_heat_flux, delta_qpar, qpar_fldline_avg,
-        pp, run_names, time, itime_min, itime_max, nwrite_movie, z, iz0, n_ion_species)
-    # load particle distribution function (pdf) data
+        pp, run_names, time, itime_min, itime_max, nwrite_movie, z, iz0, n_ion_species,
+        "ion")
+    if maximum(n_neutral_species) > 0
+        # analyze the velocity neutral moments data
+        neutral_density_fldline_avg, neutral_uz_fldline_avg, neutral_pz_fldline_avg,
+        neutral_vth_fldline_avg, neutral_qz_fldline_avg, delta_neutral_density,
+        delta_neutral_uz, delta_neutral_pz, delta_neutral_vth, delta_neutral_qz =
+            get_tuple_of_return_values(analyze_moments_data, neutral_density, neutral_uz,
+                neutral_pz, neutral_thermal_speed, neutral_qz, ntime, n_neutral_species, z)
+        # create the requested plots of the neutral moments
+        plot_moments(neutral_density, delta_neutral_density, neutral_density_fldline_avg,
+            neutral_uz, delta_neutral_uz, neutral_uz_fldline_avg, neutral_pz,
+            delta_neutral_pz, neutral_pz_fldline_avg, neutral_thermal_speed,
+            delta_neutral_vth, neutral_vth_fldline_avg, neutral_qz, delta_neutral_qz,
+            neutral_qz_fldline_avg, pp, run_names, time, itime_min, itime_max,
+            nwrite_movie, z, iz0, n_ion_species, "neutral")
+    end
+
     # analyze the pdf data
     f_fldline_avg, delta_f, dens_moment, upar_moment, ppar_moment =
         get_tuple_of_return_values(analyze_pdf_data, ff, n_ion_species, ntime_pdfs, z,
                                    vpa, thermal_speed_at_pdf_times, 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
-    n_ion_species_max = maximum(n_ion_species)
+    plot_dfns(density_at_pdf_times, parallel_flow_at_pdf_times,
+        parallel_pressure_at_pdf_times, thermal_speed_at_pdf_times, ff, dens_moment,
+        upar_moment, ppar_moment, time_pdfs, n_ion_species, z, vpa, evolve_density,
+        evolve_upar, evolve_ppar, run_names, itime_min_pdfs, itime_max_pdfs,
+        nwrite_movie_pdfs, iz0, "ion")
+
+    if maximum(n_neutral_species) > 0
+        # analyze the neutral pdf data
+        neutral_f_fldline_avg, delta_neutral_f, neutral_dens_moment, neutral_uz_moment,
+        neutral_pz_moment =
+            get_tuple_of_return_values(analyze_pdf_data, neutral_ff, n_neutral_species,
+                                       ntime_pdfs, z, vpa, neutral_thermal_speed_at_pdf_times,
+                                       evolve_ppar)
+
+        plot_dfns(neutral_density_at_pdf_times, neutral_uz_at_pdf_times,
+            neutral_pz_at_pdf_times, neutral_thermal_speed_at_pdf_times, neutral_ff,
+            neutral_dens_moment, neutral_uz_moment, neutral_pz_moment, time_pdfs,
+            n_neutral_species, z, vz, evolve_density, evolve_upar, evolve_ppar, run_names,
+            itime_min_pdfs, itime_max_pdfs, nwrite_movie_pdfs, iz0, "neutral")
+    end
+
+    println("done.")
+end
+
+"""
+"""
+function calculate_and_write_frequencies(run_name, ntime, time, z, itime_min, itime_max,
+                                         iz0, delta_phi, pp)
+    if pp.calculate_frequencies
+        println("Calculating the frequency and damping/growth rate...")
+        # shifted_time = t - t0
+        shifted_time = allocate_float(ntime)
+        @. shifted_time = time - time[itime_min]
+        # assume phi(z0,t) = A*exp(growth_rate*t)*cos(ω*t + φ)
+        # and fit phi(z0,t)/phi(z0,t0), which eliminates the constant A pre-factor
+        @views phi_fit = fit_delta_phi_mode(shifted_time[itime_min:itime_max], z,
+                                            delta_phi[:, itime_min:itime_max])
+        frequency = phi_fit.frequency
+        growth_rate = phi_fit.growth_rate
+
+        # write info related to fit to file
+        io = open_ascii_output_file(run_name, "frequency_fit.txt")
+        println(io, "#growth_rate: ", phi_fit.growth_rate,
+                "  frequency: ", phi_fit.frequency,
+                " fit_errors: ", phi_fit.amplitude_fit_error, " ",
+                phi_fit.offset_fit_error, " ", phi_fit.cosine_fit_error)
+        println(io)
+
+        # Calculate the fitted phi as a function of time at index iz0
+        L = z[end] - z[begin]
+        fitted_delta_phi =
+            @. (phi_fit.amplitude0 * cos(2.0 * π * (z[iz0] + phi_fit.offset0) / L)
+                * exp(phi_fit.growth_rate * shifted_time)
+                * cos(phi_fit.frequency * shifted_time + phi_fit.phase))
+        for i ∈ 1:ntime
+            println(io, "time: ", time[i], "  delta_phi: ", delta_phi[iz0,i],
+                    "  fitted_delta_phi: ", fitted_delta_phi[i])
+        end
+        close(io)
+    else
+        frequency = 0.0
+        growth_rate = 0.0
+        phase = 0.0
+        shifted_time = allocate_float(ntime)
+        @. shifted_time = time - time[itime_min]
+        fitted_delta_phi = zeros(ntime)
+
+    end
+    return frequency, growth_rate, shifted_time, fitted_delta_phi
+end
+
+"""
+"""
+function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
+    z, iz0, run_names, fitted_delta_phi, pp)
+
+    println("Plotting fields data...")
+
+    n_runs = length(run_names)
     if n_runs == 1
         prefix = run_names[1]
         legend = false
@@ -971,547 +1092,122 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         prefix = default_compare_prefix
         legend = true
     end
-    for is ∈ 1:n_ion_species_max
-        if n_ion_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
+
+    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)
         plot(legend=legend)
-        for (t, n, n_int, run_label) ∈ zip(time_pdfs, density_at_pdf_times, dens_moment, run_names)
-            @views plot!(t, n[iz0,is,:] .- n_int[iz0,is,:], label=run_label)
+        for (t, p, run_label) ∈ zip(time, phi, run_names)
+            @views plot!(t, p[iz0,:], label=run_label)
         end
-        outfile = string(prefix, "_intf0_vs_t", spec_string, ".pdf")
+        outfile = string(prefix, "_phi0_vs_t.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
         plot(legend=legend)
-        for (t, upar_int, run_label) ∈ zip(time_pdfs, upar_moment, run_names)
-            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)
+        for (t, dp, fit, run_label) ∈ zip(time, delta_phi, fitted_delta_phi, run_names)
+            # 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(prefix, "_intwf0_vs_t", spec_string, ".pdf")
+        outfile = string(prefix, "_delta_phi0_vs_t.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
-        plot(legend=legend)
-        for (t, ppar, ppar_int, run_label) ∈ zip(time_pdfs, parallel_pressure_at_pdf_times, ppar_moment, run_names)
-            @views plot(t, ppar[iz0,is,:] .- ppar_int[iz0,is,:], label=run_label)
-        end
-        outfile = string(prefix, "_intw2f0_vs_t", spec_string, ".pdf")
+    end
+    if pp.plot_phi_vs_z_t
+        # make a heatmap plot of ϕ(z,t)
+        subplots = (heatmap(t, this_z.grid, p, xlabel="time", ylabel="z", title=run_label,
+                            c = :deep)
+                    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names))
+        plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+        outfile = string(prefix, "_phi_vs_z_t.pdf")
         savefig(outfile)
-        #fmin = minimum(ff[:,:,is,:])
-        #fmax = maximum(ff[:,:,is,:])
-        if pp.animate_f_vs_vpa_z
-            # make a gif animation of ln f(vpa,z,t)
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
-                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)
-            # make a gif animation of f(vpa,z,t)
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
-                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
-                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+    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
+            plot(legend=legend)
+            for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
+                @views plot!(this_z.grid, p[:,i], xlabel="z", ylabel="ϕ",
+                             ylims=(phimin, phimax), label=run_label)
             end
-            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")
-            subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
-            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
-            outfile = string(prefix, "_f_vs_z_vpa_final", spec_string, ".pdf")
-            savefig(outfile)
+        end
+        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(prefix, "_harrison_comparison.pdf")
+    # savefig(outfile)
+    plot(legend=legend)
+    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
+        plot!(this_z.grid, p[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label=run_label,
+              linewidth=2)
+    end
+    outfile = string(prefix, "_phi_final.pdf")
+    savefig(outfile)
 
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa.grid, 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)
+    println("done.")
+end
 
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa.grid, 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.plot_f_unnormalized_vs_vpa_z
-            PyPlot.clf()
-            fig = PyPlot.figure(1, figsize=(6*n_runs,4))
-            it = itime_max_pdfs
-            # 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_at_pdf_times,
-                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
-                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
-                                  run_names)
+"""
+"""
+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_names, time, itime_min, itime_max, nwrite_movie, z, iz0, n_species, label)
 
-                PyPlot.subplot(1, n_runs, run_ind)
-                @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                    f[:,:,is,it], this_z.grid, this_vpa.grid, n[:,is,it],
-                    upar[:,is,it], vth[:,is,it], ev_n, ev_u, ev_p)
-                plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; title=run_label,
-                                      plot_log=false)
-            end
-            PyPlot.savefig(string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".pdf"))
-            PyPlot.clf()
+    println("Plotting velocity moments data...")
 
-            plot(legend=legend)
-            it = itime_max_pdfs
-            # 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_at_pdf_times,
-                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
-                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
-                                  run_names)
+    n_runs = length(run_names)
+    if n_runs == 1
+        prefix = run_names[1]
+        legend = false
+    else
+        prefix = default_compare_prefix
+        legend = true
+    end
 
-                @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                    f[:,1,is,it], this_vpa.grid, n[1,is,it], upar[1,is,it], vth[1,is,it],
-                    ev_n, ev_u, ev_p)
-                @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=0)",
-                             label=run_label)
+    # plot the species-summed, field-line averaged vs time
+    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_names)
+        @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, "_$(label)_denstot_vs_t.pdf")
+    savefig(outfile)
+    for is ∈ 1:maximum(n_species)
+        spec_string = string(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)
+            plot(legend=legend)
+            for (t, n, run_label) ∈ zip(time, density, run_names)
+                @views plot!(t, n[iz0,is,:], label=run_label)
             end
-            savefig(string(prefix, "_f0_unnorm_vs_vpa", spec_string, ".pdf"))
-
-            plot(legend=legend)
-            it = itime_max_pdfs
-            # 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_at_pdf_times,
-                                  parallel_flow_at_pdf_times, thermal_speed_at_pdf_times,
-                                  evolve_density, evolve_upar, evolve_ppar, z, vpa,
-                                  run_names)
-
-                @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                    f[:,end,is,it], this_vpa.grid, n[end,is,it], upar[end,is,it], vth[end,is,it],
-                    ev_n, ev_u, ev_p)
-                @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=L)",
-                             label=run_label)
-            end
-            savefig(string(prefix, "_fL_unnorm_vs_vpa", spec_string, ".pdf"))
-        end
-        if pp.animate_f_unnormalized
-            ## 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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid,
-            #                           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_names))
-            #    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)
-            ## make a gif animation of log(f_unnorm)(v_parallel_unnorm,z,t)
-            #anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid, 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_names))
-            #    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)
-
-            matplotlib = pyimport("matplotlib")
-            matplotlib.use("agg")
-            matplotlib_animation = pyimport("matplotlib.animation")
-            iframes = collect(itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs)
-            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_at_pdf_times,
-                                      parallel_flow_at_pdf_times,
-                                      thermal_speed_at_pdf_times, evolve_density,
-                                      evolve_upar, evolve_ppar, z, vpa, run_names)
-
-                    PyPlot.subplot(1, n_runs, run_ind)
-                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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*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))
-            PyPlot.clf()
-
-            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_at_pdf_times,
-                                      parallel_flow_at_pdf_times,
-                                      thermal_speed_at_pdf_times, evolve_density,
-                                      evolve_upar, evolve_ppar, z, vpa, run_names)
-
-                    PyPlot.subplot(1, n_runs, run_ind)
-                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
-                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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*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))
-
-            # Ensure PyPlot figure is cleared
-            closeall()
-
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (f, n, upar, vth, ev_n, ev_u, ev_p, this_vpa, run_label) ∈
-                    zip(ff, density_at_pdf_times, parallel_flow_at_pdf_times,
-                        thermal_speed_at_pdf_times, evolve_density, evolve_upar,
-                        evolve_ppar, vpa, run_names)
-                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                        f[:,1,is,i], this_vpa.grid, 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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (f, n, upar, vth, ev_n, ev_u, ev_p, this_vpa, run_label) ∈
-                    zip(ff, density_at_pdf_times, parallel_flow_at_pdf_times,
-                        thermal_speed_at_pdf_times, evolve_density, evolve_upar,
-                        evolve_ppar, vpa, run_names)
-                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
-                        f[:,end,is,i], this_vpa.grid, 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)
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                subplots = (@views heatmap(this_z, this_vpa.grid, 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_names))
-                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)
-        end
-        if pp.animate_f_vs_vpa_z0
-            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (f, this_z, run_label) ∈ zip(ff, z, run_names)
-                    @views plot!(this_z, f[ivpa0,:,is,i], ylims = (fmin,fmax),
-                                 label=run_label)
-                end
-            end
-            outfile = string(prefix, "_f_vs_z", spec_string, ".gif")
-            gif(anim, outfile, fps=5)
-        end
-        if pp.animate_deltaf_vs_vpa_z0
-            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (df, this_z, run_label) ∈ zip(delta_f, z, run_names)
-                    @views plot!(this_z, df[ivpa0,:,is,i], ylims = (fmin,fmax),
-                                 label=run_label)
-                end
-            end
-            outfile = string(prefix, "_deltaf_vs_z", spec_string, ".gif")
-            gif(anim, outfile, fps=5)
-        end
-        if pp.animate_f_vs_vpa_z0
-            fmin = minimum(minimum(f[:,iz0,is,:]) for f ∈ ff)
-            fmax = maximum(maximum(f[:,iz0,is,:]) for f ∈ ff)
-
-            # if is == 1
-            #     tmp = copy(ff)
-            #     @. tmp[:,1,1,:] /= vpa^2
-            #     bohm_integral = copy(time)
-            #     for i ∈ 1:ntime
-            #         @views bohm_integral[i] = integrate_over_vspace(tmp[1:cld(nvpa,2)-1,1,1,i],vpa_wgts[1:cld(nvpa,2)-1])/2.0
-            #     end
-            #     plot(time, bohm_integral, xlabel="time", label="Bohm integral")
-            #     plot!(time, density[1,1,:], label="nᵢ(zmin)")
-            #     outfile = string(prefix, "_Bohm_criterion.pdf")
-            #     savefig(outfile)
-            #     println()
-            #     if bohm_integral[end] <= density[1,1,end]
-            #         println("Bohm criterion: ", bohm_integral[end], " <= ", density[1,1,end], " is satisfied!")
-            #     else
-            #         println("Bohm criterion: ", bohm_integral[end], " <= ", density[1,1,end], " is not satisfied!")
-            #     end
-            #     println()
-            #     for j ∈ 0:10
-            #         println("j: ", j, "  Bohm integral: ", integrate_over_vspace(tmp[1:cld(nvpa,2)-j,1,1,end],vpa_wgts[1:cld(nvpa,2)-j,end])/2.0)
-            #     end
-            # end
-            # make a gif animation of f(vpa,z0,t)
-            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                #@views plot(vpa, ff[iz0,:,is,i], ylims = (fmin,fmax))
-                plot(legend=legend)
-                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
-                    @views plot!(this_vpa.grid, f[:,iz0,is,i], label=run_label)
-                end
-            end
-            outfile = string(prefix, "_f_vs_vpa", spec_string, ".gif")
-            gif(anim, outfile, fps=5)
-        end
-        if pp.animate_deltaf_vs_vpa_z0
-            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
-                plot(legend=legend)
-                for (df, this_vpa, fn, fx, run_label) ∈
-                        zip(delta_f, vpa, fmin, fmax, run_names)
-                    @views plot!(this_vpa.grid, delta_f[:,iz0,is,i], ylims = (fn,fx),
-                                label=run_label)
-                end
-            end
-            outfile = string(prefix, "_deltaf_vs_vpa", spec_string, ".gif")
-            gif(anim, outfile, fps=5)
-        end
-    end
-    println("done.")
-
-end
-
-"""
-"""
-function calculate_and_write_frequencies(run_name, ntime, time, z, itime_min, itime_max,
-                                         iz0, delta_phi, pp)
-    if pp.calculate_frequencies
-        println("Calculating the frequency and damping/growth rate...")
-        # shifted_time = t - t0
-        shifted_time = allocate_float(ntime)
-        @. shifted_time = time - time[itime_min]
-        # assume phi(z0,t) = A*exp(growth_rate*t)*cos(ω*t + φ)
-        # and fit phi(z0,t)/phi(z0,t0), which eliminates the constant A pre-factor
-        @views phi_fit = fit_delta_phi_mode(shifted_time[itime_min:itime_max], z,
-                                            delta_phi[:, itime_min:itime_max])
-        frequency = phi_fit.frequency
-        growth_rate = phi_fit.growth_rate
-
-        # write info related to fit to file
-        io = open_ascii_output_file(run_name, "frequency_fit.txt")
-        println(io, "#growth_rate: ", phi_fit.growth_rate,
-                "  frequency: ", phi_fit.frequency,
-                " fit_errors: ", phi_fit.amplitude_fit_error, " ",
-                phi_fit.offset_fit_error, " ", phi_fit.cosine_fit_error)
-        println(io)
-
-        # Calculate the fitted phi as a function of time at index iz0
-        L = z[end] - z[begin]
-        fitted_delta_phi =
-            @. (phi_fit.amplitude0 * cos(2.0 * π * (z[iz0] + phi_fit.offset0) / L)
-                * exp(phi_fit.growth_rate * shifted_time)
-                * cos(phi_fit.frequency * shifted_time + phi_fit.phase))
-        for i ∈ 1:ntime
-            println(io, "time: ", time[i], "  delta_phi: ", delta_phi[iz0,i],
-                    "  fitted_delta_phi: ", fitted_delta_phi[i])
-        end
-        close(io)
-    else
-        frequency = 0.0
-        growth_rate = 0.0
-        phase = 0.0
-        shifted_time = allocate_float(ntime)
-        @. shifted_time = time - time[itime_min]
-        fitted_delta_phi = zeros(ntime)
-
-    end
-    return frequency, growth_rate, shifted_time, fitted_delta_phi
-end
-
-"""
-"""
-function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
-    z, iz0, run_names, fitted_delta_phi, pp)
-
-    println("Plotting fields data...")
-
-    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)
-        plot(legend=legend)
-        for (t, p, run_label) ∈ zip(time, phi, run_names)
-            @views plot!(t, p[iz0,:], label=run_label)
-        end
-        outfile = string(prefix, "_phi0_vs_t.pdf")
-        savefig(outfile)
-        plot(legend=legend)
-        for (t, dp, fit, run_label) ∈ zip(time, delta_phi, fitted_delta_phi, run_names)
-            # 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(prefix, "_delta_phi0_vs_t.pdf")
-        savefig(outfile)
-    end
-    if pp.plot_phi_vs_z_t
-        # make a heatmap plot of ϕ(z,t)
-        subplots = (heatmap(t, this_z.grid, p, xlabel="time", ylabel="z", title=run_label,
-                            c = :deep)
-                    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names))
-        plot(subplots..., layout=(1,n_runs), size=(600*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
-            plot(legend=legend)
-            for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
-                @views plot!(this_z.grid, p[:,i], xlabel="z", ylabel="ϕ",
-                             ylims=(phimin, phimax), label=run_label)
-            end
-        end
-        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(prefix, "_harrison_comparison.pdf")
-    # savefig(outfile)
-    plot(legend=legend)
-    for (t, this_z, p, run_label) ∈ zip(time, z, phi, run_names)
-        plot!(this_z.grid, p[:,end], xlabel="z/Lz", ylabel="eϕ/Te", label=run_label,
-              linewidth=2)
-    end
-    outfile = string(prefix, "_phi_final.pdf")
-    savefig(outfile)
-
-    println("done.")
-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_names, 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 = 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_names)
-        @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:maximum(n_species)
-        spec_string = string(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)
-            plot(legend=legend)
-            for (t, n, run_label) ∈ zip(time, density, run_names)
-                @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
+            outfile = string(prefix, "_$(label)_dens0_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+            # plot the time trace of n_s(z=z0)-density_fldline_avg
             plot(legend=legend)
             for (t, dn, run_label) ∈ zip(time, delta_density, run_names)
                 @views plot!(t, abs.(dn[iz0,is,:]), yaxis=:log, label=run_label)
             end
-            outfile = string(prefix, "_delta_dens0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_delta_dens0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of density_fldline_avg
             plot(legend=legend)
@@ -1519,7 +1215,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                 @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")
+            outfile = string(prefix, "_$(label)_fldline_avg_dens_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the deviation from conservation of density_fldline_avg
             plot(legend=legend)
@@ -1527,7 +1223,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                 @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")
+            outfile = string(prefix, "_$(label)_conservation_dens_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         upar_min = minimum(minimum(upar[:,is,:]) for upar ∈ parallel_flow)
@@ -1538,14 +1234,14 @@ function plot_moments(density, delta_density, density_fldline_avg,
             for (t, upar, run_label) ∈ zip(time, parallel_flow, run_names)
                 @views plot!(t, upar[iz0,is,:], label=run_label)
             end
-            outfile = string(prefix, "_upar0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_upar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
             plot(legend=legend)
             for (t, dupar, run_label) ∈ zip(time, delta_upar, run_names)
                 @views plot!(t, abs.(du[iz0,is,:]), yaxis=:log, label=run_label)
             end
-            outfile = string(prefix, "_delta_upar0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_delta_upar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of ppar_fldline_avg
             plot(legend=legend)
@@ -1554,7 +1250,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                              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")
+            outfile = string(prefix, "_$(label)_fldline_avg_upar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         ppar_min = minimum(minimum(ppar[:,is,:]) for ppar ∈ parallel_pressure)
@@ -1565,14 +1261,14 @@ function plot_moments(density, delta_density, density_fldline_avg,
             for (t, ppar, run_label) ∈ zip(time, parallel_pressure, run_names)
                 @views plot!(t, ppar[iz0,is,:], label=run_label)
             end
-            outfile = string(prefix, "_ppar0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_ppar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
             plot(legend=legend)
             for (t, dppar, run_label) ∈ zip(time, delta_ppar, run_names)
                 @views plot!(t, abs.(dppar[iz0,is,:]), yaxis=:log, label=run_label)
             end
-            outfile = string(prefix, "_delta_ppar0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_delta_ppar0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of ppar_fldline_avg
             plot(legend=legend)
@@ -1580,7 +1276,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
                 @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")
+            outfile = string(prefix, "_$(label)_fldline_avg_ppar_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
         end
         vth_min = minimum(minimum(vth[:,is,:]) for vth ∈ thermal_speed)
@@ -1591,14 +1287,14 @@ function plot_moments(density, delta_density, density_fldline_avg,
             for (t, vth, run_label) ∈ zip(time, thermal_speed, run_names)
                 @views plot!(t, vth[iz0,is,:], label=run_label)
             end
-            outfile = string(prefix, "_vth0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_vth0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of n_s(z=z0)-density_fldline_avg
             plot(legend=legend)
             for (t, dvth, run_label) ∈ zip(time, delta_vth, run_names)
                 @views plot!(t, abs.(dvth[iz0,is,:]), yaxis=:log, label=run_label)
             end
-            outfile = string(prefix, "_delta_vth0_vs_t_spec", spec_string, ".pdf")
+            outfile = string(prefix, "_$(label)_delta_vth0_vs_t_spec", spec_string, ".pdf")
             savefig(outfile)
             # plot the time trace of vth_fldline_avg
             plot(legend=legend)
@@ -1606,137 +1302,510 @@ function plot_moments(density, delta_density, density_fldline_avg,
                 @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)
+            outfile = string(prefix, "_$(label)_fldline_avg_vth_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+        end
+        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)
+            plot(legend=legend)
+            for (t, qpar, run_label) ∈ zip(time, parallel_heat_flux, run_names)
+                @views plot!(t, qpar[iz0,is,:], label=run_label)
+            end
+            outfile = string(prefix, "_$(label)_qpar0_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+            # plot the time trace of n_s(z=z0)-density_fldline_avg
+            plot(legend=legend)
+            for (t, dqpar, run_label) ∈ zip(time, delta_qpar, run_names)
+                @views plot!(t, abs.(dqpar[iz0,is,:]), yaxis=:log, label=run_label)
+            end
+            outfile = string(prefix, "_$(label)_delta_qpar0_vs_t_spec", spec_string, ".pdf")
+            savefig(outfile)
+            # plot the time trace of ppar_fldline_avg
+            plot(legend=legend)
+            for (t, qpar_avg, run_label) ∈ zip(time, qpar_fldline_avg, run_names)
+                @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, "_$(label)_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)
+            subplots = (heatmap(t, this_z.grid, n[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            outfile = string(prefix, "_$(label)_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)
+            subplots = (heatmap(t, this_z.grid, upar[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow,
+                                                              run_names))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            outfile = string(prefix, "_$(label)_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)
+            subplots = (heatmap(t, this_z.grid, ppar[:,is,:], xlabel="time", ylabel="z",
+                                title=run_label, c = :deep)
+                        for (t, this_z, ppar, run_label) ∈
+                            zip(time, z, parallel_pressure, run_names))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            outfile = string(prefix, "_$(label)_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)
+            subplots = (heatmap(t, this_z.grid, 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_names))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            outfile = string(prefix, "_$(label)_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
+                plot(legend=legend)
+                for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names)
+                    @views plot!(this_z.grid, n[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ",
+                                 ylims=(dens_min, dens_max), label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_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
+                plot(legend=legend)
+                for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow, run_names)
+                    @views plot!(this_z.grid, upar[:,is,i], xlabel="z", ylabel="upars/vt",
+                                 ylims=(upar_min, upar_max), label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_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
+                plot(legend=legend)
+                for (t, this_z, ppar, run_label) ∈ zip(time, z, parallel_pressure, run_names)
+                    @views plot!(this_z.grid, ppar[:,is,i], xlabel="z", ylabel="ppars",
+                                 ylims=(ppar_min, ppar_max), label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_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
+                plot(legend=legend)
+                for (t, this_z, vth, run_label) ∈ zip(time, z, thermal_speed, run_names)
+                    @views plot!(this_z.grid, vth[:,is,i], xlabel="z", ylabel="vths",
+                                 ylims=(vth_min, vth_max), label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_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
+                plot(legend=legend)
+                for (t, this_z, qpar, run_label) ∈ zip(time, z, parallel_heat_flux,
+                                                       run_names)
+                    @views plot!(this_z.grid, qpar[:,is,i], xlabel="z", ylabel="qpars",
+                                 ylims=(qpar_min, qpar_max), label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_qpar_vs_z_spec", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
+    end
+    println("done.")
+end
+
+"""
+"""
+function plot_dfns(density, parallel_flow, parallel_pressure, thermal_speed, ff,
+                   dens_moment, upar_moment, ppar_moment, time_pdfs, n_species, z, vpa,
+                   evolve_density, evolve_upar, evolve_ppar, run_names, itime_min_pdfs,
+                   itime_max_pdfs, nwrite_movie_pdfs, iz0, label)
+    println("Plotting distribution function data...")
+    cmlog(cmlin::ColorGradient) = RGB[cmlin[x] for x=LinRange(0,1,30)]
+    logdeep = cgrad(:deep, scale=:log) |> cmlog
+    n_species_max = maximum(n_species)
+    n_runs = length(run_names)
+    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
+        plot(legend=legend)
+        for (t, n, n_int, run_label) ∈ zip(time_pdfs, density, dens_moment, run_names)
+            @views plot!(t, n[iz0,is,:] .- n_int[iz0,is,:], label=run_label)
+        end
+        outfile = string(prefix, "_$(label)_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
+        plot(legend=legend)
+        for (t, upar_int, run_label) ∈ zip(time_pdfs, upar_moment, run_names)
+            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(prefix, "_$(label)_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
+        plot(legend=legend)
+        for (t, ppar, ppar_int, run_label) ∈ zip(time_pdfs, parallel_pressure, ppar_moment, run_names)
+            @views plot(t, ppar[iz0,is,:] .- ppar_int[iz0,is,:], label=run_label)
+        end
+        outfile = string(prefix, "_$(label)_intw2f0_vs_t", spec_string, ".pdf")
+        savefig(outfile)
+        #fmin = minimum(ff[:,:,is,:])
+        #fmax = maximum(ff[:,:,is,:])
+        if pp.animate_f_vs_vpa_z
+            # make a gif animation of ln f(vpa,z,t)
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
+                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            end
+            outfile = string(prefix, "_$(label)_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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
+                subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            end
+            outfile = string(prefix, "_$(label)_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")
+            subplots = (@views heatmap(this_z.grid, this_vpa.grid, 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_names))
+            plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            outfile = string(prefix, "_$(label)_f_vs_z_vpa_final", spec_string, ".pdf")
+            savefig(outfile)
+
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                plot(legend=legend)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
+                    @views plot!(this_vpa.grid, f[:,1,is,i], xlabel="vpa", ylabel="f(z=0)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_f0_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                plot(legend=legend)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
+                    @views plot!(this_vpa.grid, f[:,end,is,i], xlabel="vpa", ylabel="f(z=L)",
+                                 label=run_label)
+                end
+            end
+            outfile = string(prefix, "_$(label)_fL_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
+        if pp.plot_f_unnormalized_vs_vpa_z
+            PyPlot.clf()
+            fig = PyPlot.figure(1, figsize=(6*n_runs,4))
+            it = itime_max_pdfs
+            # 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_names)
+
+                PyPlot.subplot(1, n_runs, run_ind)
+                @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
+                    f[:,:,is,it], this_z.grid, this_vpa.grid, n[:,is,it],
+                    upar[:,is,it], vth[:,is,it], ev_n, ev_u, ev_p)
+                plot_unnormalised_f2d(f_unnorm, z2d, dzdt2d; title=run_label,
+                                      plot_log=false)
+            end
+            PyPlot.savefig(string(prefix, "_f_unnorm_vs_vpa_z", spec_string, ".pdf"))
+            PyPlot.clf()
+
+            plot(legend=legend)
+            it = itime_max_pdfs
+            # 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_names)
+
+                @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                    f[:,1,is,it], this_vpa.grid, n[1,is,it], upar[1,is,it], vth[1,is,it],
+                    ev_n, ev_u, ev_p)
+                @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=0)",
+                             label=run_label)
+            end
+            savefig(string(prefix, "_f0_unnorm_vs_vpa", spec_string, ".pdf"))
+
+            plot(legend=legend)
+            it = itime_max_pdfs
+            # 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_names)
+
+                @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                    f[:,end,is,it], this_vpa.grid, n[end,is,it], upar[end,is,it], vth[end,is,it],
+                    ev_n, ev_u, ev_p)
+                @views plot!(dzdt, f_unnorm, xlabel="vpa", ylabel="f_unnorm(z=L)",
+                             label=run_label)
+            end
+            savefig(string(prefix, "_fL_unnorm_vs_vpa", spec_string, ".pdf"))
         end
-        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)
-            plot(legend=legend)
-            for (t, qpar, run_label) ∈ zip(time, parallel_heat_flux, run_names)
-                @views plot!(t, qpar[iz0,is,:], label=run_label)
+        if pp.animate_f_unnormalized
+            ## 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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid,
+            #                           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_names))
+            #    plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            #end
+            #outfile = string(prefix, "_$(label)_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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+            #    subplots = (@views plot_unnormalised(f[:,:,is,i], this_z.grid, this_vpa.grid, 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_names))
+            #    plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            #end
+            #outfile = string(prefix, "_$(label)_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_pdfs:nwrite_movie_pdfs:itime_max_pdfs)
+            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_names)
+
+                    PyPlot.subplot(1, n_runs, run_ind)
+                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
+                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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
-            outfile = string(prefix, "_qpar0_vs_t_spec", spec_string, ".pdf")
-            savefig(outfile)
-            # plot the time trace of n_s(z=z0)-density_fldline_avg
-            plot(legend=legend)
-            for (t, dqpar, run_label) ∈ zip(time, delta_qpar, run_names)
-                @views plot!(t, abs.(dqpar[iz0,is,:]), yaxis=:log, label=run_label)
+            fig = PyPlot.figure(1, figsize=(6*n_runs,4))
+            myanim = matplotlib_animation.FuncAnimation(fig, make_frame, frames=nframes)
+            outfile = string(prefix, "_$(label)_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()
+                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_names)
+
+                    PyPlot.subplot(1, n_runs, run_ind)
+                    @views f_unnorm, z2d, dzdt2d = get_unnormalised_f_coords_2d(
+                        f[:,:,is,iframe], this_z.grid, this_vpa.grid, 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
-            outfile = string(prefix, "_delta_qpar0_vs_t_spec", spec_string, ".pdf")
-            savefig(outfile)
-            # plot the time trace of ppar_fldline_avg
-            plot(legend=legend)
-            for (t, qpar_avg, run_label) ∈ zip(time, qpar_fldline_avg, run_names)
-                @views plot!(t, qpar_avg[is,:], xlabel="time", ylabel="<qpars/NₑTₑvth>",
-                             ylims=(qpar_min,qpar_max), label=run_label)
+            fig = PyPlot.figure(figsize=(6*n_runs,4))
+            myanim = matplotlib_animation.FuncAnimation(fig, make_frame_log, frames=nframes)
+            outfile = string(prefix, "_$(label)_logf_unnorm_vs_vpa_z", spec_string, ".gif")
+            myanim.save(outfile, writer=matplotlib_animation.PillowWriter(fps=30))
+
+            # Ensure PyPlot figure is cleared
+            closeall()
+
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                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_names)
+                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                        f[:,1,is,i], this_vpa.grid, 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, "_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)
-            subplots = (heatmap(t, this_z.grid, n[:,is,:], xlabel="time", ylabel="z",
-                                title=run_label, c = :deep)
-                        for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names))
-            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
-        if pp.plot_upar_vs_z_t
-            # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z.grid, upar[:,is,:], xlabel="time", ylabel="z",
-                                title=run_label, c = :deep)
-                        for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow,
-                                                              run_names))
-            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
-        if pp.plot_ppar_vs_z_t
-            # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z.grid, ppar[:,is,:], xlabel="time", ylabel="z",
-                                title=run_label, c = :deep)
-                        for (t, this_z, ppar, run_label) ∈
-                            zip(time, z, parallel_pressure, run_names))
-            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
-        if pp.plot_qpar_vs_z_t
-            # make a heatmap plot of upar_s(z,t)
-            subplots = (heatmap(t, this_z.grid, 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_names))
-            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
-        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
+            outfile = string(prefix, "_$(label)_f0_unnorm_vs_vpa", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
-                for (t, this_z, n, run_label) ∈ zip(time, z, density, run_names)
-                    @views plot!(this_z.grid, n[:,is,i], xlabel="z", ylabel="nᵢ/Nₑ",
-                                 ylims=(dens_min, dens_max), label=run_label)
+                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_names)
+                    @views f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(
+                        f[:,end,is,i], this_vpa.grid, 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, "_dens_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_$(label)_fL_unnorm_vs_vpa", 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
+        if pp.animate_deltaf_vs_vpa_z
+            # make a gif animation of δf(vpa,z,t)
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                subplots = (@views heatmap(this_z, this_vpa.grid, 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_names))
+                plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+            end
+            outfile = string(prefix, "_$(label)_deltaf_vs_vpa_z", spec_string, ".gif")
+            gif(anim, outfile, fps=5)
+        end
+        if pp.animate_f_vs_vpa_z0
+            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
-                for (t, this_z, upar, run_label) ∈ zip(time, z, parallel_flow, run_names)
-                    @views plot!(this_z.grid, upar[:,is,i], xlabel="z", ylabel="upars/vt",
-                                 ylims=(upar_min, upar_max), label=run_label)
+                for (f, this_z, run_label) ∈ zip(ff, z, run_names)
+                    @views plot!(this_z, f[ivpa0,:,is,i], ylims = (fmin,fmax),
+                                 label=run_label)
                 end
             end
-            outfile = string(prefix, "_upar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_$(label)_f_vs_z", 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
+        if pp.animate_deltaf_vs_vpa_z0
+            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
-                for (t, this_z, ppar, run_label) ∈ zip(time, z, parallel_pressure, run_names)
-                    @views plot!(this_z.grid, ppar[:,is,i], xlabel="z", ylabel="ppars",
-                                 ylims=(ppar_min, ppar_max), label=run_label)
+                for (df, this_z, run_label) ∈ zip(delta_f, z, run_names)
+                    @views plot!(this_z, df[ivpa0,:,is,i], ylims = (fmin,fmax),
+                                 label=run_label)
                 end
             end
-            outfile = string(prefix, "_ppar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_$(label)_deltaf_vs_z", 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
+        if pp.animate_f_vs_vpa_z0
+            fmin = minimum(minimum(f[:,iz0,is,:]) for f ∈ ff)
+            fmax = maximum(maximum(f[:,iz0,is,:]) for f ∈ ff)
+
+            # if is == 1
+            #     tmp = copy(ff)
+            #     @. tmp[:,1,1,:] /= vpa^2
+            #     bohm_integral = copy(time)
+            #     for i ∈ 1:ntime
+            #         @views bohm_integral[i] = integrate_over_vspace(tmp[1:cld(nvpa,2)-1,1,1,i],vpa_wgts[1:cld(nvpa,2)-1])/2.0
+            #     end
+            #     plot(time, bohm_integral, xlabel="time", label="Bohm integral")
+            #     plot!(time, density[1,1,:], label="nᵢ(zmin)")
+            #     outfile = string(prefix, "_$(label)_Bohm_criterion.pdf")
+            #     savefig(outfile)
+            #     println()
+            #     if bohm_integral[end] <= density[1,1,end]
+            #         println("Bohm criterion: ", bohm_integral[end], " <= ", density[1,1,end], " is satisfied!")
+            #     else
+            #         println("Bohm criterion: ", bohm_integral[end], " <= ", density[1,1,end], " is not satisfied!")
+            #     end
+            #     println()
+            #     for j ∈ 0:10
+            #         println("j: ", j, "  Bohm integral: ", integrate_over_vspace(tmp[1:cld(nvpa,2)-j,1,1,end],vpa_wgts[1:cld(nvpa,2)-j,end])/2.0)
+            #     end
+            # end
+            # make a gif animation of f(vpa,z0,t)
+            anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                #@views plot(vpa, ff[iz0,:,is,i], ylims = (fmin,fmax))
                 plot(legend=legend)
-                for (t, this_z, vth, run_label) ∈ zip(time, z, thermal_speed, run_names)
-                    @views plot!(this_z.grid, vth[:,is,i], xlabel="z", ylabel="vths",
-                                 ylims=(vth_min, vth_max), label=run_label)
+                for (f, this_vpa, run_label) ∈ zip(ff, vpa, run_names)
+                    @views plot!(this_vpa.grid, f[:,iz0,is,i], label=run_label)
                 end
             end
-            outfile = string(prefix, "_vth_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_$(label)_f_vs_vpa", 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
+        if pp.animate_deltaf_vs_vpa_z0
+            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_pdfs:nwrite_movie_pdfs:itime_max_pdfs
                 plot(legend=legend)
-                for (t, this_z, qpar, run_label) ∈ zip(time, z, parallel_heat_flux,
-                                                       run_names)
-                    @views plot!(this_z.grid, qpar[:,is,i], xlabel="z", ylabel="qpars",
-                                 ylims=(qpar_min, qpar_max), label=run_label)
+                for (df, this_vpa, fn, fx, run_label) ∈
+                        zip(delta_f, vpa, fmin, fmax, run_names)
+                    @views plot!(this_vpa.grid, delta_f[:,iz0,is,i], ylims = (fn,fx),
+                                label=run_label)
                 end
             end
-            outfile = string(prefix, "_qpar_vs_z_spec", spec_string, ".gif")
+            outfile = string(prefix, "_$(label)_deltaf_vs_vpa", spec_string, ".gif")
             gif(anim, outfile, fps=5)
         end
     end
-    println("done.")
+    return nothing
 end
 
 """

From 87b82f69461b709a4900277296d149c08604b2d8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 25 Apr 2023 09:48:13 +0100
Subject: [PATCH 025/428] Create 'global' coordinates in post-processing

Fixes post-processing when using distributed-memory MPI.
---
 src/plot_MMS_sequence.jl |   5 +-
 src/plot_sequence.jl     |  17 +++---
 src/post_processing.jl   | 123 ++++++++++++++++-----------------------
 3 files changed, 60 insertions(+), 85 deletions(-)

diff --git a/src/plot_MMS_sequence.jl b/src/plot_MMS_sequence.jl
index d006ea10c..51cb75d10 100644
--- a/src/plot_MMS_sequence.jl
+++ b/src/plot_MMS_sequence.jl
@@ -16,7 +16,7 @@ using LaTeXStrings
 using ..post_processing_input: pp
 using ..post_processing: compare_charged_pdf_symbolic_test, compare_fields_symbolic_test
 using ..post_processing: compare_moments_symbolic_test, compare_neutral_pdf_symbolic_test
-using ..post_processing: read_distributed_zr_data!, construct_global_zr_grids
+using ..post_processing: read_distributed_zr_data!, construct_global_zr_coords
 using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_charged_moments
 using ..post_processing: allocate_global_zr_fields#, get_geometry_and_composition, get_coords_nelement
 using ..array_allocation: allocate_float
@@ -174,8 +174,7 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         end 
         # read in the data from different block netcdf files
         # grids 
-        z, z_wgts, r, r_wgts = construct_global_zr_grids(run_name,
-           "moments",z.n_global,r.n_global,nblocks)
+        r_global, z_global = construct_global_zr_coords(r, z)
         #println("z: ",z)
         #println("r: ",r)
         
diff --git a/src/plot_sequence.jl b/src/plot_sequence.jl
index 914c112e9..370a13b16 100644
--- a/src/plot_sequence.jl
+++ b/src/plot_sequence.jl
@@ -11,7 +11,7 @@ using Statistics: mean
 using SpecialFunctions: erfi
 using LaTeXStrings
 # modules
-using ..post_processing: read_distributed_zr_data!, construct_global_zr_grids
+using ..post_processing: read_distributed_zr_data!, construct_global_zr_coords
 using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_charged_moments
 using ..post_processing: allocate_global_zr_fields#, get_coords_nelement
 using ..array_allocation: allocate_float
@@ -54,8 +54,8 @@ function plot_sequence_fields_data(path_list)
         # load local sizes of grids stored on each netCDF file 
         # z z_wgts r r_wgts may take different values on different blocks
         # we need to construct the global grid below
-        nz_local, nz_global, z_local, z_wgts, Lz = load_coordinate_data(fid, "z")
-        nr_local, nr_global, r_local, r_wgts, Lr = load_coordinate_data(fid, "r")
+        z_local, z_local_spectral = load_coordinate_data(fid, "z")
+        r_local, r_local_spectral = load_coordinate_data(fid, "r")
         # load time data 
         ntime, time = load_time_data(fid)
         # load species data 
@@ -66,17 +66,16 @@ function plot_sequence_fields_data(path_list)
         phi, Ez, Er = allocate_global_zr_fields(nz_global,nr_global,ntime)
         # read in the data from different block netcdf files
         # grids 
-        z, z_wgts, r, r_wgts = construct_global_zr_grids(run_name,
-           "moments",nz_global,nr_global,nblocks)
+        r, r_spectral, z, z_spectral = construct_global_zr_coords(r_local, z_local)
         # fields 
-        read_distributed_zr_data!(phi,"phi",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(Ez,"Ez",run_name,"moments",nblocks,nz_local,nr_local) 
-        read_distributed_zr_data!(Er,"Er",run_name,"moments",nblocks,nz_local,nr_local) 
+        read_distributed_zr_data!(phi,"phi",run_name,"moments",nblocks,z_local.n,r_local.n)
+        read_distributed_zr_data!(Ez,"Ez",run_name,"moments",nblocks,z_local.n,r_local.n)
+        read_distributed_zr_data!(Er,"Er",run_name,"moments",nblocks,z_local.n,r_local.n)
         
         # store phi, Ez at ir = 1 and itime = end
         push!(phi_z_list,phi[:,1,end])
         push!(Ez_z_list,Ez[:,1,end])
-        push!(z_list,z[:])
+        push!(z_list,z.grid[:])
         push!(phi_sheath_entrance_list,phi[1,1,end])
         push!(Ez_sheath_entrance_list,Ez[1,1,end])
     end
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 8939b07ca..953907d69 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -7,7 +7,7 @@ export compare_charged_pdf_symbolic_test
 export compare_moments_symbolic_test
 export compare_neutral_pdf_symbolic_test
 export compare_fields_symbolic_test
-export construct_global_zr_grids
+export construct_global_zr_coords
 export allocate_global_zr_charged_moments
 export allocate_global_zr_neutral_moments
 export allocate_global_zr_fields
@@ -26,10 +26,12 @@ using Statistics: mean
 using SpecialFunctions: erfi
 using LaTeXStrings
 using Measures
+using MPI
 # modules
 using ..post_processing_input: pp
 using ..quadrature: composite_simpson_weights
 using ..array_allocation: allocate_float
+using ..coordinates: define_coordinate
 using ..file_io: open_ascii_output_file
 using ..type_definitions: mk_float, mk_int
 using ..initial_conditions: vpagrid_to_dzdt
@@ -46,7 +48,7 @@ using ..analysis: analyze_fields_data, analyze_moments_data, analyze_pdf_data,
 using ..velocity_moments: integrate_over_vspace
 using ..manufactured_solns: manufactured_solutions, manufactured_electric_fields
 using ..moment_kinetics_input: mk_input, get
-using ..input_structs: geometry_input, species_composition
+using ..input_structs: geometry_input, grid_input, species_composition
 using ..input_structs: electron_physics_type, boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath
 using TOML
 import Base: get
@@ -192,43 +194,19 @@ function iglobal_func(ilocal,irank,nlocal)
     return iglobal
 end
 
-function construct_global_zr_grids(run_name::String,file_key::String,nz_global::mk_int,nr_global::mk_int,nblocks::mk_int)
-    zgrid = allocate_float(nz_global)
-    zwgts = allocate_float(nz_global)
-    rgrid = allocate_float(nr_global)
-    rwgts = allocate_float(nr_global)
+function construct_global_zr_coords(r_local, z_local)
 
-    # current routine loops over blocks
-    # whereas the optimal routine would loop over a single z/r group
-    for iblock in 0:nblocks-1
-        fid = open_readonly_output_file(run_name,file_key,iblock=iblock,printout=false)
-        z, z_spectral = load_coordinate_data(fid, "z")
-        r, r_spectral = load_coordinate_data(fid, "r")
+    function make_global_input(coord_local)
+        return grid_input(coord_local.name, coord_local.ngrid,
+            coord_local.nelement_global, coord_local.nelement_global, 1, 0, coord_local.L,
+            coord_local.discretization, coord_local.fd_option, coord_local.bc,
+            coord_local.advection, MPI.COMM_NULL)
+    end
 
-        z_irank, r_irank = load_rank_data(fid)
-        # MRH should wgts at element boundaries be doubled
-        # MRH in the main code duplicated points have half the integration wgt
-        if z_irank == 0
-            zgrid[1] = z.grid[1]
-            zwgts[1] = z.wgts[1]
-        end
-        for iz_local in 2:z.n
-            iz_global = iglobal_func(iz_local,z_irank,z.n)
-            zgrid[iz_global] = z.grid[iz_local]
-            zwgts[iz_global] = z.wgts[iz_local]
-        end
+    r_global, r_global_spectral = define_coordinate(make_global_input(r_local))
+    z_global, z_global_spectral = define_coordinate(make_global_input(z_local))
 
-        if r_irank == 0
-            rgrid[1] = r.grid[1]
-            rwgts[1] = r.wgts[1]
-        end
-        for ir_local in 2:r.n
-            ir_global = iglobal_func(ir_local,r_irank,r.n)
-            rgrid[ir_global] = r.grid[ir_local]
-            rwgts[ir_global] = r.wgts[ir_local]
-        end
-    end
-    return zgrid, zwgts, rgrid, rwgts
+    return r_global, r_global_spectral, z_global, z_global_spectral
 end
 
 """
@@ -438,11 +416,8 @@ function analyze_and_plot_data(prefix...)
     end
     # read in the data from different block files
     # grids
-    z_global, z_wgts, r_global, r_wgts =
-        get_tuple_of_return_values(construct_global_zr_grids, run_names, "moments",
-                                   Tuple(this_z.n_global for this_z ∈ z),
-                                   Tuple(this_r.n_global for this_r ∈ r),
-                                   nblocks)
+    r_global, r_global_spectral, z_global, z_global_spectral =
+        get_tuple_of_return_values(construct_global_zr_coords, r, z)
 
     # fields
     get_tuple_of_return_values(read_distributed_zr_data!, phi, "phi", run_names, "moments",
@@ -652,7 +627,7 @@ function analyze_and_plot_data(prefix...)
         #evaluate 1D-1V diagnostics at fixed ir0
         plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
             nwrite_movie_pdfs, itime_min_pdfs, itime_max_pdfs, ivpa0, iz0,
-            ir0, r,
+            ir0, r_global,
             Tuple(p[:,ir0,:] for p ∈ phi),
             Tuple(n[:,ir0,:,:] for n ∈ density),
             Tuple(upar[:,ir0,:,:] for upar ∈ parallel_flow),
@@ -678,7 +653,7 @@ function analyze_and_plot_data(prefix...)
             Tuple(neutral_vth[:,ir0,:,:] for neutral_vth ∈ neutral_thermal_speed_at_pdf_times),
             Tuple(neutral_f[:,ivr0,ivzeta0,:,ir0,:,:] for neutral_f ∈ neutral_ff),
             n_ion_species, n_neutral_species, evolve_density, evolve_upar, evolve_ppar,
-            vz, vpa, z, ntime, time, ntime_pdfs, time_pdfs)
+            vz, vpa, z_global, ntime, time, ntime_pdfs, time_pdfs)
     end
 
     diagnostics_2d = false
@@ -686,9 +661,9 @@ function analyze_and_plot_data(prefix...)
         # analyze the fields data
         phi_iz0 = Tuple(p[iz0,:,:] for p ∈ phi)
         phi_fldline_avg, delta_phi =
-            get_tuple_of_return_values(analyze_fields_data, phi_iz0, ntime, r)
+            get_tuple_of_return_values(analyze_fields_data, phi_iz0, ntime, r_global)
         get_tuple_of_return_values(plot_fields_rt, phi_iz0, delta_phi, time, itime_min,
-                                   itime_max, nwrite_movie, r, ir0, run_name, delta_phi,
+                                   itime_max, nwrite_movie, r_global, ir0, run_name, delta_phi,
                                    pp)
     end
 
@@ -723,6 +698,8 @@ function analyze_and_plot_data(prefix...)
     time = time[1]
     z = z[1]
     r = r[1]
+    z_global = z_global[1]
+    r_global = r_global[1]
     n_ion_species = n_ion_species[1]
     run_name = run_names[1]
     phi = phi[1]
@@ -734,12 +711,12 @@ function analyze_and_plot_data(prefix...)
     composition = composition[1]
 
     # make plots and animations of the phi, Ez and Er
-    plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time, z.grid, r.grid, iz0,
-                            ir0, n_ion_species, itime_min, itime_max, nwrite_movie,
-                            run_name, pp)
+    plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time,
+                            z_global.grid, r_global.grid, iz0, ir0, n_ion_species,
+                            itime_min, itime_max, nwrite_movie, run_name, pp)
     # make plots and animations of the phi, Ez and Er
-    plot_fields_2D(phi, Ez, Er, time, z.grid, r.grid, iz0, ir0, itime_min, itime_max, nwrite_movie,
-                   run_name, pp, "")
+    plot_fields_2D(phi, Ez, Er, time, z_global.grid, r_global.grid, iz0, ir0, itime_min,
+                   itime_max, nwrite_movie, run_name, pp, "")
     # make plots and animations of the ion pdf
     # only if ntime == ntime_pdfs & data on one shared memory process
     if ntime == ntime_pdfs && r.n_global == r.n && z.n_global == z.n
@@ -749,14 +726,14 @@ function analyze_and_plot_data(prefix...)
         read_distributed_zr_data!(ff, "f", run_names, "dfns", nblocks, z.n, r.n)
 
         spec_type = "ion"
-        plot_charged_pdf(ff, vpa_local, vperp_local, z.grid, r.grid, ivpa0, ivperp0,
-                         iz0, ir0, spec_type, n_ion_species, itime_min_pdfs,
+        plot_charged_pdf(ff, vpa_local, vperp_local, z_global.grid, r_global.grid, ivpa0,
+                         ivperp0, iz0, ir0, spec_type, n_ion_species, itime_min_pdfs,
                          itime_max_pdfs, nwrite_movie_pdfs, run_name, pp)
         # make plots and animations of the neutral pdf
         if n_neutral_species > 0
             spec_type = "neutral"
-            plot_neutral_pdf(neutral_ff, vz.grid, vr.grid, vzeta.grid, z.grid,
-                             r.grid, ivz0, ivr0, ivzeta0, iz0, ir0, spec_type,
+            plot_neutral_pdf(neutral_ff, vz.grid, vr.grid, vzeta.grid, z_global.grid,
+                             r_global.grid, ivz0, ivr0, ivzeta0, iz0, ir0, spec_type,
                              n_neutral_species, itime_min_pdfs, itime_max_pdfs,
                              nwrite_movie_pdfs, run_name, pp)
         end
@@ -774,20 +751,20 @@ function analyze_and_plot_data(prefix...)
         r_bc = get(scan_input, "r_bc", "periodic")
         z_bc = get(scan_input, "z_bc", "periodic")
         # avoid passing Lr = 0 into manufactured_solns functions
-        if r.n > 1
-            Lr_in = r.L
+        if r_global.n > 1
+            Lr_in = r_global.L
         else
             Lr_in = 1.0
         end
 
         #geometry, composition = get_geometry_and_composition(scan_input,n_ion_species,n_neutral_species)
 
-        manufactured_solns_list = manufactured_solutions(Lr_in,z.L,r_bc,z_bc,geometry,composition,r.n)
+        manufactured_solns_list = manufactured_solutions(Lr_in,z_global.L,r_bc,z_bc,geometry,composition,r_global.n)
         dfni_func = manufactured_solns_list.dfni_func
         densi_func = manufactured_solns_list.densi_func
         dfnn_func = manufactured_solns_list.dfnn_func
         densn_func = manufactured_solns_list.densn_func
-        manufactured_E_fields = manufactured_electric_fields(Lr_in,z.L,r_bc,z_bc,composition,r.n)
+        manufactured_E_fields = manufactured_electric_fields(Lr_in,z_global.L,r_bc,z_bc,composition,r_global.n)
         Er_func = manufactured_E_fields.Er_func
         Ez_func = manufactured_E_fields.Ez_func
         phi_func = manufactured_E_fields.phi_func
@@ -797,36 +774,36 @@ function analyze_and_plot_data(prefix...)
         Er_sym = copy(phi[:,:,:])
         Ez_sym = copy(phi[:,:,:])
         for it in 1:ntime
-            for ir in 1:nr_global
-                for iz in 1:nz_global
-                    phi_sym[iz,ir,it] = phi_func(z_grid[iz],r_grid[ir],time[it])
-                    Ez_sym[iz,ir,it] = Ez_func(z_grid[iz],r_grid[ir],time[it])
-                    Er_sym[iz,ir,it] = Er_func(z_grid[iz],r_grid[ir],time[it])
+            for ir in 1:r_global.n
+                for iz in 1:z_global.n
+                    phi_sym[iz,ir,it] = phi_func(z_global.grid[iz],r_global.grid[ir],time[it])
+                    Ez_sym[iz,ir,it] = Ez_func(z_global.grid[iz],r_global.grid[ir],time[it])
+                    Er_sym[iz,ir,it] = Er_func(z_global.grid[iz],r_global.grid[ir],time[it])
                 end
             end
         end
         # make plots and animations of the phi, Ez and Er
-        #plot_fields_2D(phi_sym, Ez_sym, Er_sym, time, z_grid, r_grid, iz0, ir0,
+        #plot_fields_2D(phi_sym, Ez_sym, Er_sym, time, z_global.grid, r_global.grid, iz0, ir0,
         #    itime_min, itime_max, nwrite_movie, run_name, pp, "_sym")
         println("time/ (Lref/cref): ", time)
-        compare_fields_symbolic_test(run_name,phi,phi_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,phi,phi_sym,z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{\phi}",L"\widetilde{\phi}^{sym}",L"\varepsilon(\widetilde{\phi})","phi")
-        compare_fields_symbolic_test(run_name,Er,Er_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,Er,Er_sym,z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{E}_r",L"\widetilde{E}^{sym}_r",L"\varepsilon(\widetilde{E}_r)","Er")
-        compare_fields_symbolic_test(run_name,Ez,Ez_sym,z_grid,r_grid,time,nz_global,nr_global,ntime,
+        compare_fields_symbolic_test(run_name,Ez,Ez_sym,z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{E}_z",L"\widetilde{E}^{sym}_z",L"\varepsilon(\widetilde{E}_z)","Ez")
 
         # ion test
         density_sym = copy(density[:,:,:,:])
         is = 1
         for it in 1:ntime
-            for ir in 1:nr_global
-                for iz in 1:nz_global
-                    density_sym[iz,ir,is,it] = densi_func(z_grid[iz],r_grid[ir],time[it])
+            for ir in 1:r_global.n
+                for iz in 1:z_global.n
+                    density_sym[iz,ir,is,it] = densi_func(z_global.grid[iz],r_global.grid[ir],time[it])
                 end
             end
         end
-        compare_moments_symbolic_test(run_name,density,density_sym,"ion",z_grid,r_grid,time,nz_global,nr_global,ntime,
+        compare_moments_symbolic_test(run_name,density,density_sym,"ion",z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{n}_i",L"\widetilde{n}_i^{sym}",L"\varepsilon(\widetilde{n}_i)","dens")
 
         compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
@@ -838,11 +815,11 @@ function analyze_and_plot_data(prefix...)
             for it in 1:ntime
                 for ir in 1:nr_global
                     for iz in 1:nz_global
-                        neutral_density_sym[iz,ir,is,it] = densn_func(z_grid[iz],r_grid[ir],time[it])
+                        neutral_density_sym[iz,ir,is,it] = densn_func(z_global.grid[iz],r_global.grid[ir],time[it])
                     end
                 end
             end
-            compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z_grid,r_grid,time,nz_global,nr_global,ntime,
+            compare_moments_symbolic_test(run_name,neutral_density,neutral_density_sym,"neutral",z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
              L"\widetilde{n}_n",L"\widetilde{n}_n^{sym}",L"\varepsilon(\widetilde{n}_n)","dens")
 
             compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,"neutral",

From 6f23455179be9135e9ce5ad6367120db160594c7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 May 2023 11:59:49 +0100
Subject: [PATCH 026/428] Fix `irank` value output when using parallel I/O

Want the output file to have `irank=0` written for every coordinate, so
it's as if it was produced by a serial run. Otherwise the
post-processing will get confused when trying to load 'distributed'
data.
---
 src/file_io.jl | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index bc97d064a..abb3ca257 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -325,9 +325,15 @@ function define_io_coordinate!(parent, coord, coord_name, description, parallel_
         write_single_value!(group, "n_global", coord.n_global; parallel_io=parallel_io,
                             description="total number of $coord_name grid points")
 
-        # write the rank in the coord-direction of this process
-        write_single_value!(group, "irank", coord.irank, parallel_io=parallel_io,
-                            description="rank of this block in the $(coord.name) grid communicator")
+        if parallel_io
+            # write the rank as if whole file was written by rank-0
+            write_single_value!(group, "irank", 0, parallel_io=parallel_io,
+                                description="rank of this block in the $(coord.name) grid communicator")
+        else
+            # write the rank in the coord-direction of this process
+            write_single_value!(group, "irank", coord.irank, parallel_io=parallel_io,
+                                description="rank of this block in the $(coord.name) grid communicator")
+        end
 
         # write the global length of this coordinate to variable "L"
         # within "coords/coord_name" group

From fdadae73e931f2f9871ed03ea75b016cced8d47c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 May 2023 15:43:27 +0100
Subject: [PATCH 027/428] Save input settings to output files, and load them in
 post-processing

---
 src/file_io.jl         | 60 +++++++++++++++++++++++++++++++++++++-----
 src/file_io_hdf5.jl    | 12 +++++++++
 src/file_io_netcdf.jl  | 13 +++++++++
 src/load_data.jl       | 25 +++++++++++++++++-
 src/moment_kinetics.jl |  2 +-
 src/post_processing.jl | 14 +++++-----
 6 files changed, 112 insertions(+), 14 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index abb3ca257..d52f00bb0 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -102,7 +102,8 @@ function io_has_parallel() end
 open the necessary output files
 """
 function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vperp, z, r,
-                       composition, collisions, evolve_density, evolve_upar, evolve_ppar)
+                       composition, collisions, evolve_density, evolve_upar, evolve_ppar,
+                       input_dict)
     begin_serial_region()
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -124,12 +125,13 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
 
         io_moments = setup_moments_io(out_prefix, io_input.binary_format, r, z,
                                       composition, collisions, evolve_density,
-                                      evolve_upar, evolve_ppar, io_input.parallel_io,
-                                      comm_inter_block[])
+                                      evolve_upar, evolve_ppar, input_dict,
+                                      io_input.parallel_io, comm_inter_block[])
         io_dfns = setup_dfns_io(out_prefix, io_input.binary_format,
                                 boundary_distributions, r, z, vperp, vpa, vzeta, vr, vz,
                                 composition, collisions, evolve_density, evolve_upar,
-                                evolve_ppar, io_input.parallel_io, comm_inter_block[])
+                                evolve_ppar, input_dict, io_input.parallel_io,
+                                comm_inter_block[])
 
         return ascii, io_moments, io_dfns
     end
@@ -142,6 +144,21 @@ Get a (sub-)group from a file or group
 """
 function get_group() end
 
+"""
+Test if a member of a (sub-)group is a group
+"""
+function is_group() end
+
+"""
+Get names of all subgroups
+"""
+function get_subgroup_keys() end
+
+"""
+Get names of all variables
+"""
+function get_variable_keys() end
+
 """
     write_single_value!(file_or_group, name, value; description=nothing)
 
@@ -189,6 +206,30 @@ function write_overview!(fid, composition, collisions, parallel_io, evolve_densi
     end
 end
 
+"""
+Save info from the dict with input settings to the output file
+
+Note: assumes all keys in `input_dict` are strings.
+"""
+function write_input!(fid, input_dict, parallel_io)
+    function write_dict(io, section_dict, parallel_io)
+        # Function that can be called recursively to write nested Dicts into sub-groups in
+        # the output file
+        for (key, value) ∈ section_dict
+            if isa(value, Dict)
+                subsection_io = create_io_group(io, key)
+                write_dict(subsection_io, value, parallel_io)
+            else
+                write_single_value!(io, key, value, parallel_io=parallel_io)
+            end
+        end
+    end
+    @serial_region begin
+        input_io = create_io_group(fid, "input")
+        write_dict(input_io, input_dict, parallel_io)
+    end
+end
+
 """
 Write the distributions that may be used for boundary conditions to the output file
 """
@@ -542,7 +583,8 @@ end
 setup file i/o for moment variables
 """
 function setup_moments_io(prefix, binary_format, r, z, composition, collisions,
-                          evolve_density, evolve_upar, evolve_ppar, parallel_io, io_comm)
+                          evolve_density, evolve_upar, evolve_ppar, input_dict,
+                          parallel_io, io_comm)
     @serial_region begin
         moments_prefix = string(prefix, ".moments")
         if !parallel_io
@@ -557,6 +599,9 @@ function setup_moments_io(prefix, binary_format, r, z, composition, collisions,
         write_overview!(fid, composition, collisions, parallel_io, evolve_density,
                         evolve_upar, evolve_ppar)
 
+        # write the input settings
+        write_input!(fid, input_dict, parallel_io)
+
         ### define coordinate dimensions ###
         define_spatial_coordinates!(fid, z, r, parallel_io)
 
@@ -577,7 +622,7 @@ setup file i/o for distribution function variables
 """
 function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vperp, vpa,
                        vzeta, vr, vz, composition, collisions, evolve_density,
-                       evolve_upar, evolve_ppar, parallel_io, io_comm)
+                       evolve_upar, evolve_ppar, input_dict, parallel_io, io_comm)
 
     @serial_region begin
         dfns_prefix = string(prefix, ".dfns")
@@ -594,6 +639,9 @@ function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vper
         write_overview!(fid, composition, collisions, parallel_io, evolve_density,
                         evolve_upar, evolve_ppar)
 
+        # write the input settings
+        write_input!(fid, input_dict, parallel_io)
+
         # write the distributions that may be used for boundary conditions to the output
         # file
         write_boundary_distributions!(fid, boundary_distributions, parallel_io,
diff --git a/src/file_io_hdf5.jl b/src/file_io_hdf5.jl
index 2d1d68a47..472662ec2 100644
--- a/src/file_io_hdf5.jl
+++ b/src/file_io_hdf5.jl
@@ -60,6 +60,18 @@ function get_group(file_or_group::HDF5.H5DataStore, name::String)
     end
 end
 
+function is_group(file_or_group::HDF5.H5DataStore, name::String)
+    return isa(file_or_group[name], HDF5.H5DataStore)
+end
+
+function get_subgroup_keys(file_or_group::HDF5.H5DataStore)
+    return collect(k for k ∈ keys(file_or_group) if is_group(file_or_group, k))
+end
+
+function get_variable_keys(file_or_group::HDF5.H5DataStore)
+    return collect(k for k ∈ keys(file_or_group) if !is_group(file_or_group, k))
+end
+
 # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group
 function write_single_value!(file_or_group::HDF5.H5DataStore, name,
                              data::Union{Number, AbstractString, AbstractArray{T,N}},
diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index f0db2c2c7..1869cd833 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -41,6 +41,19 @@ function get_group(file_or_group::NCDataset, name::String)
     end
 end
 
+function is_group(file_or_group::NCDataset, name::String)
+    println("check groups ", NCDatasets.groupnames(file_or_group))
+    return name ∈ NCDatasets.groupnames(file_or_group)
+end
+
+function get_subgroup_keys(file_or_group::NCDataset)
+    return NCDatasets.groupnames(file_or_group)
+end
+
+function get_variable_keys(file_or_group::NCDataset)
+    return keys(file_or_group)
+end
+
 function add_attribute!(file_or_group::NCDataset, name, value)
     file_or_group.attrib[name] = value
 end
diff --git a/src/load_data.jl b/src/load_data.jl
index 6b62e629d..281965b6c 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -15,7 +15,7 @@ export load_rank_data
 export load_species_data
 
 using ..coordinates: define_coordinate
-using ..file_io: get_group
+using ..file_io: get_group, get_subgroup_keys, get_variable_keys
 using ..input_structs: advection_input, grid_input
 using ..looping
 
@@ -136,6 +136,29 @@ function load_slice(file_or_group::NCDataset, name::String, slices_or_indices...
     end
 end
 
+"""
+Load saved input settings
+"""
+function load_input(fid)
+    function read_dict(io, section_name)
+        # Function that can be called recursively to read nested Dicts from sub-groups in
+        # the output file
+        section_io = get_group(io, section_name)
+        section = Dict{String,Any}()
+
+        for key ∈ get_variable_keys(section_io)
+            section[key] = load_variable(section_io, key)
+        end
+        for key ∈ get_subgroup_keys(section_io)
+            section[key] = read_dict(section_io, key)
+        end
+
+        return section
+    end
+
+    return read_dict(fid, "input")
+end
+
 """
 Load data for a coordinate
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 211479f76..c321d51c7 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -386,7 +386,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
     # setup i/o
     ascii_io, io_moments, io_dfns = setup_file_io(io_input, boundary_distributions, vz,
         vr, vzeta, vpa, vperp, z, r, composition, collisions, moments.evolve_density,
-        moments.evolve_upar, moments.evolve_ppar)
+        moments.evolve_upar, moments.evolve_ppar, input_dict)
     # write initial data to ascii files
     write_data_to_ascii(moments, fields, vpa, vperp, z, r, code_time, composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 953907d69..e5c7fa267 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -35,7 +35,7 @@ using ..coordinates: define_coordinate
 using ..file_io: open_ascii_output_file
 using ..type_definitions: mk_float, mk_int
 using ..initial_conditions: vpagrid_to_dzdt
-using ..load_data: open_readonly_output_file, get_group, load_time_data
+using ..load_data: open_readonly_output_file, get_group, load_input, load_time_data
 using ..load_data: get_nranks
 using ..load_data: load_fields_data, load_pdf_data
 using ..load_data: load_charged_particle_moments_data, load_neutral_particle_moments_data
@@ -380,6 +380,9 @@ function analyze_and_plot_data(prefix...)
     # load block data on iblock=0
     nblocks, iblock = get_tuple_of_return_values(load_block_data, moments_files0)
 
+    # load input used for the run(s)
+    scan_input = get_tuple_of_return_values(load_input, moments_files0)
+
     # load global and local sizes of grids stored on each output file
     # z z_wgts r r_wgts may take different values on different blocks
     # we need to construct the global grid below
@@ -579,9 +582,9 @@ function analyze_and_plot_data(prefix...)
         close(f)
     end
 
-    #geometry, composition =
-    #    get_tuple_of_return_values(get_geometry_and_composition, scan_input,
-    #                               n_ion_species, n_neutral_species)
+    geometry, composition =
+        get_tuple_of_return_values(get_geometry_and_composition, scan_input,
+                                   n_ion_species, n_neutral_species)
 
     # initialise the post-processing input options
     nwrite_movie, itime_min, itime_max, nwrite_movie_pdfs, itime_min_pdfs, itime_max_pdfs,
@@ -691,6 +694,7 @@ function analyze_and_plot_data(prefix...)
     end
 
     # For now, don't support multi-run comparison in remaining 2D and MMS diagnostics
+    scan_input = scan_input[1]
     density = density[1]
     parallel_flow = parallel_flow[1]
     parallel_pressure = parallel_pressure[1]
@@ -757,8 +761,6 @@ function analyze_and_plot_data(prefix...)
             Lr_in = 1.0
         end
 
-        #geometry, composition = get_geometry_and_composition(scan_input,n_ion_species,n_neutral_species)
-
         manufactured_solns_list = manufactured_solutions(Lr_in,z_global.L,r_bc,z_bc,geometry,composition,r_global.n)
         dfni_func = manufactured_solns_list.dfni_func
         densi_func = manufactured_solns_list.densi_func

From f33ee403a867d09dbc72a1fba035b22e897ef8db Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 May 2023 21:34:11 +0100
Subject: [PATCH 028/428] Convert \pi to mk_float in test inputs

Julia's \pi on its own is not a type that HDF5 can write, so explicitly
convert to mk_float.
---
 test/nonlinear_sound_wave_tests.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/nonlinear_sound_wave_tests.jl b/test/nonlinear_sound_wave_tests.jl
index ca0ff8105..2742760cd 100644
--- a/test/nonlinear_sound_wave_tests.jl
+++ b/test/nonlinear_sound_wave_tests.jl
@@ -143,10 +143,10 @@ test_input_finite_difference = Dict("n_ion_species" => 1,
                                     "z_IC_upar_amplitude1" => 0.0,
                                     "z_IC_upar_phase1" => 0.0,
                                     "z_IC_temperature_amplitude1" => 0.5,
-                                    "z_IC_temperature_phase1" => π,
+                                    "z_IC_temperature_phase1" => mk_float(π),
                                     "z_IC_option2" => "sinusoid",
                                     "z_IC_density_amplitude2" => 0.5,
-                                    "z_IC_density_phase2" => π,
+                                    "z_IC_density_phase2" => mk_float(π),
                                     "z_IC_upar_amplitude2" => 0.0,
                                     "z_IC_upar_phase2" => 0.0,
                                     "z_IC_temperature_amplitude2" => 0.5,

From 5a6ad8b06d24ab03d0e7ea3f4c2134cbf324aef6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 May 2023 22:37:58 +0100
Subject: [PATCH 029/428] Do not write variable if number of species is zero

If n_ion_species (n_neutral_species) is zero, then the arry for an ion
(neutral) variable has no data, so does not need to be written and might
cause an error due to having a size-zero dimension.
---
 src/file_io_hdf5.jl   | 38 ++++++++++++++++++++++++++++++--------
 src/file_io_netcdf.jl | 38 ++++++++++++++++++++++++++++++--------
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/src/file_io_hdf5.jl b/src/file_io_hdf5.jl
index 472662ec2..19a153bb2 100644
--- a/src/file_io_hdf5.jl
+++ b/src/file_io_hdf5.jl
@@ -89,8 +89,20 @@ function write_single_value!(file_or_group::HDF5.H5DataStore, name,
     end
 
     if n_ion_species !== nothing
+        if n_ion_species < 0
+            error("n_ion_species must be non-negative, got $n_ion_species")
+        elseif n_ion_species == 0
+            # No data to write
+            return nothing
+        end
         coords = tuple(coords..., n_ion_species)
     elseif n_neutral_species !== nothing
+        if n_neutral_species < 0
+            error("n_neutral_species must be non-negative, got $n_neutral_species")
+        elseif n_neutral_species == 0
+            # No data to write
+            return nothing
+        end
         coords = tuple(coords..., n_neutral_species)
     end
     dim_sizes, chunk_sizes = hdf5_get_fixed_dim_sizes(coords, parallel_io)
@@ -182,21 +194,31 @@ end
 
 function create_dynamic_variable!(file_or_group::HDF5.H5DataStore, name, type,
                                   coords::coordinate...; parallel_io,
-                                  n_ion_species=0, n_neutral_species=0,
+                                  n_ion_species=nothing, n_neutral_species=nothing,
                                   description=nothing, units=nothing)
 
-    if n_ion_species != 0 && n_neutral_species != 0
+    if n_ion_species !== nothing && n_neutral_species !== nothing
         error("Variable should not contain both ion and neutral species dimensions. "
               * "Got n_ion_species=$n_ion_species and "
               * "n_neutral_species=$n_neutral_species")
     end
-    n_ion_species < 0 && error("n_ion_species must be non-negative, got $n_ion_species")
-    n_neutral_species < 0 && error("n_neutral_species must be non-negative, got $n_neutral_species")
 
     # Add the number of species to the spatial/velocity-space coordinates
-    if n_ion_species > 0
+    if n_ion_species !== nothing
+        if n_ion_species < 0
+            error("n_ion_species must be non-negative, got $n_ion_species")
+        elseif n_ion_species == 0
+            # No data to write
+            return nothing
+        end
         fixed_coords = tuple(coords..., n_ion_species)
-    elseif n_neutral_species > 0
+    elseif n_neutral_species !== nothing
+        if n_neutral_species < 0
+            error("n_neutral_species must be non-negative, got $n_neutral_species")
+        elseif n_neutral_species == 0
+            # No data to write
+            return nothing
+        end
         fixed_coords = tuple(coords..., n_neutral_species)
     else
         fixed_coords = coords
@@ -208,9 +230,9 @@ function create_dynamic_variable!(file_or_group::HDF5.H5DataStore, name, type,
 
     # Add attribute listing the dimensions belonging to this variable
     dim_names = Tuple(c.name for c ∈ coords)
-    if n_ion_species > 0
+    if n_ion_species !== nothing
         dim_names = tuple(dim_names..., "ion_species")
-    elseif n_neutral_species > 0
+    elseif n_neutral_species !== nothing
         dim_names = tuple(dim_names..., "neutral_species")
     end
     add_attribute!(var, "dims", join(dim_names, ","))
diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 1869cd833..7cede9af9 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -98,9 +98,21 @@ function write_single_value!(file_or_group::NCDataset, name,
         dims = Tuple(c.name for c in coords)
 
         if n_ion_species !== nothing
+            if n_ion_species < 0
+                error("n_ion_species must be non-negative, got $n_ion_species")
+            elseif n_ion_species == 0
+                # No data to write
+                return nothing
+            end
             maybe_create_netcdf_dim(file_or_group, "ion_species", n_ion_species)
             dims = tuple(dims..., "ion_species")
         elseif n_neutral_species !== nothing
+            if n_neutral_species < 0
+                error("n_neutral_species must be non-negative, got $n_neutral_species")
+            elseif n_neutral_species == 0
+                # No data to write
+                return nothing
+            end
             maybe_create_netcdf_dim(file_or_group, "neutral_species", n_neutral_species)
             dims = tuple(dims..., "neutral_species")
         end
@@ -120,25 +132,35 @@ end
 
 function create_dynamic_variable!(file_or_group::NCDataset, name, type,
                                   coords::coordinate...; parallel_io,
-                                  n_ion_species=0, n_neutral_species=0,
+                                  n_ion_species=nothing, n_neutral_species=nothing,
                                   description=nothing, units=nothing)
 
-    if n_ion_species != 0 && n_neutral_species != 0
+    if n_ion_species !== nothing && n_neutral_species !== nothing
         error("Variable should not contain both ion and neutral species dimensions. "
               * "Got n_ion_species=$n_ion_species and "
               * "n_neutral_species=$n_neutral_species")
     end
-    n_ion_species < 0 && error("n_ion_species must be non-negative, got $n_ion_species")
-    n_neutral_species < 0 && error("n_neutral_species must be non-negative, got $n_neutral_species")
 
     # Create time dimension if necessary
     maybe_create_netcdf_dim(file_or_group, "time", Inf)
 
     # Create species dimension if necessary
-    if n_ion_species > 0
+    if n_ion_species !== nothing
+        if n_ion_species < 0
+            error("n_ion_species must be non-negative, got $n_ion_species")
+        elseif n_ion_species == 0
+            # No data to write
+            return nothing
+        end
         maybe_create_netcdf_dim(file_or_group, "ion_species", n_ion_species)
     end
-    if n_neutral_species > 0
+    if n_neutral_species !== nothing
+        if n_neutral_species < 0
+            error("n_neutral_species must be non-negative, got $n_neutral_species")
+        elseif n_neutral_species == 0
+            # No data to write
+            return nothing
+        end
         maybe_create_netcdf_dim(file_or_group, "neutral_species", n_neutral_species)
     end
 
@@ -150,9 +172,9 @@ function create_dynamic_variable!(file_or_group::NCDataset, name, type,
     # create the variable so it can be expanded indefinitely (up to the largest unsigned
     # integer in size) in the time dimension
     coord_dims = Tuple(c.name for c ∈ coords)
-    if n_ion_species > 0
+    if n_ion_species !== nothing
         fixed_dims = tuple(coord_dims..., "ion_species")
-    elseif n_neutral_species > 0
+    elseif n_neutral_species !== nothing
         fixed_dims = tuple(coord_dims..., "neutral_species")
     else
         fixed_dims = coord_dims

From ea8407b380b079708e55f7104295dc0002817d07 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Jun 2023 09:17:27 +0100
Subject: [PATCH 030/428] Fix NetCDF output when value happens to be a
 SubString{String}

Passing a variable with type SubString{String} to NCDatasets causes an
error, so convert any AbstractString to a String before writing.
---
 src/file_io_netcdf.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/file_io_netcdf.jl b/src/file_io_netcdf.jl
index 7cede9af9..8553e8f1f 100644
--- a/src/file_io_netcdf.jl
+++ b/src/file_io_netcdf.jl
@@ -86,8 +86,15 @@ function write_single_value!(file_or_group::NCDataset, name,
               "Got n_ion_species=$n_ion_species, n_neutral_species=$n_neutral_species")
     end
 
-    if isa(value, Number) || isa(value, String)
+    if isa(value, Number) || isa(value, AbstractString)
         coords !== () && error("cannot pass coordinates with a scalar")
+
+        if isa(value, AbstractString)
+            # Trying to write a SubString{String} causes an error, so force anything
+            # string-like to be a String
+            value = String(value)
+        end
+
         type = typeof(value)
         dims = ()
     else

From c86346df0ddef246918f2020fc07be23f7df53f8 Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Wed, 7 Jun 2023 11:01:00 +0100
Subject: [PATCH 031/428] replaced 'charged' with 'ion' in numerous structs and
 other variables throughout the code in preparation for adding a separate
 'electron' species that my be treated quite differently than the ions.

---
 src/charge_exchange.jl             |  20 ++---
 src/continuity.jl                  |   8 +-
 src/derivatives.jl                 |  16 ++--
 src/energy_equation.jl             |   8 +-
 src/file_io.jl                     |  74 ++++++++--------
 src/force_balance.jl               |   8 +-
 src/initial_conditions.jl          | 123 ++++++++++++-------------
 src/ionization.jl                  |   4 +-
 src/load_data.jl                   |  76 ++++++++--------
 src/moment_kinetics.jl             |   2 +-
 src/moment_kinetics_input.jl       |  90 +++++++++----------
 src/moment_kinetics_structs.jl     |   6 +-
 src/plot_MMS_sequence.jl           |  12 +--
 src/plot_sequence.jl               |   2 +-
 src/post_processing.jl             |  46 +++++-----
 src/r_advection.jl                 |   6 +-
 src/source_terms.jl                |  10 +--
 src/time_advance.jl                | 138 ++++++++++++++---------------
 src/velocity_moments.jl            |  86 +++++++++---------
 src/vpa_advection.jl               |  20 ++---
 src/z_advection.jl                 |   6 +-
 test/nonlinear_sound_wave_tests.jl |  98 ++++++++++----------
 22 files changed, 430 insertions(+), 429 deletions(-)

diff --git a/src/charge_exchange.jl b/src/charge_exchange.jl
index 1e4b5a267..66c1bb7fa 100644
--- a/src/charge_exchange.jl
+++ b/src/charge_exchange.jl
@@ -28,7 +28,7 @@ function charge_exchange_collisions_1V!(f_out, f_neutral_out, fvec_in, moments,
                 f_out[:,1,:,:,is], fvec_in.pdf[:,1,:,:,is],
                 fvec_in.pdf_neutral[:,1,1,:,:,is],
                 fvec_in.density_neutral[:,:,is], fvec_in.upar[:,:,is],
-                fvec_in.uz_neutral[:,:,is], moments.charged.vth[:,:,is],
+                fvec_in.uz_neutral[:,:,is], moments.ion.vth[:,:,is],
                 moments.neutral.vth[:,:,is], moments, vpa, vz, charge_exchange_frequency,
                 vz_spectral, dt)
         end
@@ -42,7 +42,7 @@ function charge_exchange_collisions_1V!(f_out, f_neutral_out, fvec_in, moments,
                 f_neutral_out[:,1,1,:,:,isn], fvec_in.pdf_neutral[:,1,1,:,:,isn],
                 fvec_in.pdf[:,1,:,:,isn], fvec_in.density[:,:,isn],
                 fvec_in.uz_neutral[:,:,isn], fvec_in.upar[:,:,isn],
-                moments.neutral.vth[:,:,isn], moments.charged.vth[:,:,isn], moments,
+                moments.neutral.vth[:,:,isn], moments.ion.vth[:,:,isn], moments,
                 vz, vpa, charge_exchange_frequency, vpa_spectral, dt)
         end
     else
@@ -135,7 +135,7 @@ function charge_exchange_collisions_single_species!(f_out, pdf_in, pdf_other,
     end
 end
 
-function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in, f_charged_vrvzvzeta_in, fvec_in, composition, vz, vr, vzeta, vpa, vperp, z, r,
+function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in, f_ion_vrvzvzeta_in, fvec_in, composition, vz, vr, vzeta, vpa, vperp, z, r,
                                      charge_exchange_frequency, dt)
     # This routine assumes a 3V model with:
     @boundscheck vz.n == size(f_neutral_out,1) || throw(BoundsError(f_neutral_out))
@@ -144,12 +144,12 @@ function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in,
     @boundscheck z.n == size(f_neutral_out,4) || throw(BoundsError(f_neutral_out))
     @boundscheck r.n == size(f_neutral_out,5) || throw(BoundsError(f_neutral_out))
     @boundscheck composition.n_neutral_species == size(f_neutral_out,6) || throw(BoundsError(f_neutral_out))
-    @boundscheck vz.n == size(f_charged_vrvzvzeta_in,1) || throw(BoundsError(f_charged_vrvzvzeta_in))
-    @boundscheck vr.n == size(f_charged_vrvzvzeta_in,2) || throw(BoundsError(f_charged_vrvzvzeta_in))
-    @boundscheck vzeta.n == size(f_charged_vrvzvzeta_in,3) || throw(BoundsError(f_charged_vrvzvzeta_in))
-    @boundscheck z.n == size(f_charged_vrvzvzeta_in,4) || throw(BoundsError(f_charged_vrvzvzeta_in))
-    @boundscheck r.n == size(f_charged_vrvzvzeta_in,5) || throw(BoundsError(f_charged_vrvzvzeta_in))
-    @boundscheck composition.n_neutral_species == size(f_charged_vrvzvzeta_in,6) || throw(BoundsError(f_charged_vrvzvzeta_in))
+    @boundscheck vz.n == size(f_ion_vrvzvzeta_in,1) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck vr.n == size(f_ion_vrvzvzeta_in,2) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck vzeta.n == size(f_ion_vrvzvzeta_in,3) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck z.n == size(f_ion_vrvzvzeta_in,4) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck r.n == size(f_ion_vrvzvzeta_in,5) || throw(BoundsError(f_ion_vrvzvzeta_in))
+    @boundscheck composition.n_neutral_species == size(f_ion_vrvzvzeta_in,6) || throw(BoundsError(f_ion_vrvzvzeta_in))
     @boundscheck vpa.n == size(f_out,1) || throw(BoundsError(f_out))
     @boundscheck vperp.n == size(f_out,2) || throw(BoundsError(f_out))
     @boundscheck z.n == size(f_out,3) || throw(BoundsError(f_out))
@@ -181,7 +181,7 @@ function charge_exchange_collisions_3V!(f_out, f_neutral_out, f_neutral_gav_in,
         for is ∈ 1:composition.n_ion_species
             f_neutral_out[ivz,ivr,ivzeta,iz,ir,isn] +=
                 dt*charge_exchange_frequency*(
-                    f_charged_vrvzvzeta_in[ivz,ivr,ivzeta,iz,ir,is]*fvec_in.density_neutral[iz,ir,isn]
+                    f_ion_vrvzvzeta_in[ivz,ivr,ivzeta,iz,ir,is]*fvec_in.density_neutral[iz,ir,isn]
                     - fvec_in.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn]*fvec_in.density[iz,ir,is])
         end
     end
diff --git a/src/continuity.jl b/src/continuity.jl
index a1bcb9180..ca7ff95cd 100644
--- a/src/continuity.jl
+++ b/src/continuity.jl
@@ -8,7 +8,7 @@ using ..calculus: derivative!
 using ..looping
 
 """
-use the continuity equation dn/dt + d(n*upar)/dz to update the density n for all charged
+use the continuity equation dn/dt + d(n*upar)/dz to update the density n for all ion
 species
 """
 function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spectral,
@@ -18,8 +18,8 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     @loop_s_r_z is ir iz begin
         # Use ddens_dz is upwinded using upar
         dens_out[iz,ir,is] -=
-            dt*(fvec_in.upar[iz,ir,is]*moments.charged.ddens_dz_upwind[iz,ir,is] +
-                fvec_in.density[iz,ir,is]*moments.charged.dupar_dz[iz,ir,is])
+            dt*(fvec_in.upar[iz,ir,is]*moments.ion.ddens_dz_upwind[iz,ir,is] +
+                fvec_in.density[iz,ir,is]*moments.ion.dupar_dz[iz,ir,is])
     end
 
     # update the density to account for ionization collisions;
@@ -34,7 +34,7 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
-            dens_out[iz,ir,is] += dt*diffusion_coefficient*moments.charged.d2dens_dz[iz,ir,is]
+            dens_out[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2dens_dz[iz,ir,is]
         end
     end
 end
diff --git a/src/derivatives.jl b/src/derivatives.jl
index d7f840442..8ef4bc21a 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -19,7 +19,7 @@ using ..looping
 Centered derivatives
 df/dr group of rountines for
 fields & moments -> [z,r]
-dfns (charged) -> [vpa,vperp,z,r,s]
+dfns (ion) -> [vpa,vperp,z,r,s]
 dfns (neutrals) -> [vz,vr,vzeta,z,r,sn]
 """
 
@@ -49,7 +49,7 @@ function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
 end
 
 #df/dr
-#5D version for f[vpa,vperp,z,r,s] -> charged particle dfn (species indexing taken outside this loop)
+#5D version for f[vpa,vperp,z,r,s] -> ion particle dfn (species indexing taken outside this loop)
 function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5},
         dfdr_lower_endpoints::AbstractArray{mk_float,4},
         dfdr_upper_endpoints::AbstractArray{mk_float,4},
@@ -101,7 +101,7 @@ end
 Centered derivatives
 df/dz group of rountines for
 fields & moments -> [z,r]
-dfns (charged) -> [vpa,vperp,z,r,s]
+dfns (ion) -> [vpa,vperp,z,r,s]
 dfns (neutrals) -> [vz,vr,vzeta,z,r,sn]
 """
 
@@ -130,7 +130,7 @@ function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
 	
 end
 
-#5D version for f[vpa,vperp,z,r,s] -> dfn charged particles
+#5D version for f[vpa,vperp,z,r,s] -> dfn ion particles
 function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5},
         dfdz_lower_endpoints::AbstractArray{mk_float,4},
         dfdz_upper_endpoints::AbstractArray{mk_float,4},
@@ -182,7 +182,7 @@ end
 Upwind derivatives
 df/dr group of rountines for
 fields & moments -> [z,r]
-dfns (charged) -> [vpa,vperp,z,r,s]
+dfns (ion) -> [vpa,vperp,z,r,s]
 dfns (neutrals) -> [vz,vr,vzeta,z,r,sn]
 """
 
@@ -215,7 +215,7 @@ function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
 end
 
 #df/dr
-#5D version for f[vpa,vperp,z,r,s] -> charged particle dfn (species indexing taken outside this loop)
+#5D version for f[vpa,vperp,z,r,s] -> ion particle dfn (species indexing taken outside this loop)
 function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5},
         advect, adv_fac_lower_buffer::AbstractArray{mk_float,4},
         adv_fac_upper_buffer::AbstractArray{mk_float,4},
@@ -278,7 +278,7 @@ end
 Upwind derivatives
 df/dz group of rountines for
 fields & moments -> [z,r]
-dfns (charged) -> [vpa,vperp,z,r,s]
+dfns (ion) -> [vpa,vperp,z,r,s]
 dfns (neutrals) -> [vz,vr,vzeta,z,r,sn]
 """
 
@@ -309,7 +309,7 @@ function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
     end
 end
 
-#5D version for f[vpa,vperp,z,r,s] -> dfn charged particles
+#5D version for f[vpa,vperp,z,r,s] -> dfn ion particles
 function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_float,5},
         advect, adv_fac_lower_buffer::AbstractArray{mk_float,4},
         adv_fac_upper_buffer::AbstractArray{mk_float,4},
diff --git a/src/energy_equation.jl b/src/energy_equation.jl
index 2d8fb578a..7b795ebef 100644
--- a/src/energy_equation.jl
+++ b/src/energy_equation.jl
@@ -17,15 +17,15 @@ function energy_equation!(ppar, fvec, moments, collisions, dt, spectral, composi
     begin_s_r_z_region()
 
     @loop_s_r_z is ir iz begin
-        ppar[iz,ir,is] += dt*(-fvec.upar[iz,ir,is]*moments.charged.dppar_dz_upwind[iz,ir,is]
-                              - moments.charged.dqpar_dz[iz,ir,is]
-                              - 3.0*fvec.ppar[iz,ir,is]*moments.charged.dupar_dz[iz,ir,is])
+        ppar[iz,ir,is] += dt*(-fvec.upar[iz,ir,is]*moments.ion.dppar_dz_upwind[iz,ir,is]
+                              - moments.ion.dqpar_dz[iz,ir,is]
+                              - 3.0*fvec.ppar[iz,ir,is]*moments.ion.dupar_dz[iz,ir,is])
     end
 
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
-            ppar[iz,ir,is] += dt*diffusion_coefficient*moments.charged.d2ppar_dz2[iz,ir,is]
+            ppar[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2ppar_dz2[iz,ir,is]
         end
     end
 
diff --git a/src/file_io.jl b/src/file_io.jl
index d52f00bb0..10b212f90 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -26,7 +26,7 @@ struct ascii_ios{T <: Union{IOStream,Nothing}}
     #ff::T
     # corresponds to the ascii file to which velocity space moments of the
     # distribution function such as density and pressure are written
-    moments_charged::T
+    moments_ion::T
     moments_neutral::T
     # corresponds to the ascii file to which electromagnetic fields
     # such as the electrostatic potential are written
@@ -48,15 +48,15 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn}
     Er::Tphi
     # handle for the z electric field variable
     Ez::Tphi
-    # handle for the charged species density
+    # handle for the ion species density
     density::Tmomi
-    # handle for the charged species parallel flow
+    # handle for the ion species parallel flow
     parallel_flow::Tmomi
-    # handle for the charged species parallel pressure
+    # handle for the ion species parallel pressure
     parallel_pressure::Tmomi
-    # handle for the charged species parallel heat flux
+    # handle for the ion species parallel heat flux
     parallel_heat_flux::Tmomi
-    # handle for the charged species thermal speed
+    # handle for the ion species thermal speed
     thermal_speed::Tmomi
 
     # handle for the neutral species density
@@ -77,7 +77,7 @@ distribution function data only
 struct io_dfns_info{Tfile, Tfi, Tfn, Tmoments}
     # file identifier for the binary file to which data is written
     fid::Tfile
-    # handle for the charged species distribution function variable
+    # handle for the ion species distribution function variable
     f::Tfi
     # handle for the neutral species distribution function variable
     f_neutral::Tfn
@@ -115,10 +115,10 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
 
         if io_input.ascii_output
             #ff_io = open_ascii_output_file(out_prefix, "f_vs_t")
-            mom_chrg_io = open_ascii_output_file(out_prefix, "moments_charged_vs_t")
+            mom_ion_io = open_ascii_output_file(out_prefix, "moments_ion_vs_t")
             mom_ntrl_io = open_ascii_output_file(out_prefix, "moments_neutral_vs_t")
             fields_io = open_ascii_output_file(out_prefix, "fields_vs_t")
-            ascii = ascii_ios(mom_chrg_io, mom_ntrl_io, fields_io)
+            ascii = ascii_ios(mom_ion_io, mom_ntrl_io, fields_io)
         else
             ascii = ascii_ios(nothing, nothing, nothing)
         end
@@ -238,14 +238,14 @@ function write_boundary_distributions!(fid, boundary_distributions, parallel_io,
     @serial_region begin
         boundary_distributions_io = create_io_group(fid, "boundary_distributions")
 
-        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_left",
-            boundary_distributions.pdf_rboundary_charged[:,:,:,1,:], vpa, vperp, z,
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_ion_left",
+            boundary_distributions.pdf_rboundary_ion[:,:,:,1,:], vpa, vperp, z,
             parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
-            description="Initial charged-particle pdf at left radial boundary")
-        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_right",
-            boundary_distributions.pdf_rboundary_charged[:,:,:,2,:], vpa, vperp, z,
+            description="Initial ion-particle pdf at left radial boundary")
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_ion_right",
+            boundary_distributions.pdf_rboundary_ion[:,:,:,2,:], vpa, vperp, z,
             parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
-            description="Initial charged-particle pdf at right radial boundary")
+            description="Initial ion-particle pdf at right radial boundary")
         write_single_value!(boundary_distributions_io, "pdf_rboundary_neutral_left",
             boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:], vz, vr, vzeta, z,
             parallel_io=parallel_io, n_neutral_species=composition.n_neutral_species,
@@ -452,35 +452,35 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_density = create_dynamic_variable!(dynamic, "density", mk_float, z, r;
                                               n_ion_species=n_ion_species,
                                               parallel_io=parallel_io,
-                                              description="charged species density",
+                                              description="ion species density",
                                               units="n_ref")
 
         # io_upar is the handle for the ion parallel flow density
         io_upar = create_dynamic_variable!(dynamic, "parallel_flow", mk_float, z, r;
                                            n_ion_species=n_ion_species,
                                            parallel_io=parallel_io,
-                                           description="charged species parallel flow",
+                                           description="ion species parallel flow",
                                            units="c_ref = sqrt(2*T_ref/mi)")
 
         # io_ppar is the handle for the ion parallel pressure
         io_ppar = create_dynamic_variable!(dynamic, "parallel_pressure", mk_float, z, r;
                                            n_ion_species=n_ion_species,
                                            parallel_io=parallel_io,
-                                           description="charged species parallel pressure",
+                                           description="ion species parallel pressure",
                                            units="n_ref*T_ref")
 
         # io_qpar is the handle for the ion parallel heat flux
         io_qpar = create_dynamic_variable!(dynamic, "parallel_heat_flux", mk_float, z, r;
                                            n_ion_species=n_ion_species,
                                            parallel_io=parallel_io,
-                                           description="charged species parallel heat flux",
+                                           description="ion species parallel heat flux",
                                            units="n_ref*T_ref*c_ref")
 
         # io_vth is the handle for the ion thermal speed
         io_vth = create_dynamic_variable!(dynamic, "thermal_speed", mk_float, z, r;
                                           n_ion_species=n_ion_species,
                                           parallel_io=parallel_io,
-                                          description="charged species thermal speed",
+                                          description="ion species thermal speed",
                                           units="c_ref")
 
         # io_density_neutral is the handle for the neutral particle density
@@ -546,7 +546,7 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz,
         io_f = create_dynamic_variable!(dynamic, "f", mk_float, vpa, vperp, z, r;
                                         n_ion_species=n_ion_species,
                                         parallel_io=parallel_io,
-                                        description="charged species distribution function")
+                                        description="ion species distribution function")
 
         # io_f_neutral is the handle for the neutral pdf
         io_f_neutral = create_dynamic_variable!(dynamic, "f_neutral", mk_float, vz, vr, vzeta, z, r;
@@ -691,15 +691,15 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
         append_to_dynamic_var(io_moments.Ez, fields.Ez, t_idx, z, r)
 
         # add the density data at this time slice to the output file
-        append_to_dynamic_var(io_moments.density, moments.charged.dens, t_idx, z, r,
+        append_to_dynamic_var(io_moments.density, moments.ion.dens, t_idx, z, r,
                               n_ion_species)
-        append_to_dynamic_var(io_moments.parallel_flow, moments.charged.upar, t_idx, z, r,
+        append_to_dynamic_var(io_moments.parallel_flow, moments.ion.upar, t_idx, z, r,
                               n_ion_species)
-        append_to_dynamic_var(io_moments.parallel_pressure, moments.charged.ppar, t_idx,
+        append_to_dynamic_var(io_moments.parallel_pressure, moments.ion.ppar, t_idx,
                               z, r, n_ion_species)
-        append_to_dynamic_var(io_moments.parallel_heat_flux, moments.charged.qpar, t_idx,
+        append_to_dynamic_var(io_moments.parallel_heat_flux, moments.ion.qpar, t_idx,
                               z, r, n_ion_species)
-        append_to_dynamic_var(io_moments.thermal_speed, moments.charged.vth, t_idx, z, r,
+        append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth, t_idx, z, r,
                               n_ion_species)
         if n_neutral_species > 0
             append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
@@ -758,15 +758,15 @@ end
             append_to_dynamic_var(io_moments.Ez, fields.Ez.data, t_idx, z, r)
 
             # add the density data at this time slice to the output file
-            append_to_dynamic_var(io_moments.density, moments.charged.dens.data, t_idx, z,
+            append_to_dynamic_var(io_moments.density, moments.ion.dens.data, t_idx, z,
                                   r, n_ion_species)
-            append_to_dynamic_var(io_moments.parallel_flow, moments.charged.upar.data,
+            append_to_dynamic_var(io_moments.parallel_flow, moments.ion.upar.data,
                                   t_idx, z, r, n_ion_species)
-            append_to_dynamic_var(io_moments.parallel_pressure, moments.charged.ppar.data,
+            append_to_dynamic_var(io_moments.parallel_pressure, moments.ion.ppar.data,
                                   t_idx, z, r, n_ion_species)
             append_to_dynamic_var(io_moments.parallel_heat_flux,
-                                  moments.charged.qpar.data, t_idx, z, r, n_ion_species)
-            append_to_dynamic_var(io_moments.thermal_speed, moments.charged.vth.data,
+                                  moments.ion.qpar.data, t_idx, z, r, n_ion_species)
+            append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth.data,
                                   t_idx, z, r, n_ion_species)
             if n_neutral_species > 0
                 append_to_dynamic_var(io_moments.density_neutral,
@@ -849,7 +849,7 @@ include("file_io_hdf5.jl")
 """
 function write_data_to_ascii(moments, fields, vpa, vperp, z, r, t, n_ion_species,
                              n_neutral_species, ascii_io::Union{ascii_ios,Nothing})
-    if ascii_io === nothing || ascii_io.moments_charged === nothing
+    if ascii_io === nothing || ascii_io.moments_ion === nothing
         # ascii I/O is disabled
         return nothing
     end
@@ -858,7 +858,7 @@ function write_data_to_ascii(moments, fields, vpa, vperp, z, r, t, n_ion_species
         # Only read/write from first process in each 'block'
 
         #write_f_ascii(ff, z, vpa, t, ascii_io.ff)
-        write_moments_charged_ascii(moments.charged, z, r, t, n_ion_species, ascii_io.moments_charged)
+        write_moments_ion_ascii(moments.ion, z, r, t, n_ion_species, ascii_io.moments_ion)
         if n_neutral_species > 0
             write_moments_neutral_ascii(moments.neutral, z, r, t, n_neutral_species, ascii_io.moments_neutral)
         end
@@ -893,9 +893,9 @@ function write_f_ascii(f, z, vpa, t, ascii_io)
 end
 
 """
-write moments of the charged species distribution function f at this time slice
+write moments of the ion species distribution function f at this time slice
 """
-function write_moments_charged_ascii(mom, z, r, t, n_species, ascii_io)
+function write_moments_ion_ascii(mom, z, r, t, n_species, ascii_io)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
@@ -1092,9 +1092,9 @@ function debug_dump(vz::coordinate, vr::coordinate, vzeta::coordinate, vpa::coor
         debug_output_file.label[debug_output_counter[]] = label
         # add the distribution function data at this time slice to the netcdf file
         if ff === nothing
-            debug_output_file.dfns.charged_f[:,:,:,:,:,debug_output_counter[]] = 0.0
+            debug_output_file.dfns.ion_f[:,:,:,:,:,debug_output_counter[]] = 0.0
         else
-            debug_output_file.dfns.charged_f[:,:,:,:,:,debug_output_counter[]] = ff
+            debug_output_file.dfns.ion_f[:,:,:,:,:,debug_output_counter[]] = ff
         end
         # add the moments data at this time slice to the netcdf file
         if dens === nothing
diff --git a/src/force_balance.jl b/src/force_balance.jl
index f0fcb9227..6044431f0 100644
--- a/src/force_balance.jl
+++ b/src/force_balance.jl
@@ -22,9 +22,9 @@ function force_balance!(pflx, density_out, fvec, moments, fields, collisions, dt
     upar = fvec.upar
     @loop_s_r_z is ir iz begin
         pflx[iz,ir,is] = density[iz,ir,is]*upar[iz,ir,is] -
-                         dt*(moments.charged.dppar_dz[iz,ir,is] +
-                             upar[iz,ir,is]*upar[iz,ir,is]*moments.charged.ddens_dz_upwind[iz,ir,is] +
-                             2.0*density[iz,ir,is]*upar[iz,ir,is]*moments.charged.dupar_dz_upwind[iz,ir,is] -
+                         dt*(moments.ion.dppar_dz[iz,ir,is] +
+                             upar[iz,ir,is]*upar[iz,ir,is]*moments.ion.ddens_dz_upwind[iz,ir,is] +
+                             2.0*density[iz,ir,is]*upar[iz,ir,is]*moments.ion.dupar_dz_upwind[iz,ir,is] -
                              0.5*geometry.bzed*fields.Ez[iz,ir]*density[iz,ir,is])
     end
 
@@ -32,7 +32,7 @@ function force_balance!(pflx, density_out, fvec, moments, fields, collisions, dt
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
-            pflx[iz,ir,is] += dt*diffusion_coefficient*moments.charged.d2upar_dz2[iz,ir,is]*density[iz,ir,is]
+            pflx[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2upar_dz2[iz,ir,is]*density[iz,ir,is]
         end
     end
 
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 171a41ba7..bf1145ba5 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -27,8 +27,8 @@ using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
-using ..velocity_moments: create_moments_charged, create_moments_neutral, update_qpar!
-using ..velocity_moments: moments_charged_substruct, moments_neutral_substruct
+using ..velocity_moments: create_moments_ion, create_moments_neutral, update_qpar!
+using ..velocity_moments: moments_ion_substruct, moments_neutral_substruct
 using ..velocity_moments: update_neutral_density!, update_neutral_pz!, update_neutral_pr!, update_neutral_pzeta!
 using ..velocity_moments: update_neutral_uz!, update_neutral_ur!, update_neutral_uzeta!, update_neutral_qz!
 using ..velocity_moments: update_ppar!, update_upar!, update_density!
@@ -46,14 +46,14 @@ end
 
 # struct of structs neatly contains i+n info?
 struct pdf_struct
-    #charged particles: s + r + z + vperp + vpa
-    charged::pdf_substruct{5}
+    #ion particles: s + r + z + vperp + vpa
+    ion::pdf_substruct{5}
     #neutral particles: s + r + z + vzeta + vr + vz
     neutral::pdf_substruct{6}
 end
 
 struct moments_struct
-    charged::moments_charged_substruct
+    ion::moments_ion_substruct
     neutral::moments_neutral_substruct
     # flag that indicates if the density should be evolved via continuity equation
     evolve_density::Bool
@@ -72,8 +72,8 @@ end
 struct boundary_distributions_struct
     # knudsen cosine distribution for imposing the neutral wall boundary condition
     knudsen::MPISharedArray{mk_float,3}
-    # charged particle r boundary values (vpa,vperp,z,r,s)
-    pdf_rboundary_charged::MPISharedArray{mk_float,5}
+    # ion particle r boundary values (vpa,vperp,z,r,s)
+    pdf_rboundary_ion::MPISharedArray{mk_float,5}
     # neutral particle r boundary values (vz,vr,vzeta,z,r,s)
     pdf_rboundary_neutral::MPISharedArray{mk_float,6}
 end
@@ -90,7 +90,7 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
     # the time-dependent entries are not initialised.
     # moments arrays have same r and z grids for both ion and neutral species
     # and so are included in the same struct
-    charged = create_moments_charged(z.n, r.n, composition.n_ion_species,
+    ion = create_moments_ion(z.n, r.n, composition.n_ion_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, numerical_dissipation)
     neutral = create_moments_neutral(z.n, r.n, composition.n_neutral_species,
@@ -106,7 +106,7 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
         particle_number_conserved = true
     end
 
-    moments = moments_struct(charged, neutral, evolve_moments.density,
+    moments = moments_struct(ion, neutral, evolve_moments.density,
                              particle_number_conserved,
                              evolve_moments.conservation,
                              evolve_moments.parallel_flow,
@@ -123,12 +123,12 @@ Allocate arrays for pdfs
 """
 function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     # allocate pdf arrays
-    pdf_charged_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_ion_species)
-    pdf_charged_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
+    pdf_ion_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_ion_species)
+    pdf_ion_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
 
-    return pdf_struct(pdf_substruct(pdf_charged_norm, pdf_charged_buffer),
+    return pdf_struct(pdf_substruct(pdf_ion_norm, pdf_ion_buffer),
                       pdf_substruct(pdf_neutral_norm, pdf_neutral_buffer))
 
 end
@@ -149,12 +149,12 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         n_neutral_species = composition.n_neutral_species
         @serial_region begin
             # initialise the density profile
-            init_density!(moments.charged.dens, z, r, species.charged, n_ion_species)
+            init_density!(moments.ion.dens, z, r, species.ion, n_ion_species)
             # initialise the parallel flow profile
-            init_upar!(moments.charged.upar, z, r, species.charged, n_ion_species)
+            init_upar!(moments.ion.upar, z, r, species.ion, n_ion_species)
             # initialise the parallel thermal speed profile
-            init_vth!(moments.charged.vth, z, r, species.charged, n_ion_species)
-            @. moments.charged.ppar = 0.5 * moments.charged.dens * moments.charged.vth^2
+            init_vth!(moments.ion.vth, z, r, species.ion, n_ion_species)
+            @. moments.ion.ppar = 0.5 * moments.ion.dens * moments.ion.vth^2
 
             if n_neutral_species > 0
                 #neutral particles
@@ -167,9 +167,9 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
             end
         end
-        moments.charged.dens_updated .= true
-        moments.charged.upar_updated .= true
-        moments.charged.ppar_updated .= true
+        moments.ion.dens_updated .= true
+        moments.ion.upar_updated .= true
+        moments.ion.ppar_updated .= true
         moments.neutral.dens_updated .= true
         moments.neutral.uz_updated .= true
         moments.neutral.pz_updated .= true
@@ -182,9 +182,9 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                         vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
         begin_s_r_z_region()
         # calculate the initial parallel heat flux from the initial un-normalised pdf
-        update_qpar!(moments.charged.qpar, moments.charged.qpar_updated,
-                     moments.charged.dens, moments.charged.upar, moments.charged.vth,
-                     pdf.charged.norm, vpa, vperp, z, r, composition,
+        update_qpar!(moments.ion.qpar, moments.ion.qpar_updated,
+                     moments.ion.dens, moments.ion.upar, moments.ion.vth,
+                     pdf.ion.norm, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
 
         if n_neutral_species > 0
@@ -221,30 +221,30 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
         for is ∈ 1:composition.n_ion_species, ir ∈ 1:r.n
             # Add ion contributions to wall flux here. Neutral contributions will be
             # added in init_neutral_pdf_over_density!()
-            if species.charged[is].z_IC.initialization_option == "bgk" || species.charged[is].vpa_IC.initialization_option == "bgk"
-                @views init_bgk_pdf!(pdf.charged.norm[:,1,:,ir,is], 0.0, species.charged[is].initial_temperature, z.grid, z.L, vpa.grid)
+            if species.ion[is].z_IC.initialization_option == "bgk" || species.ion[is].vpa_IC.initialization_option == "bgk"
+                @views init_bgk_pdf!(pdf.ion.norm[:,1,:,ir,is], 0.0, species.ion[is].initial_temperature, z.grid, z.L, vpa.grid)
             else
                 # updates pdf_norm to contain pdf / density, so that ∫dvpa pdf.norm = 1,
                 # ∫dwpa wpa * pdf.norm = 0, and ∫dwpa m_s (wpa/vths)^2 pdf.norm = 1/2
                 # to machine precision
-                @views init_charged_pdf_over_density!(
-                    pdf.charged.norm[:,:,:,ir,is], species.charged[is], composition, vpa, vperp,
-                    z, vpa_spectral, moments.charged.dens[:,ir,is],
-                    moments.charged.upar[:,ir,is], moments.charged.ppar[:,ir,is],
-                    moments.charged.vth[:,ir,is],
-                    moments.charged.v_norm_fac[:,ir,is], moments.evolve_density,
+                @views init_ion_pdf_over_density!(
+                    pdf.ion.norm[:,:,:,ir,is], species.ion[is], composition, vpa, vperp,
+                    z, vpa_spectral, moments.ion.dens[:,ir,is],
+                    moments.ion.upar[:,ir,is], moments.ion.ppar[:,ir,is],
+                    moments.ion.vth[:,ir,is],
+                    moments.ion.v_norm_fac[:,ir,is], moments.evolve_density,
                     moments.evolve_upar, moments.evolve_ppar)
             end
-            @views wall_flux_0[ir,is] = -(moments.charged.dens[1,ir,is] *
-                                          moments.charged.upar[1,ir,is])
-            @views wall_flux_L[ir,is] = moments.charged.dens[end,ir,is] *
-                                        moments.charged.upar[end,ir,is]
+            @views wall_flux_0[ir,is] = -(moments.ion.dens[1,ir,is] *
+                                          moments.ion.upar[1,ir,is])
+            @views wall_flux_L[ir,is] = moments.ion.dens[end,ir,is] *
+                                        moments.ion.upar[end,ir,is]
 
             @loop_z iz begin
                 if moments.evolve_ppar
-                    @. pdf.charged.norm[:,:,iz,ir,is] *= moments.charged.vth[iz,ir,is]
+                    @. pdf.ion.norm[:,:,iz,ir,is] *= moments.ion.vth[iz,ir,is]
                 elseif moments.evolve_density == false
-                    @. pdf.charged.norm[:,:,iz,ir,is] *= moments.charged.dens[iz,ir,is]
+                    @. pdf.ion.norm[:,:,iz,ir,is] *= moments.ion.dens[iz,ir,is]
                 end
             end
         end
@@ -395,7 +395,7 @@ end
 
 """
 """
-function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
+function init_ion_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
         vpa_spectral, density, upar, ppar, vth, v_norm_fac, evolve_density, evolve_upar,
         evolve_ppar)
 
@@ -429,7 +429,7 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
             end
 
             # Only do this correction for runs without wall bc, because consistency of
-            # pdf and moments is taken care of by convert_full_f_charged_to_normalised!()
+            # pdf and moments is taken care of by convert_full_f_ion_to_normalised!()
             # for wall bc cases.
             for iz ∈ 1:z.n
                 # densfac = the integral of the pdf over v-space, which should be unity,
@@ -489,7 +489,7 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
                     @. pdf[:,ivperp,iz] *= 1.0 - exp(-vpa.grid^2*inverse_width)
                 end
             end
-            # Can use non-shared memory here because `init_charged_pdf_over_density!()` is
+            # Can use non-shared memory here because `init_ion_pdf_over_density!()` is
             # called inside a `@serial_region`
             lower_z_pdf_buffer = allocate_float(vpa.n, vperp.n)
             upper_z_pdf_buffer = allocate_float(vpa.n, vperp.n)
@@ -531,7 +531,7 @@ function init_charged_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
 
             # Get the unnormalised pdf and the moments of the constructed full-f
             # distribution function (which will be modified from the input moments).
-            convert_full_f_charged_to_normalised!(pdf, density, upar, ppar, vth, vperp,
+            convert_full_f_ion_to_normalised!(pdf, density, upar, ppar, vth, vperp,
                                                   vpa, vpa_spectral, evolve_density,
                                                   evolve_upar, evolve_ppar)
 
@@ -665,7 +665,7 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
                     @. pdf[:,ivr,ivzeta,iz] *= 1.0 - exp(-vz.grid^2*inverse_width)
                 end
             end
-            # Can use non-shared memory here because `init_charged_pdf_over_density!()` is
+            # Can use non-shared memory here because `init_ion_pdf_over_density!()` is
             # called inside a `@serial_region`
             lower_z_pdf_buffer = allocate_float(vz.n, vr.n, vzeta.n)
             upper_z_pdf_buffer = allocate_float(vz.n, vr.n, vzeta.n)
@@ -785,23 +785,23 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     #nb manufactured functions not functions of species
     begin_s_r_z_region()
     @loop_s_r_z is ir iz begin
-        moments.charged.dens[iz,ir,is] = densi_func(z.grid[iz],r.grid[ir],0.0)
+        moments.ion.dens[iz,ir,is] = densi_func(z.grid[iz],r.grid[ir],0.0)
         @loop_vperp_vpa ivperp ivpa begin
-            pdf.charged.norm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
+            pdf.ion.norm[ivpa,ivperp,iz,ir,is] = dfni_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],0.0)
         end
     end
     # update upar, ppar, qpar, vth consistent with manufactured solns
-    update_density!(moments.charged.dens, pdf.charged.norm, vpa, vperp, z, r, composition)
-    update_qpar!(moments.charged.qpar, pdf.charged.norm, vpa, vperp, z, r, composition)
-    update_ppar!(moments.charged.ppar, pdf.charged.norm, vpa, vperp, z, r, composition)
+    update_density!(moments.ion.dens, pdf.ion.norm, vpa, vperp, z, r, composition)
+    update_qpar!(moments.ion.qpar, pdf.ion.norm, vpa, vperp, z, r, composition)
+    update_ppar!(moments.ion.ppar, pdf.ion.norm, vpa, vperp, z, r, composition)
     # get particle flux
-    update_upar!(moments.charged.upar, pdf.charged.norm, vpa, vperp, z, r, composition)
+    update_upar!(moments.ion.upar, pdf.ion.norm, vpa, vperp, z, r, composition)
     # convert from particle particle flux to parallel flow
     begin_s_r_z_region()
     @loop_s_r_z is ir iz begin
-        moments.charged.upar[iz,ir,is] /= moments.charged.dens[iz,ir,is]
+        moments.ion.upar[iz,ir,is] /= moments.ion.dens[iz,ir,is]
     # update the thermal speed
-        moments.charged.vth[iz,ir,is] = sqrt(2.0*moments.charged.ppar[iz,ir,is]/moments.charged.dens[iz,ir,is])
+        moments.ion.vth[iz,ir,is] = sqrt(2.0*moments.ion.ppar[iz,ir,is]/moments.ion.dens[iz,ir,is])
     end
 
     if n_neutral_species > 0
@@ -838,7 +838,7 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
             moments.neutral.ur[iz,ir,isn] /= moments.neutral.dens[iz,ir,isn]
             moments.neutral.uzeta[iz,ir,isn] /= moments.neutral.dens[iz,ir,isn]
             # get vth for neutrals
-            moments.charged.vth[iz,ir,isn] = sqrt(2.0*moments.neutral.ptot[iz,ir,isn]/moments.neutral.dens[iz,ir,isn])
+            moments.ion.vth[iz,ir,isn] = sqrt(2.0*moments.neutral.ptot[iz,ir,isn]/moments.neutral.dens[iz,ir,isn])
         end
     end
     return nothing
@@ -908,7 +908,7 @@ function init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composi
     return knudsen_cosine
 end
 
-function init_rboundary_pdfs!(rboundary_charged, rboundary_neutral, pdf::pdf_struct, vz,
+function init_rboundary_pdfs!(rboundary_ion, rboundary_neutral, pdf::pdf_struct, vz,
                               vr, vzeta, vpa, vperp, z, r, composition)
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
@@ -916,8 +916,8 @@ function init_rboundary_pdfs!(rboundary_charged, rboundary_neutral, pdf::pdf_str
 
     begin_s_z_region() #do not parallelise r here
     @loop_s_z_vperp_vpa is iz ivperp ivpa begin
-        rboundary_charged[ivpa,ivperp,iz,1,is] = pdf.charged.norm[ivpa,ivperp,iz,1,is]
-        rboundary_charged[ivpa,ivperp,iz,end,is] = pdf.charged.norm[ivpa,ivperp,iz,end,is]
+        rboundary_ion[ivpa,ivperp,iz,1,is] = pdf.ion.norm[ivpa,ivperp,iz,1,is]
+        rboundary_ion[ivpa,ivperp,iz,end,is] = pdf.ion.norm[ivpa,ivperp,iz,end,is]
     end
     if n_neutral_species > 0
         begin_sn_z_region() #do not parallelise r here
@@ -926,7 +926,7 @@ function init_rboundary_pdfs!(rboundary_charged, rboundary_neutral, pdf::pdf_str
             rboundary_neutral[ivz,ivr,ivzeta,iz,end,isn] = pdf.neutral.norm[ivz,ivr,ivzeta,iz,end,isn]
         end
     end
-    return rboundary_charged, rboundary_neutral
+    return rboundary_ion, rboundary_neutral
 end
 
 """
@@ -941,17 +941,17 @@ function create_boundary_distributions(vz, vr, vzeta, vpa, vperp, z, composition
     #depends on T_wall, which has already been set
     init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composition)
     #initialise fixed-in-time radial boundary condition based on initial condition values
-    pdf_rboundary_charged = allocate_shared_float(vpa.n, vperp.n, z.n, 2,
+    pdf_rboundary_ion = allocate_shared_float(vpa.n, vperp.n, z.n, 2,
                                                   composition.n_ion_species)
     pdf_rboundary_neutral =  allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, 2,
                                                    composition.n_neutral_species)
 
-    return boundary_distributions_struct(knudsen_cosine, pdf_rboundary_charged, pdf_rboundary_neutral)
+    return boundary_distributions_struct(knudsen_cosine, pdf_rboundary_ion, pdf_rboundary_neutral)
 end
 
 function init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp, z, r, composition)
     #initialise fixed-in-time radial boundary condition based on initial condition values
-    init_rboundary_pdfs!(boundary_distributions.pdf_rboundary_charged,
+    init_rboundary_pdfs!(boundary_distributions.pdf_rboundary_ion,
                          boundary_distributions.pdf_rboundary_neutral, pdf, vz, vr, vzeta,
                          vpa, vperp, z, r, composition)
     return nothing
@@ -1046,7 +1046,7 @@ function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc:
 end
 
 """
-enforce boundary conditions on charged particle f in z
+enforce boundary conditions on ion particle f in z
 """
 function enforce_z_boundary_condition!(pdf, density, upar, ppar, moments, bc::String, adv,
                                        z, vperp, vpa, composition)
@@ -1105,7 +1105,8 @@ end
 """
 enforce boundary conditions on neutral particle distribution function
 """
-function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
+#f_ion does not seem to be used here -- can it be removed?
+function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
         boundary_distributions, density_neutral, uz_neutral, pz_neutral, moments,
         density_ion, upar_ion, Er, r_adv, z_adv, vzeta_adv, vr_adv, vz_adv, r, z, vzeta,
         vr, vz, composition, geometry, scratch_dummy, r_diffusion, vz_diffusion)
@@ -1763,7 +1764,7 @@ function vpagrid_to_dzdt(vpagrid, vth, upar, evolve_ppar, evolve_upar)
 end
 
 """
-Take the full charged-particle distribution function, calculate the moments, then
+Take the full ion distribution function, calculate the moments, then
 normalise and shift to the moment-kinetic grid.
 
 Uses input value of `f` and modifies in place to the normalised distribution functions.
@@ -1772,7 +1773,7 @@ the moments of `f`.
 
 Inputs/outputs depend on z, vperp, and vpa (should be inside loops over species, r)
 """
-function convert_full_f_charged_to_normalised!(f, density, upar, ppar, vth, vperp, vpa,
+function convert_full_f_ion_to_normalised!(f, density, upar, ppar, vth, vperp, vpa,
         vpa_spectral, evolve_density, evolve_upar, evolve_ppar)
 
     @loop_z iz begin
diff --git a/src/ionization.jl b/src/ionization.jl
index 147cb75ce..66e442f58 100644
--- a/src/ionization.jl
+++ b/src/ionization.jl
@@ -87,7 +87,7 @@ function ionization_collisions_1V!(f_out, f_neutral_out, fvec_in, vz, vpa, vperp
             if moments.evolve_ppar
                 # will need the ratio of thermal speeds both to interpolate between vpa grids
                 # for different species and to account for different normalizations of each species' pdf
-                vth_ratio = moments.charged.vth[iz,ir,is]/moments.neutral.vth[iz,ir,isn]
+                vth_ratio = moments.ion.vth[iz,ir,is]/moments.neutral.vth[iz,ir,isn]
             else
                 vth_ratio = 1.0
             end
@@ -121,7 +121,7 @@ function ionization_collisions_1V!(f_out, f_neutral_out, fvec_in, vz, vpa, vperp
                     # to get f_{s'}(wpahat_s), need to obtain wpahat_s grid locations
                     # in terms of the wpahat_{s'} coordinate:
                     # (wpahat_{s'})_j = ((wpahat_{s})_j * vth_{s} + upar_{s} - upar_{s'}) / vth_{s'}
-                    @. vpa.scratch = (vpa.grid * moments.charged.vth[iz,ir,is] + fvec_in.upar[iz,ir,is] - fvec_in.uz_neutral[iz,ir,isn]) / moments.neutral.vth[iz,ir,isn]
+                    @. vpa.scratch = (vpa.grid * moments.ion.vth[iz,ir,is] + fvec_in.upar[iz,ir,is] - fvec_in.uz_neutral[iz,ir,isn]) / moments.neutral.vth[iz,ir,isn]
                 end
                 # interpolate to the new grid (passed in as vpa.scratch)
                 # and return interpolated values in vpa.scratch2
diff --git a/src/load_data.jl b/src/load_data.jl
index 281965b6c..922ddb8df 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -4,7 +4,7 @@ module load_data
 
 export open_readonly_output_file
 export load_fields_data
-export load_charged_particle_moments_data
+export load_ion_particle_moments_data
 export load_neutral_particle_moments_data
 export load_pdf_data
 export load_neutral_pdf_data
@@ -315,26 +315,26 @@ end
 
 """
 """
-function load_charged_particle_moments_data(fid; printout=false)
+function load_ion_particle_moments_data(fid; printout=false)
     if printout
-        print("Loading charged particle velocity moments data...")
+        print("Loading ion particle velocity moments data...")
     end
 
     group = get_group(fid, "dynamic_data")
 
-    # Read charged species density
+    # Read ion species density
     density = load_variable(group, "density")
 
-    # Read charged species parallel flow
+    # Read ion species parallel flow
     parallel_flow = load_variable(group, "parallel_flow")
 
-    # Read charged species parallel pressure
+    # Read ion species parallel pressure
     parallel_pressure = load_variable(group, "parallel_pressure")
 
-    # Read charged_species parallel heat flux
+    # Read ion_species parallel heat flux
     parallel_heat_flux = load_variable(group, "parallel_heat_flux")
 
-    # Read charged species thermal speed
+    # Read ion species thermal speed
     thermal_speed = load_variable(group, "thermal_speed")
 
     if printout
@@ -377,12 +377,12 @@ end
 """
 function load_pdf_data(fid; printout=false)
     if printout
-        print("Loading charged particle distribution function data...")
+        print("Loading ion particle distribution function data...")
     end
 
     group = get_group(fid, "dynamic_data")
 
-    # Read charged distribution function
+    # Read ion distribution function
     pdf = load_variable(group, "f")
 
     if printout
@@ -475,29 +475,29 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                 vr_range = get_range(vr)
                 vz_range = get_range(vz)
 
-                pdf.charged.norm .= load_slice(dynamic, "f", vpa_range, vperp_range,
+                pdf.ion.norm .= load_slice(dynamic, "f", vpa_range, vperp_range,
                                                z_range, r_range, :, time_index)
-                moments.charged.dens .= load_slice(dynamic, "density", z_range, r_range,
+                moments.ion.dens .= load_slice(dynamic, "density", z_range, r_range,
                                                    :, time_index)
-                moments.charged.dens_updated .= true
-                moments.charged.upar .= load_slice(dynamic, "parallel_flow", z_range,
+                moments.ion.dens_updated .= true
+                moments.ion.upar .= load_slice(dynamic, "parallel_flow", z_range,
                                                    r_range, :, time_index)
-                moments.charged.upar_updated .= true
-                moments.charged.ppar .= load_slice(dynamic, "parallel_pressure", z_range,
+                moments.ion.upar_updated .= true
+                moments.ion.ppar .= load_slice(dynamic, "parallel_pressure", z_range,
                                                    r_range, :, time_index)
-                moments.charged.ppar_updated .= true
-                moments.charged.qpar .= load_slice(dynamic, "parallel_heat_flux", z_range,
+                moments.ion.ppar_updated .= true
+                moments.ion.qpar .= load_slice(dynamic, "parallel_heat_flux", z_range,
                                                    r_range, :, time_index)
-                moments.charged.qpar_updated .= true
-                moments.charged.vth .= load_slice(dynamic, "thermal_speed", z_range,
+                moments.ion.qpar_updated .= true
+                moments.ion.vth .= load_slice(dynamic, "thermal_speed", z_range,
                                                   r_range, :, time_index)
 
                 boundary_distributions_io = get_group(fid, "boundary_distributions")
-                boundary_distributions.pdf_rboundary_charged[:,:,:,1,:] .=
-                load_slice(boundary_distributions_io, "pdf_rboundary_charged_left",
+                boundary_distributions.pdf_rboundary_ion[:,:,:,1,:] .=
+                load_slice(boundary_distributions_io, "pdf_rboundary_ion_left",
                            vpa_range, vperp_range, z_range, :)
-                boundary_distributions.pdf_rboundary_charged[:,:,:,2,:] .=
-                load_slice(boundary_distributions_io, "pdf_rboundary_charged_right",
+                boundary_distributions.pdf_rboundary_ion[:,:,:,2,:] .=
+                load_slice(boundary_distributions_io, "pdf_rboundary_ion_right",
                            vpa_range, vperp_range, z_range, :)
 
                 if composition.n_neutral_species > 0
@@ -554,27 +554,27 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
 
                 code_time = load_slice(dynamic, "time", time_index)
 
-                pdf.charged.norm .= load_slice(dynamic, "f", :, :, :, :, :, time_index)
-                moments.charged.dens .= load_slice(dynamic, "density", :, :, :,
+                pdf.ion.norm .= load_slice(dynamic, "f", :, :, :, :, :, time_index)
+                moments.ion.dens .= load_slice(dynamic, "density", :, :, :,
                                                    time_index)
-                moments.charged.dens_updated .= true
-                moments.charged.upar .= load_slice(dynamic, "parallel_flow", :, :, :,
+                moments.ion.dens_updated .= true
+                moments.ion.upar .= load_slice(dynamic, "parallel_flow", :, :, :,
                                                    time_index)
-                moments.charged.upar_updated .= true
-                moments.charged.ppar .= load_slice(dynamic, "parallel_pressure", :, :, :,
+                moments.ion.upar_updated .= true
+                moments.ion.ppar .= load_slice(dynamic, "parallel_pressure", :, :, :,
                                                    time_index)
-                moments.charged.ppar_updated .= true
-                moments.charged.qpar .= load_slice(dynamic, "parallel_heat_flux", :, :, :,
+                moments.ion.ppar_updated .= true
+                moments.ion.qpar .= load_slice(dynamic, "parallel_heat_flux", :, :, :,
                                                    time_index)
-                moments.charged.qpar_updated .= true
-                moments.charged.vth .= load_slice(dynamic, "thermal_speed", :, :, :,
+                moments.ion.qpar_updated .= true
+                moments.ion.vth .= load_slice(dynamic, "thermal_speed", :, :, :,
                                                   time_index)
 
                 boundary_distributions_io = get_group(fid, "boundary_distributions")
-                boundary_distributions.pdf_rboundary_charged[:,:,:,1,:] .=
-                load_variable(boundary_distributions_io, "pdf_rboundary_charged_left")
-                boundary_distributions.pdf_rboundary_charged[:,:,:,2,:] .=
-                load_variable(boundary_distributions_io, "pdf_rboundary_charged_right")
+                boundary_distributions.pdf_rboundary_ion[:,:,:,1,:] .=
+                load_variable(boundary_distributions_io, "pdf_rboundary_ion_left")
+                boundary_distributions.pdf_rboundary_ion[:,:,:,2,:] .=
+                load_variable(boundary_distributions_io, "pdf_rboundary_ion_right")
 
                 if composition.n_neutral_species > 0
                     pdf.neutral.norm .= load_slice(dynamic, "f_neutral", :, :, :, :, :, :,
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index c321d51c7..dfaad2a3c 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -393,7 +393,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
 
     write_moments_data_to_binary(moments, fields, code_time, composition.n_ion_species,
         composition.n_neutral_species, io_moments, 1, r, z)
-    write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, moments, fields,
+    write_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
          code_time, composition.n_ion_species, composition.n_neutral_species, io_dfns, 1,
          r, z, vperp, vpa, vzeta, vr, vz)
 
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index ad2026ddd..7284c04b2 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -122,30 +122,30 @@ function mk_input(scan_input=Dict())
     #println("Info: rhostar is ",geometry.rhostar)
     
     ispecies = 1
-    species.charged[1].z_IC.initialization_option = get(scan_input, "z_IC_option$ispecies", "gaussian")
-    species.charged[1].initial_density = get(scan_input, "initial_density$ispecies", 1.0)
-    species.charged[1].initial_temperature = get(scan_input, "initial_temperature$ispecies", 1.0)
-    species.charged[1].z_IC.width = get(scan_input, "z_IC_width$ispecies", 0.125)
-    species.charged[1].z_IC.density_amplitude = get(scan_input, "z_IC_density_amplitude$ispecies", 0.001)
-    species.charged[1].z_IC.density_phase = get(scan_input, "z_IC_density_phase$ispecies", 0.0)
-    species.charged[1].z_IC.upar_amplitude = get(scan_input, "z_IC_upar_amplitude$ispecies", 0.0)
-    species.charged[1].z_IC.upar_phase = get(scan_input, "z_IC_upar_phase$ispecies", 0.0)
-    species.charged[1].z_IC.temperature_amplitude = get(scan_input, "z_IC_temperature_amplitude$ispecies", 0.0)
-    species.charged[1].z_IC.temperature_phase = get(scan_input, "z_IC_temperature_phase$ispecies", 0.0)
-    species.charged[1].vpa_IC.initialization_option = get(scan_input, "vpa_IC_option$ispecies", "gaussian")
-    species.charged[1].vpa_IC.density_amplitude = get(scan_input, "vpa_IC_density_amplitude$ispecies", 1.000)
-    species.charged[1].vpa_IC.width = get(scan_input, "vpa_IC_width$ispecies", 1.0)
-    species.charged[1].vpa_IC.density_phase = get(scan_input, "vpa_IC_density_phase$ispecies", 0.0)
-    species.charged[1].vpa_IC.upar_amplitude = get(scan_input, "vpa_IC_upar_amplitude$ispecies", 0.0)
-    species.charged[1].vpa_IC.upar_phase = get(scan_input, "vpa_IC_upar_phase$ispecies", 0.0)
-    species.charged[1].vpa_IC.temperature_amplitude = get(scan_input, "vpa_IC_temperature_amplitude$ispecies", 0.0)
-    species.charged[1].vpa_IC.temperature_phase = get(scan_input, "vpa_IC_temperature_phase$ispecies", 0.0)
+    species.ion[1].z_IC.initialization_option = get(scan_input, "z_IC_option$ispecies", "gaussian")
+    species.ion[1].initial_density = get(scan_input, "initial_density$ispecies", 1.0)
+    species.ion[1].initial_temperature = get(scan_input, "initial_temperature$ispecies", 1.0)
+    species.ion[1].z_IC.width = get(scan_input, "z_IC_width$ispecies", 0.125)
+    species.ion[1].z_IC.density_amplitude = get(scan_input, "z_IC_density_amplitude$ispecies", 0.001)
+    species.ion[1].z_IC.density_phase = get(scan_input, "z_IC_density_phase$ispecies", 0.0)
+    species.ion[1].z_IC.upar_amplitude = get(scan_input, "z_IC_upar_amplitude$ispecies", 0.0)
+    species.ion[1].z_IC.upar_phase = get(scan_input, "z_IC_upar_phase$ispecies", 0.0)
+    species.ion[1].z_IC.temperature_amplitude = get(scan_input, "z_IC_temperature_amplitude$ispecies", 0.0)
+    species.ion[1].z_IC.temperature_phase = get(scan_input, "z_IC_temperature_phase$ispecies", 0.0)
+    species.ion[1].vpa_IC.initialization_option = get(scan_input, "vpa_IC_option$ispecies", "gaussian")
+    species.ion[1].vpa_IC.density_amplitude = get(scan_input, "vpa_IC_density_amplitude$ispecies", 1.000)
+    species.ion[1].vpa_IC.width = get(scan_input, "vpa_IC_width$ispecies", 1.0)
+    species.ion[1].vpa_IC.density_phase = get(scan_input, "vpa_IC_density_phase$ispecies", 0.0)
+    species.ion[1].vpa_IC.upar_amplitude = get(scan_input, "vpa_IC_upar_amplitude$ispecies", 0.0)
+    species.ion[1].vpa_IC.upar_phase = get(scan_input, "vpa_IC_upar_phase$ispecies", 0.0)
+    species.ion[1].vpa_IC.temperature_amplitude = get(scan_input, "vpa_IC_temperature_amplitude$ispecies", 0.0)
+    species.ion[1].vpa_IC.temperature_phase = get(scan_input, "vpa_IC_temperature_phase$ispecies", 0.0)
     ispecies += 1
     if n_neutral_species > 0
         species.neutral[1].z_IC.initialization_option = get(scan_input, "z_IC_option$ispecies", "gaussian")
         species.neutral[1].initial_density = get(scan_input, "initial_density$ispecies", 1.0)
         species.neutral[1].initial_temperature = get(scan_input, "initial_temperature$ispecies", 1.0)
-        species.neutral[1].z_IC.width = get(scan_input, "z_IC_width$ispecies", species.charged[1].z_IC.width)
+        species.neutral[1].z_IC.width = get(scan_input, "z_IC_width$ispecies", species.ion[1].z_IC.width)
         species.neutral[1].z_IC.density_amplitude = get(scan_input, "z_IC_density_amplitude$ispecies", 0.001)
         species.neutral[1].z_IC.density_phase = get(scan_input, "z_IC_density_phase$ispecies", 0.0)
         species.neutral[1].z_IC.upar_amplitude = get(scan_input, "z_IC_upar_amplitude$ispecies", 0.0)
@@ -154,7 +154,7 @@ function mk_input(scan_input=Dict())
         species.neutral[1].z_IC.temperature_phase = get(scan_input, "z_IC_temperature_phase$ispecies", 0.0)
         species.neutral[1].vpa_IC.initialization_option = get(scan_input, "vpa_IC_option$ispecies", "gaussian")
         species.neutral[1].vpa_IC.density_amplitude = get(scan_input, "vpa_IC_density_amplitude$ispecies", 1.000)
-        species.neutral[1].vpa_IC.width = get(scan_input, "vpa_IC_width$ispecies", species.charged[1].vpa_IC.width)
+        species.neutral[1].vpa_IC.width = get(scan_input, "vpa_IC_width$ispecies", species.ion[1].vpa_IC.width)
         species.neutral[1].vpa_IC.density_phase = get(scan_input, "vpa_IC_density_phase$ispecies", 0.0)
         species.neutral[1].vpa_IC.upar_amplitude = get(scan_input, "vpa_IC_upar_amplitude$ispecies", 0.0)
         species.neutral[1].vpa_IC.upar_phase = get(scan_input, "vpa_IC_upar_phase$ispecies", 0.0)
@@ -164,13 +164,13 @@ function mk_input(scan_input=Dict())
     end
     #################### end specification of species inputs #####################
 
-    collisions.charge_exchange = get(scan_input, "charge_exchange_frequency", 2.0*sqrt(species.charged[1].initial_temperature))
+    collisions.charge_exchange = get(scan_input, "charge_exchange_frequency", 2.0*sqrt(species.ion[1].initial_temperature))
     collisions.ionization = get(scan_input, "ionization_frequency", collisions.charge_exchange)
     collisions.constant_ionization_rate = get(scan_input, "constant_ionization_rate", false)
 
     # parameters related to the time stepping
     nstep = get(scan_input, "nstep", 5)
-    dt = get(scan_input, "dt", 0.00025/sqrt(species.charged[1].initial_temperature))
+    dt = get(scan_input, "dt", 0.00025/sqrt(species.ion[1].initial_temperature))
     nwrite_moments = get(scan_input, "nwrite", 1)
     nwrite_dfns = get(scan_input, "nwrite_dfns", nstep)
     # options are n_rk_stages = 1, 2, 3 or 4 (corresponding to forward Euler,
@@ -224,7 +224,7 @@ function mk_input(scan_input=Dict())
 	# do not parallelise vpa with distributed-memory MPI
     vpa.nelement_local = vpa.nelement_global 
     # L is the box length in units of vthermal_species
-    vpa.L = get(scan_input, "vpa_L", 8.0*sqrt(species.charged[1].initial_temperature))
+    vpa.L = get(scan_input, "vpa_L", 8.0*sqrt(species.ion[1].initial_temperature))
     # determine the boundary condition
     # only supported option at present is "zero" and "periodic"
     vpa.bc = get(scan_input, "vpa_bc", "periodic")
@@ -244,7 +244,7 @@ function mk_input(scan_input=Dict())
     # do not parallelise vperp with distributed-memory MPI
     vperp.nelement_local = vperp.nelement_global 
     # L is the box length in units of vthermal_species
-    vperp.L = get(scan_input, "vperp_L", 8.0*sqrt(species.charged[1].initial_temperature))
+    vperp.L = get(scan_input, "vperp_L", 8.0*sqrt(species.ion[1].initial_temperature))
     # determine the boundary condition
     # only supported option at present is "zero" and "periodic"
     # MRH probably need to add new bc option here
@@ -288,7 +288,7 @@ function mk_input(scan_input=Dict())
         # do not parallelise vz with distributed-memory MPI
         vr.nelement_local = vr.nelement_global
         # L is the box length in units of vthermal_species
-        vr.L = get(scan_input, "vr_L", 8.0*sqrt(species.charged[1].initial_temperature))
+        vr.L = get(scan_input, "vr_L", 8.0*sqrt(species.ion[1].initial_temperature))
         # determine the boundary condition
         # only supported option at present is "zero" and "periodic"
         vr.bc = get(scan_input, "vr_bc", "none")
@@ -304,7 +304,7 @@ function mk_input(scan_input=Dict())
         # do not parallelise vz with distributed-memory MPI
         vzeta.nelement_local = vzeta.nelement_global
         # L is the box length in units of vthermal_species
-        vzeta.L = get(scan_input, "vzeta_L", 8.0*sqrt(species.charged[1].initial_temperature))
+        vzeta.L = get(scan_input, "vzeta_L", 8.0*sqrt(species.ion[1].initial_temperature))
         # determine the boundary condition
         # only supported option at present is "zero" and "periodic"
         vzeta.bc = get(scan_input, "vzeta_bc", "none")
@@ -374,26 +374,26 @@ function mk_input(scan_input=Dict())
     vzeta_immutable = grid_input("vzeta", vzeta.ngrid, vzeta.nelement_global, vzeta.nelement_local, 1, 0, vzeta.L,
         vzeta.discretization, vzeta.fd_option, vzeta.bc, vzeta_advection_immutable, MPI.COMM_NULL)
     
-    species_charged_immutable = Array{species_parameters,1}(undef,n_ion_species)
+    species_ion_immutable = Array{species_parameters,1}(undef,n_ion_species)
     species_neutral_immutable = Array{species_parameters,1}(undef,n_neutral_species)
     
     for is ∈ 1:n_ion_species
         species_type = "ion"
         #    species_type = "electron"
-        z_IC = initial_condition_input(species.charged[is].z_IC.initialization_option,
-            species.charged[is].z_IC.width, species.charged[is].z_IC.wavenumber,
-            species.charged[is].z_IC.density_amplitude, species.charged[is].z_IC.density_phase,
-            species.charged[is].z_IC.upar_amplitude, species.charged[is].z_IC.upar_phase,
-            species.charged[is].z_IC.temperature_amplitude, species.charged[is].z_IC.temperature_phase,
-            species.charged[is].z_IC.monomial_degree)
-        vpa_IC = initial_condition_input(species.charged[is].vpa_IC.initialization_option,
-            species.charged[is].vpa_IC.width, species.charged[is].vpa_IC.wavenumber,
-            species.charged[is].vpa_IC.density_amplitude, species.charged[is].vpa_IC.density_phase,
-            species.charged[is].vpa_IC.upar_amplitude, species.charged[is].vpa_IC.upar_phase,
-            species.charged[is].vpa_IC.temperature_amplitude,
-            species.charged[is].vpa_IC.temperature_phase, species.charged[is].vpa_IC.monomial_degree)
-        species_charged_immutable[is] = species_parameters(species_type, species.charged[is].initial_temperature,
-            species.charged[is].initial_density, z_IC, vpa_IC)
+        z_IC = initial_condition_input(species.ion[is].z_IC.initialization_option,
+            species.ion[is].z_IC.width, species.ion[is].z_IC.wavenumber,
+            species.ion[is].z_IC.density_amplitude, species.ion[is].z_IC.density_phase,
+            species.ion[is].z_IC.upar_amplitude, species.ion[is].z_IC.upar_phase,
+            species.ion[is].z_IC.temperature_amplitude, species.ion[is].z_IC.temperature_phase,
+            species.ion[is].z_IC.monomial_degree)
+        vpa_IC = initial_condition_input(species.ion[is].vpa_IC.initialization_option,
+            species.ion[is].vpa_IC.width, species.ion[is].vpa_IC.wavenumber,
+            species.ion[is].vpa_IC.density_amplitude, species.ion[is].vpa_IC.density_phase,
+            species.ion[is].vpa_IC.upar_amplitude, species.ion[is].vpa_IC.upar_phase,
+            species.ion[is].vpa_IC.temperature_amplitude,
+            species.ion[is].vpa_IC.temperature_phase, species.ion[is].vpa_IC.monomial_degree)
+        species_ion_immutable[is] = species_parameters(species_type, species.ion[is].initial_temperature,
+            species.ion[is].initial_density, z_IC, vpa_IC)
     end
     if n_neutral_species > 0
         for is ∈ 1:n_neutral_species
@@ -414,7 +414,7 @@ function mk_input(scan_input=Dict())
                 species.neutral[is].initial_density, z_IC, vpa_IC)
         end
     end 
-    species_immutable = (charged = species_charged_immutable, neutral = species_neutral_immutable)
+    species_immutable = (ion = species_ion_immutable, neutral = species_neutral_immutable)
     
     force_Er_zero = get(scan_input, "force_Er_zero_at_wall", false)
     drive_immutable = drive_input(drive.force_phi, drive.amplitude, drive.frequency, force_Er_zero)
@@ -779,7 +779,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
         epsilon_offset, use_vpabar_in_mms_dfni, alpha_switch, mn_over_mi, me_over_mi,
         allocate_float(n_species))
     
-    species_charged = Array{species_parameters_mutable,1}(undef,n_ion_species)
+    species_ion = Array{species_parameters_mutable,1}(undef,n_ion_species)
     species_neutral = Array{species_parameters_mutable,1}(undef,n_neutral_species)
     
     # initial temperature for each species defaults to Tₑ
@@ -830,7 +830,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
 
     # fill in entries in species struct corresponding to ion species
     for is ∈ 1:n_ion_species
-        species_charged[is] = species_parameters_mutable("ion", initial_temperature, initial_density,
+        species_ion[is] = species_parameters_mutable("ion", initial_temperature, initial_density,
             deepcopy(z_initial_conditions), deepcopy(vpa_initial_conditions))
     end
     # if there are neutrals, fill in corresponding entries in species struct
@@ -840,7 +840,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
                 initial_density, deepcopy(z_initial_conditions), deepcopy(vpa_initial_conditions))
         end
     end
-    species = (charged = species_charged, neutral = species_neutral)
+    species = (ion = species_ion, neutral = species_neutral)
     
     # if drive_phi = true, include external electrostatic potential of form
     # phi(z,t=0)*drive_amplitude*sinpi(time*drive_frequency)
diff --git a/src/moment_kinetics_structs.jl b/src/moment_kinetics_structs.jl
index fa63cd9d0..56829f815 100644
--- a/src/moment_kinetics_structs.jl
+++ b/src/moment_kinetics_structs.jl
@@ -10,9 +10,9 @@ using ..type_definitions: mk_float
 
 """
 """
-struct scratch_pdf{n_distribution_charged, n_moment, n_distribution_neutral,n_moment_neutral}
-    # charged particles
-    pdf::MPISharedArray{mk_float, n_distribution_charged}
+struct scratch_pdf{n_distribution_ion, n_moment, n_distribution_neutral,n_moment_neutral}
+    # ion particles
+    pdf::MPISharedArray{mk_float, n_distribution_ion}
     density::MPISharedArray{mk_float, n_moment}
     upar::MPISharedArray{mk_float, n_moment}
     ppar::MPISharedArray{mk_float, n_moment}
diff --git a/src/plot_MMS_sequence.jl b/src/plot_MMS_sequence.jl
index 51cb75d10..eadecbbf8 100644
--- a/src/plot_MMS_sequence.jl
+++ b/src/plot_MMS_sequence.jl
@@ -14,16 +14,16 @@ using SpecialFunctions: erfi
 using LaTeXStrings
 # modules
 using ..post_processing_input: pp
-using ..post_processing: compare_charged_pdf_symbolic_test, compare_fields_symbolic_test
+using ..post_processing: compare_ion_pdf_symbolic_test, compare_fields_symbolic_test
 using ..post_processing: compare_moments_symbolic_test, compare_neutral_pdf_symbolic_test
 using ..post_processing: read_distributed_zr_data!, construct_global_zr_coords
-using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_charged_moments
+using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_ion_moments
 using ..post_processing: allocate_global_zr_fields#, get_geometry_and_composition, get_coords_nelement
 using ..array_allocation: allocate_float
 using ..type_definitions: mk_float, mk_int
 using ..load_data: open_readonly_output_file
 using ..load_data: load_fields_data, load_pdf_data
-using ..load_data: load_charged_particle_moments_data, load_neutral_particle_moments_data
+using ..load_data: load_ion_particle_moments_data, load_neutral_particle_moments_data
 using ..load_data: load_neutral_pdf_data, load_time_data, load_species_data
 using ..load_data: load_block_data, load_coordinate_data
 using ..velocity_moments: integrate_over_vspace
@@ -167,7 +167,7 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         # allocate arrays to contain the global fields as a function of (z,r,t)
         phi, Ez, Er = allocate_global_zr_fields(z.n_global,r.n_global,ntime)
         density, parallel_flow, parallel_pressure, parallel_heat_flux,
-            thermal_speed = allocate_global_zr_charged_moments(z.n_global,r.n_global,n_ion_species,ntime)
+            thermal_speed = allocate_global_zr_ion_moments(z.n_global,r.n_global,n_ion_species,ntime)
         if n_neutral_species > 0
             neutral_density, neutral_uz, neutral_pz, 
              neutral_qz, neutral_thermal_speed = allocate_global_zr_neutral_moments(z.n_global,r.n_global,n_neutral_species,ntime)
@@ -182,7 +182,7 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         read_distributed_zr_data!(phi,"phi",run_name,"moments",nblocks,z.n,r.n) 
         read_distributed_zr_data!(Ez,"Ez",run_name,"moments",nblocks,z.n,r.n) 
         read_distributed_zr_data!(Er,"Er",run_name,"moments",nblocks,z.n,r.n) 
-        # charged particle moments
+        # ion particle moments
         read_distributed_zr_data!(density,"density",run_name,"moments",nblocks,z.n,r.n) 
         read_distributed_zr_data!(parallel_flow,"parallel_flow",run_name,"moments",nblocks,z.n,r.n) 
         read_distributed_zr_data!(parallel_pressure,"parallel_pressure",run_name,"moments",nblocks,z.n,r.n) 
@@ -257,7 +257,7 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         # use final time point for analysis
         ion_density_error_sequence[isim] = ion_density_error_t[end] 
         
-        ion_pdf_error_t = compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
+        ion_pdf_error_t = compare_ion_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
          L"\widetilde{f}_i",L"\widetilde{f}^{sym}_i",L"\sum || \widetilde{f}_i - \widetilde{f}_i^{sym} ||^2","pdf")
         ion_pdf_error_sequence[isim] = ion_pdf_error_t[end]
         
diff --git a/src/plot_sequence.jl b/src/plot_sequence.jl
index 370a13b16..378fd6ae2 100644
--- a/src/plot_sequence.jl
+++ b/src/plot_sequence.jl
@@ -12,7 +12,7 @@ using SpecialFunctions: erfi
 using LaTeXStrings
 # modules
 using ..post_processing: read_distributed_zr_data!, construct_global_zr_coords
-using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_charged_moments
+using ..post_processing: allocate_global_zr_neutral_moments, allocate_global_zr_ion_moments
 using ..post_processing: allocate_global_zr_fields#, get_coords_nelement
 using ..array_allocation: allocate_float
 using ..type_definitions: mk_float, mk_int
diff --git a/src/post_processing.jl b/src/post_processing.jl
index e5c7fa267..c793ac540 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -3,12 +3,12 @@
 module post_processing
 
 export analyze_and_plot_data
-export compare_charged_pdf_symbolic_test
+export compare_ion_pdf_symbolic_test
 export compare_moments_symbolic_test
 export compare_neutral_pdf_symbolic_test
 export compare_fields_symbolic_test
 export construct_global_zr_coords
-export allocate_global_zr_charged_moments
+export allocate_global_zr_ion_moments
 export allocate_global_zr_neutral_moments
 export allocate_global_zr_fields
 export get_coords_nelement
@@ -38,7 +38,7 @@ using ..initial_conditions: vpagrid_to_dzdt
 using ..load_data: open_readonly_output_file, get_group, load_input, load_time_data
 using ..load_data: get_nranks
 using ..load_data: load_fields_data, load_pdf_data
-using ..load_data: load_charged_particle_moments_data, load_neutral_particle_moments_data
+using ..load_data: load_ion_particle_moments_data, load_neutral_particle_moments_data
 using ..load_data: load_neutral_pdf_data
 using ..load_data: load_variable
 using ..load_data: load_coordinate_data, load_block_data, load_rank_data,
@@ -220,7 +220,7 @@ function allocate_global_zr_fields(nz_global,nr_global,ntime)
     return phi, Ez, Er
 end
 
-function allocate_global_zr_charged_moments(nz_global,nr_global,n_ion_species,ntime)
+function allocate_global_zr_ion_moments(nz_global,nr_global,n_ion_species,ntime)
     density = allocate_float(nz_global,nr_global,n_ion_species,ntime)
     parallel_flow = allocate_float(nz_global,nr_global,n_ion_species,ntime)
     parallel_pressure = allocate_float(nz_global,nr_global,n_ion_species,ntime)
@@ -229,7 +229,7 @@ function allocate_global_zr_charged_moments(nz_global,nr_global,n_ion_species,nt
     return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed
 end
 
-function allocate_global_zr_charged_dfns(nvpa_global, nvperp_global, nz_global, nr_global,
+function allocate_global_zr_ion_dfns(nvpa_global, nvperp_global, nz_global, nr_global,
                                          n_ion_species, ntime)
     f = allocate_float(nvpa_global, nvperp_global, nz_global, nr_global, n_ion_species,
                        ntime)
@@ -406,7 +406,7 @@ function analyze_and_plot_data(prefix...)
                                              Tuple(this_r.n_global for this_r ∈ r),
                                              ntime)
     density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+        get_tuple_of_return_values(allocate_global_zr_ion_moments,
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r),
                                    n_ion_species, ntime)
@@ -435,7 +435,7 @@ function analyze_and_plot_data(prefix...)
                                nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r))
-    # charged particle moments
+    # ion particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density, "density", run_names,
                                "moments", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
@@ -506,7 +506,7 @@ function analyze_and_plot_data(prefix...)
                                    Tuple(this_r.n_global for this_r ∈ r), ntime_pdfs)
     density_at_pdf_times, parallel_flow_at_pdf_times, parallel_pressure_at_pdf_times,
     parallel_heat_flux_at_pdf_times, thermal_speed_at_pdf_times =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+        get_tuple_of_return_values(allocate_global_zr_ion_moments,
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r), n_ion_species,
                                    ntime_pdfs)
@@ -531,7 +531,7 @@ function analyze_and_plot_data(prefix...)
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r))
-    # charged particle moments
+    # ion particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density_at_pdf_times, "density",
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
@@ -601,7 +601,7 @@ function analyze_and_plot_data(prefix...)
     diagnostics_1d = true
     if diagnostics_1d
         # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-        ff = get_tuple_of_return_values(allocate_global_zr_charged_dfns,
+        ff = get_tuple_of_return_values(allocate_global_zr_ion_dfns,
                                         Tuple(this_vpa.n_global for this_vpa ∈ vpa),
                                         Tuple(this_vperp.n_global for this_vperp ∈ vperp),
                                         Tuple(this_z.n_global for this_z ∈ z),
@@ -715,7 +715,7 @@ function analyze_and_plot_data(prefix...)
     composition = composition[1]
 
     # make plots and animations of the phi, Ez and Er
-    plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time,
+    plot_ion_moments_2D(density, parallel_flow, parallel_pressure, time,
                             z_global.grid, r_global.grid, iz0, ir0, n_ion_species,
                             itime_min, itime_max, nwrite_movie, run_name, pp)
     # make plots and animations of the phi, Ez and Er
@@ -725,12 +725,12 @@ function analyze_and_plot_data(prefix...)
     # only if ntime == ntime_pdfs & data on one shared memory process
     if ntime == ntime_pdfs && r.n_global == r.n && z.n_global == z.n
         # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-        ff = allocate_global_zr_charged_dfns(vpa.n_global, vperp.n_global, z.n_global,
+        ff = allocate_global_zr_ion_dfns(vpa.n_global, vperp.n_global, z.n_global,
                                              r.n_global, n_ion_species, ntime)
         read_distributed_zr_data!(ff, "f", run_names, "dfns", nblocks, z.n, r.n)
 
         spec_type = "ion"
-        plot_charged_pdf(ff, vpa_local, vperp_local, z_global.grid, r_global.grid, ivpa0,
+        plot_ion_pdf(ff, vpa_local, vperp_local, z_global.grid, r_global.grid, ivpa0,
                          ivperp0, iz0, ir0, spec_type, n_ion_species, itime_min_pdfs,
                          itime_max_pdfs, nwrite_movie_pdfs, run_name, pp)
         # make plots and animations of the neutral pdf
@@ -744,7 +744,7 @@ function analyze_and_plot_data(prefix...)
     end
     # plot ion pdf data near the wall boundary
     if pp.plot_wall_pdf
-        plot_charged_pdf_2D_at_wall(run_name)
+        plot_ion_pdf_2D_at_wall(run_name)
     end
     # MRH need to get some run-time data here without copy-paste from mk_input
 
@@ -808,7 +808,7 @@ function analyze_and_plot_data(prefix...)
         compare_moments_symbolic_test(run_name,density,density_sym,"ion",z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{n}_i",L"\widetilde{n}_i^{sym}",L"\varepsilon(\widetilde{n}_i)","dens")
 
-        compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
+        compare_ion_pdf_symbolic_test(run_name,manufactured_solns_list,"ion",
           L"\widetilde{f}_i",L"\widetilde{f}^{sym}_i",L"\varepsilon(\widetilde{f}_i)","pdf")
         if n_neutral_species > 0
             # neutral test
@@ -2210,7 +2210,7 @@ function compare_moments_symbolic_test(run_name,moment,moment_sym,spec_string,z,
 
 end
 
-function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec_string,
+function compare_ion_pdf_symbolic_test(run_name,manufactured_solns_list,spec_string,
     pdf_label,pdf_sym_label,norm_label,file_string)
     fid = open_readonly_output_file(run_name,"dfns", printout=false)
     # load block data on iblock=0
@@ -2224,7 +2224,7 @@ function compare_charged_pdf_symbolic_test(run_name,manufactured_solns_list,spec
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time = load_time_data(fid, printout=false)
     close(fid)
-    # get the charged particle pdf
+    # get the ion particle pdf
     dfni_func = manufactured_solns_list.dfni_func
     is = 1 # only one species supported currently
 
@@ -2306,7 +2306,7 @@ function compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,spec
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time = load_time_data(fid, printout=false)
     close(fid)
-    # get the charged particle pdf
+    # get the ion particle pdf
     dfnn_func = manufactured_solns_list.dfnn_func
     is = 1 # only one species supported currently
 
@@ -2389,7 +2389,7 @@ end
 """
 plots various slices of the ion pdf (1d and 2d, stills and animations)
 """
-function plot_charged_pdf(pdf, vpa, vperp, z, r,
+function plot_ion_pdf(pdf, vpa, vperp, z, r,
     ivpa0, ivperp0, iz0, ir0,
     spec_type, n_species,
     itime_min_pdfs, itime_max_pdfs, nwrite_movie_pdfs, run_name, pp)
@@ -2607,10 +2607,10 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
     println("done.")
 end
 
-function plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time, z, r, iz0, ir0, n_ion_species,
+function plot_ion_moments_2D(density, parallel_flow, parallel_pressure, time, z, r, iz0, ir0, n_ion_species,
     itime_min, itime_max, nwrite_movie, run_name, pp)
     nr = size(r,1)
-    print("Plotting charged moments data...")
+    print("Plotting ion moments data...")
     for is in 1:n_ion_species
 		description = "_ion_spec"*string(is)*"_"
 		# the density
@@ -2700,8 +2700,8 @@ function plot_charged_moments_2D(density, parallel_flow, parallel_pressure, time
     println("done.")
 end
 
-function plot_charged_pdf_2D_at_wall(run_name)
-    print("Plotting charged pdf data at wall boundaries...")
+function plot_ion_pdf_2D_at_wall(run_name)
+    print("Plotting ion pdf data at wall boundaries...")
     # open a dfn file
     fid = open_readonly_output_file(run_name,"dfns", printout=false)
     # load block data on iblock=0
diff --git a/src/r_advection.jl b/src/r_advection.jl
index dcef106e6..97b5bad63 100644
--- a/src/r_advection.jl
+++ b/src/r_advection.jl
@@ -21,7 +21,7 @@ function r_advection!(f_out, fvec_in, moments, fields, advect, r, z, vperp, vpa,
     @loop_s is begin
         # get the updated speed along the r direction using the current f
         @views update_speed_r!(advect[is], fvec_in.upar[:,:,is],
-                               moments.charged.vth[:,:,is], fields, moments.evolve_upar,
+                               moments.ion.vth[:,:,is], fields, moments.evolve_upar,
                                moments.evolve_ppar, vpa, vperp, z, r, geometry)
         # advance r-advection equation
         @loop_z_vpa iz ivpa begin
@@ -30,14 +30,14 @@ function r_advection!(f_out, fvec_in, moments, fields, advect, r, z, vperp, vpa,
         @loop_z_vperp_vpa iz ivperp ivpa begin
             @views adjust_advection_speed!(advect[is].speed[:,ivpa,ivperp,iz],
                                            fvec_in.density[iz,:,is],
-                                           moments.charged.vth[iz,:,is],
+                                           moments.ion.vth[iz,:,is],
                                            moments.evolve_density, moments.evolve_ppar)
             # take the normalized pdf contained in fvec_in.pdf and remove the normalization,
             # returning the true (un-normalized) particle distribution function in r.scratch
             @views unnormalize_pdf!(
                        scratch_dummy.buffer_vpavperpzrs_2[ivpa,ivperp,iz,:,is],
                        fvec_in.pdf[ivpa,ivperp,iz,:,is], fvec_in.density[iz,:,is],
-                       moments.charged.vth[iz,:,is], moments.evolve_density,
+                       moments.ion.vth[iz,:,is], moments.evolve_density,
                        moments.evolve_ppar)
             advect[is].adv_fac[:,ivpa,ivperp,iz] .= -dt.*advect[is].speed[:,ivpa,ivperp,iz]
         end
diff --git a/src/source_terms.jl b/src/source_terms.jl
index 8b5241d18..5f13bdb5c 100644
--- a/src/source_terms.jl
+++ b/src/source_terms.jl
@@ -21,8 +21,8 @@ function source_terms!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectral, compo
         @loop_s is begin
             @views source_terms_evolve_ppar_no_collisions!(
                 pdf_out[:,:,:,:,is], fvec_in.pdf[:,:,:,:,is], fvec_in.density[:,:,is],
-                fvec_in.upar[:,:,is], fvec_in.ppar[:,:,is], moments.charged.vth[:,:,is],
-                moments.charged.qpar[:,:,is], z, r, dt, spectral)
+                fvec_in.upar[:,:,is], fvec_in.ppar[:,:,is], moments.ion.vth[:,:,is],
+                moments.ion.qpar[:,:,is], z, r, dt, spectral)
             if composition.n_neutral_species > 0
                 if abs(collisions.charge_exchange) > 0.0 || abs(collisions.ionization) > 0.0
                     @views source_terms_evolve_ppar_collisions!(
@@ -217,7 +217,7 @@ end
 """
 advance the dfn with an arbitrary source function 
 """
-function source_terms_manufactured!(pdf_charged_out, pdf_neutral_out, vz, vr, vzeta, vpa, vperp, z, r, t, dt, composition, manufactured_source_list)
+function source_terms_manufactured!(pdf_ion_out, pdf_neutral_out, vz, vr, vzeta, vpa, vperp, z, r, t, dt, composition, manufactured_source_list)
     if manufactured_source_list.time_independent_sources
         # the (time-independent) manufactured source arrays
         Source_i = manufactured_source_list.Source_i_array
@@ -227,7 +227,7 @@ function source_terms_manufactured!(pdf_charged_out, pdf_neutral_out, vz, vr, vz
 
         @loop_s is begin
             @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-                pdf_charged_out[ivpa,ivperp,iz,ir,is] += dt*Source_i[ivpa,ivperp,iz,ir]
+                pdf_ion_out[ivpa,ivperp,iz,ir,is] += dt*Source_i[ivpa,ivperp,iz,ir]
             end
         end
 
@@ -248,7 +248,7 @@ function source_terms_manufactured!(pdf_charged_out, pdf_neutral_out, vz, vr, vz
 
         @loop_s is begin
             @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-                pdf_charged_out[ivpa,ivperp,iz,ir,is] += dt*Source_i_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],t)
+                pdf_ion_out[ivpa,ivperp,iz,ir,is] += dt*Source_i_func(vpa.grid[ivpa],vperp.grid[ivperp],z.grid[iz],r.grid[ir],t)
             end
         end
 
diff --git a/src/time_advance.jl b/src/time_advance.jl
index c2e84cfea..501d3519f 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -179,7 +179,7 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
     # that may be evolved separately via fluid equations
-    scratch = setup_scratch_arrays(moments, pdf.charged.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    scratch = setup_scratch_arrays(moments, pdf.ion.norm, pdf.neutral.norm, t_input.n_rk_stages)
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
@@ -196,7 +196,7 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     end
 
     ##
-    # Charged particle advection only
+    # ion particle advection only
     ##
 
     # create structure r_advect whose members are the arrays needed to compute
@@ -213,8 +213,8 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     # initialise the z advection speed
     begin_s_r_vperp_vpa_region()
     @loop_s is begin
-        @views update_speed_z!(z_advect[is], moments.charged.upar[:,:,is],
-                               moments.charged.vth[:,:,is], moments.evolve_upar,
+        @views update_speed_z!(z_advect[is], moments.ion.upar[:,:,is],
+                               moments.ion.vth[:,:,is], moments.evolve_upar,
                                moments.evolve_ppar, fields, vpa, vperp, z, r, 0.0,
                                geometry)
     end
@@ -300,33 +300,33 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     if !restarting
         begin_serial_region()
         # ensure initial pdf has no negative values
-        force_minimum_pdf_value!(pdf.charged.norm, num_diss_params)
+        force_minimum_pdf_value!(pdf.ion.norm, num_diss_params)
         force_minimum_pdf_value_neutral!(pdf.neutral.norm, num_diss_params)
         # enforce boundary conditions and moment constraints to ensure a consistent initial
         # condition
         enforce_boundary_conditions!(
-            pdf.charged.norm, boundary_distributions.pdf_rboundary_charged,
-            moments.charged.dens, moments.charged.upar, moments.charged.ppar, moments,
+            pdf.ion.norm, boundary_distributions.pdf_rboundary_ion,
+            moments.ion.dens, moments.ion.upar, moments.ion.ppar, moments,
             vpa.bc, z.bc, r.bc, vpa, vperp, z, r, vpa_advect, z_advect, r_advect,
             composition, scratch_dummy, advance.r_diffusion, advance.vpa_diffusion)
         # Ensure normalised pdf exactly obeys integral constraints if evolving moments
         begin_s_r_z_region()
         @loop_s_r_z is ir iz begin
-            @views hard_force_moment_constraints!(pdf.charged.norm[:,:,iz,ir,is], moments,
+            @views hard_force_moment_constraints!(pdf.ion.norm[:,:,iz,ir,is], moments,
                                                   vpa)
         end
         # update moments in case they were affected by applying boundary conditions or
         # constraints to the pdf
         reset_moments_status!(moments)
-        update_moments!(moments, pdf.charged.norm, vpa, vperp, z, r, composition)
+        update_moments!(moments, pdf.ion.norm, vpa, vperp, z, r, composition)
         # enforce boundary conditions in r and z on the neutral particle distribution function
         if n_neutral_species > 0
             # Note, so far vr and vzeta do not need advect objects, so pass `nothing` for
             # those as a placeholder
             enforce_neutral_boundary_conditions!(
-                pdf.neutral.norm, pdf.charged.norm, boundary_distributions,
+                pdf.neutral.norm, pdf.ion.norm, boundary_distributions,
                 moments.neutral.dens, moments.neutral.uz, moments.neutral.pz, moments,
-                moments.charged.dens, moments.charged.upar, fields.Er, neutral_r_advect,
+                moments.ion.dens, moments.ion.upar, fields.Er, neutral_r_advect,
                 neutral_z_advect, nothing, nothing, neutral_vz_advect, r, z, vzeta, vr,
                 vz, composition, geometry, scratch_dummy, advance.r_diffusion,
                 advance.vz_diffusion)
@@ -343,10 +343,10 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
         # or constraints to the pdf
         begin_s_r_z_region()
         @loop_s_r_z is ir iz begin
-            scratch[1].pdf[:,:,iz,ir,is] .= pdf.charged.norm[:,:,iz,ir,is]
-            scratch[1].density[iz,ir,is] = moments.charged.dens[iz,ir,is]
-            scratch[1].upar[iz,ir,is] = moments.charged.upar[iz,ir,is]
-            scratch[1].ppar[iz,ir,is] = moments.charged.ppar[iz,ir,is]
+            scratch[1].pdf[:,:,iz,ir,is] .= pdf.ion.norm[:,:,iz,ir,is]
+            scratch[1].density[iz,ir,is] = moments.ion.dens[iz,ir,is]
+            scratch[1].upar[iz,ir,is] = moments.ion.upar[iz,ir,is]
+            scratch[1].ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
         end
 
         begin_sn_r_z_region(no_synchronize=true)
@@ -619,12 +619,12 @@ end
 create an array of structs containing scratch arrays for the normalised pdf and low-order moments
 that may be evolved separately via fluid equations
 """
-function setup_scratch_arrays(moments, pdf_charged_in, pdf_neutral_in, n_rk_stages)
+function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
     # create n_rk_stages+1 structs, each of which will contain one pdf,
     # one density, and one parallel flow array
     scratch = Vector{scratch_pdf{5,3,6,3}}(undef, n_rk_stages+1)
-    pdf_dims = size(pdf_charged_in)
-    moment_dims = size(moments.charged.dens)
+    pdf_dims = size(pdf_ion_in)
+    moment_dims = size(moments.ion.dens)
     pdf_neutral_dims = size(pdf_neutral_in)
     moment_neutral_dims = size(moments.neutral.dens)
     # populate each of the structs
@@ -648,10 +648,10 @@ function setup_scratch_arrays(moments, pdf_charged_in, pdf_neutral_in, n_rk_stag
                                       pdf_neutral_array, density_neutral_array,
                                       uz_neutral_array, pz_neutral_array)
         @serial_region begin
-            scratch[istage].pdf .= pdf_charged_in
-            scratch[istage].density .= moments.charged.dens
-            scratch[istage].upar .= moments.charged.upar
-            scratch[istage].ppar .= moments.charged.ppar
+            scratch[istage].pdf .= pdf_ion_in
+            scratch[istage].density .= moments.ion.dens
+            scratch[istage].upar .= moments.ion.upar
+            scratch[istage].ppar .= moments.ion.ppar
 
             scratch[istage].pdf_neutral .= pdf_neutral_in
             scratch[istage].density_neutral .= moments.neutral.dens
@@ -768,14 +768,14 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     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.charged.norm[:,1,:,1,is],
+                        heatmap(z.grid, vpa.grid, pdf.ion.norm[:,1,:,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_ion_species]
                     for (is, p) in enumerate(f_plots)
-                        @views draw_v_parallel_zero!(p, z.grid, moments.charged.upar[:,1,is],
-                                                     moments.charged.vth[:,1,is],
+                        @views draw_v_parallel_zero!(p, z.grid, moments.ion.upar[:,1,is],
+                                                     moments.ion.vth[:,1,is],
                                                      moments.evolve_upar,
                                                      moments.evolve_ppar)
                     end
@@ -792,14 +792,14 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                                                      moments.evolve_ppar)
                     end
                     logf_plots = [
-                        heatmap(z.grid, vpa.grid, log.(abs.(pdf.charged.norm[:,1,:,1,is])),
+                        heatmap(z.grid, vpa.grid, log.(abs.(pdf.ion.norm[:,1,:,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_neutral_species]
                     for (is, p) in enumerate(logf_plots)
-                        @views draw_v_parallel_zero!(p, z.grid, moments.charged.upar[:,1,is],
-                                                     moments.charged.vth[:,1,is],
+                        @views draw_v_parallel_zero!(p, z.grid, moments.ion.upar[:,1,is],
+                                                     moments.ion.vth[:,1,is],
                                                      moments.evolve_upar,
                                                      moments.evolve_ppar)
                     end
@@ -816,43 +816,43 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                                                      moments.evolve_ppar)
                     end
                     f0_plots = [
-                        plot(vpa.grid, pdf.charged.norm[:,1,1,1,is], xlabel="vpa", ylabel="f0", legend=false)
+                        plot(vpa.grid, pdf.ion.norm[:,1,1,1,is], xlabel="vpa", ylabel="f0", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     f0_neutral_plots = [
                         plot(vz.grid, pdf.neutral.norm[:,1,1,1,1,isn], xlabel="vz", ylabel="f0_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     fL_plots = [
-                        plot(vpa.grid, pdf.charged.norm[:,1,end,1,is], xlabel="vpa", ylabel="fL", legend=false)
+                        plot(vpa.grid, pdf.ion.norm[:,1,end,1,is], xlabel="vpa", ylabel="fL", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     fL_neutral_plots = [
                         plot(vz.grid, pdf.neutral.norm[:,1,1,end,1,isn], xlabel="vz", ylabel="fL_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     density_plots = [
-                        plot(z.grid, moments.charged.dens[:,1,is], xlabel="z", ylabel="density", legend=false)
+                        plot(z.grid, moments.ion.dens[:,1,is], xlabel="z", ylabel="density", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     density_neutral_plots = [
                         plot(z.grid, moments.neutral.dens[:,1,isn], xlabel="z", ylabel="density_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     upar_plots = [
-                        plot(z.grid, moments.charged.upar[:,1,is], xlabel="z", ylabel="upar", legend=false)
+                        plot(z.grid, moments.ion.upar[:,1,is], xlabel="z", ylabel="upar", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     upar_neutral_plots = [
                         plot(z.grid, moments.neutral.uz[:,1,isn], xlabel="z", ylabel="uz_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     ppar_plots = [
-                        plot(z.grid, moments.charged.ppar[:,1,is], xlabel="z", ylabel="ppar", legend=false)
+                        plot(z.grid, moments.ion.ppar[:,1,is], xlabel="z", ylabel="ppar", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     ppar_neutral_plots = [
                         plot(z.grid, moments.neutral.pz[:,1,isn], xlabel="z", ylabel="pz_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     vth_plots = [
-                        plot(z.grid, moments.charged.vth[:,1,is], xlabel="z", ylabel="vth", legend=false)
+                        plot(z.grid, moments.ion.vth[:,1,is], xlabel="z", ylabel="vth", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     vth_neutral_plots = [
                         plot(z.grid, moments.neutral.vth[:,1,isn], xlabel="z", ylabel="vth_neutral", legend=false)
                         for isn ∈ 1:composition.n_neutral_species]
                     qpar_plots = [
-                        plot(z.grid, moments.charged.qpar[:,1,is], xlabel="z", ylabel="qpar", legend=false)
+                        plot(z.grid, moments.ion.qpar[:,1,is], xlabel="z", ylabel="qpar", legend=false)
                         for is ∈ 1:composition.n_ion_species]
                     qpar_neutral_plots = [
                         plot(z.grid, moments.neutral.qz[:,1,isn], xlabel="z", ylabel="qz_neutral", legend=false)
@@ -889,7 +889,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                                    Dates.format(now(), dateformat"H:MM:SS"))
                 end
             end
-            write_dfns_data_to_binary(pdf.charged.norm, pdf.neutral.norm, moments, fields,
+            write_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
                                       t, composition.n_ion_species,
                                       composition.n_neutral_species, io_dfns, iwrite_dfns,
                                       r, z, vperp, vpa, vzeta, vr, vz)
@@ -920,14 +920,14 @@ function time_advance_split_operators!(pdf, scratch, t, t_input, vpa, z,
     flipflop = (mod(istep,2)==0)
     if flipflop
         # advance the operator-split 1D advection equation in vpa
-        # vpa-advection only applies for charged species
+        # vpa-advection only applies for ion species
         advance.vpa_advection = true
         time_advance_no_splitting!(pdf, scratch, t, t_input, vpa, z,
             vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
             composition, collisions, num_diss_params, advance, istep)
         advance.vpa_advection = false
         # z_advection! advances the operator-split 1D advection equation in z
-        # apply z-advection operation to all species (charged and neutral)
+        # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
         time_advance_no_splitting!(pdf, scratch, t, t_input, vpa, z,
             vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
@@ -1039,14 +1039,14 @@ function time_advance_split_operators!(pdf, scratch, t, t_input, vpa, z,
             end
         end
         # z_advection! advances the operator-split 1D advection equation in z
-        # apply z-advection operation to all species (charged and neutral)
+        # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
         time_advance_no_splitting!(pdf, scratch, t, t_input, vpa, z,
             vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
             composition, collisions, num_diss_params, advance, istep)
         advance.z_advection = false
         # advance the operator-split 1D advection equation in vpa
-        # vpa-advection only applies for charged species
+        # vpa-advection only applies for ion species
         advance.vpa_advection = true
         time_advance_no_splitting!(pdf, scratch, t, t_input, vpa, z,
             vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
@@ -1101,11 +1101,11 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
     ##
-    # update the charged particle distribution and moments
+    # update the ion particle distribution and moments
     ##
 
     @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        new_scratch.pdf[ivpa,ivperp,iz,ir,is] = rk_coefs[1]*pdf.charged.norm[ivpa,ivperp,iz,ir,is] + rk_coefs[2]*old_scratch.pdf[ivpa,ivperp,iz,ir,is] + rk_coefs[3]*new_scratch.pdf[ivpa,ivperp,iz,ir,is]
+        new_scratch.pdf[ivpa,ivperp,iz,ir,is] = rk_coefs[1]*pdf.ion.norm[ivpa,ivperp,iz,ir,is] + rk_coefs[2]*old_scratch.pdf[ivpa,ivperp,iz,ir,is] + rk_coefs[3]*new_scratch.pdf[ivpa,ivperp,iz,ir,is]
     end
     # use Runge Kutta to update any velocity moments evolved separately from the pdf
     rk_update_evolved_moments!(new_scratch, old_scratch, moments, rk_coefs)
@@ -1123,7 +1123,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     # contribution from one or more of the terms.
     # NB: probably need to do the same for the evolved moments
     enforce_boundary_conditions!(new_scratch, moments,
-        boundary_distributions.pdf_rboundary_charged, vpa.bc, z.bc, r.bc, vpa, vperp, z,
+        boundary_distributions.pdf_rboundary_ion, vpa.bc, z.bc, r.bc, vpa, vperp, z,
         r, vpa_advect, z_advect, r_advect, composition, scratch_dummy,
         advance.r_diffusion, advance.vpa_diffusion)
 
@@ -1141,7 +1141,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     begin_s_r_z_region()
     try #below block causes DomainError if ppar < 0 or density, so exit cleanly if possible
         @loop_s_r_z is ir iz begin
-            moments.charged.vth[iz,ir,is] = sqrt(2.0*new_scratch.ppar[iz,ir,is]/new_scratch.density[iz,ir,is])
+            moments.ion.vth[iz,ir,is] = sqrt(2.0*new_scratch.ppar[iz,ir,is]/new_scratch.density[iz,ir,is])
         end
     catch e
         if global_size[] > 1
@@ -1155,8 +1155,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         rethrow(e)
     end
     # update the parallel heat flux
-    update_qpar!(moments.charged.qpar, moments.charged.qpar_updated, new_scratch.density,
-                 new_scratch.upar, moments.charged.vth, new_scratch.pdf, vpa, vperp, z, r,
+    update_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
+                 new_scratch.upar, moments.ion.vth, new_scratch.pdf, vpa, vperp, z, r,
                  composition, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
 
@@ -1239,26 +1239,26 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 end
 
 """
-use Runge Kutta to update any charged-particle velocity moments evolved separately from
+use Runge Kutta to update any ion-particle velocity moments evolved separately from
 the pdf
 """
 function rk_update_evolved_moments!(new_scratch, old_scratch, moments, rk_coefs)
     # if separately evolving the particle density, update using RK
     if moments.evolve_density
         @loop_s_r_z is ir iz begin
-            new_scratch.density[iz,ir,is] = rk_coefs[1]*moments.charged.dens[iz,ir,is] + rk_coefs[2]*old_scratch.density[iz,ir,is] + rk_coefs[3]*new_scratch.density[iz,ir,is]
+            new_scratch.density[iz,ir,is] = rk_coefs[1]*moments.ion.dens[iz,ir,is] + rk_coefs[2]*old_scratch.density[iz,ir,is] + rk_coefs[3]*new_scratch.density[iz,ir,is]
         end
     end
     # if separately evolving the parallel flow, update using RK
     if moments.evolve_upar
         @loop_s_r_z is ir iz begin
-            new_scratch.upar[iz,ir,is] = rk_coefs[1]*moments.charged.upar[iz,ir,is] + rk_coefs[2]*old_scratch.upar[iz,ir,is] + rk_coefs[3]*new_scratch.upar[iz,ir,is]
+            new_scratch.upar[iz,ir,is] = rk_coefs[1]*moments.ion.upar[iz,ir,is] + rk_coefs[2]*old_scratch.upar[iz,ir,is] + rk_coefs[3]*new_scratch.upar[iz,ir,is]
         end
     end
     # if separately evolving the parallel pressure, update using RK;
     if moments.evolve_ppar
         @loop_s_r_z is ir iz begin
-            new_scratch.ppar[iz,ir,is] = rk_coefs[1]*moments.charged.ppar[iz,ir,is] + rk_coefs[2]*old_scratch.ppar[iz,ir,is] + rk_coefs[3]*new_scratch.ppar[iz,ir,is]
+            new_scratch.ppar[iz,ir,is] = rk_coefs[1]*moments.ion.ppar[iz,ir,is] + rk_coefs[2]*old_scratch.ppar[iz,ir,is] + rk_coefs[3]*new_scratch.ppar[iz,ir,is]
         end
     end
 end
@@ -1289,21 +1289,21 @@ function rk_update_evolved_moments_neutral!(new_scratch, old_scratch, moments, r
 end
 
 """
-update velocity moments that are calculable from the evolved charged pdf
+update velocity moments that are calculable from the evolved ion pdf
 """
 function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition)
     if !moments.evolve_density
-        update_density!(new_scratch.density, moments.charged.dens_updated,
+        update_density!(new_scratch.density, moments.ion.dens_updated,
                         new_scratch.pdf, vpa, vperp, z, r, composition)
     end
     if !moments.evolve_upar
-        update_upar!(new_scratch.upar, moments.charged.upar_updated, new_scratch.density,
+        update_upar!(new_scratch.upar, moments.ion.upar_updated, new_scratch.density,
                      new_scratch.ppar, new_scratch.pdf, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_ppar)
     end
     if !moments.evolve_ppar
         # update_ppar! calculates (p_parallel/m_s N_e c_s^2) + (n_s/N_e)*(upar_s/c_s)^2 = (1/√π)∫d(vpa/c_s) (vpa/c_s)^2 * (√π f_s c_s / N_e)
-        update_ppar!(new_scratch.ppar, moments.charged.ppar_updated, new_scratch.density,
+        update_ppar!(new_scratch.ppar, moments.ion.ppar_updated, new_scratch.density,
                      new_scratch.upar, new_scratch.pdf, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar)
     end
@@ -1345,12 +1345,12 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
 
     first_scratch = scratch[1]
     @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        first_scratch.pdf[ivpa,ivperp,iz,ir,is] = pdf.charged.norm[ivpa,ivperp,iz,ir,is]
+        first_scratch.pdf[ivpa,ivperp,iz,ir,is] = pdf.ion.norm[ivpa,ivperp,iz,ir,is]
     end
     @loop_s_r_z is ir iz begin
-        first_scratch.density[iz,ir,is] = moments.charged.dens[iz,ir,is]
-        first_scratch.upar[iz,ir,is] = moments.charged.upar[iz,ir,is]
-        first_scratch.ppar[iz,ir,is] = moments.charged.ppar[iz,ir,is]
+        first_scratch.density[iz,ir,is] = moments.ion.dens[iz,ir,is]
+        first_scratch.upar[iz,ir,is] = moments.ion.upar[iz,ir,is]
+        first_scratch.ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
     end
 
     if composition.n_neutral_species > 0
@@ -1395,12 +1395,12 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
     begin_s_r_z_region()
     final_scratch = scratch[istage]
     @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        pdf.charged.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
+        pdf.ion.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
     end
     @loop_s_r_z is ir iz begin
-        moments.charged.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
-        moments.charged.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
-        moments.charged.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
+        moments.ion.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
+        moments.ion.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
+        moments.ion.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
     end
     if composition.n_neutral_species > 0
         # No need to synchronize here as we only change neutral quantities and previous
@@ -1473,7 +1473,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     n_neutral_species = composition.n_neutral_species
     dt = t_input.dt
     # vpa_advection! advances the 1D advection equation in vpa.
-    # only charged species have a force accelerating them in vpa;
+    # only ion species have a force accelerating them in vpa;
     # however, neutral species do have non-zero d(wpa)/dt, so there is advection in wpa
 
     vpa_spectral, r_spectral, z_spectral = spectral_objects.vpa_spectral, spectral_objects.r_spectral, spectral_objects.z_spectral
@@ -1487,7 +1487,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
 
     # z_advection! advances 1D advection equation in z
-    # apply z-advection operation to charged species
+    # apply z-advection operation to ion species
 
     if advance.z_advection
         z_advection!(fvec_out.pdf, fvec_in, moments, fields, z_advect, z, vpa, vperp, r,
@@ -1535,9 +1535,9 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
 
     if advance.cx_collisions || advance.ionization_collisions
-        # gyroaverage neutral dfn and place it in the charged.buffer array for use in the collisions step
-        vzvrvzeta_to_vpavperp!(pdf.charged.buffer, fvec_in.pdf_neutral, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, geometry, composition)
-        # interpolate charged particle dfn and place it in the neutral.buffer array for use in the collisions step
+        # gyroaverage neutral dfn and place it in the ion.buffer array for use in the collisions step
+        vzvrvzeta_to_vpavperp!(pdf.ion.buffer, fvec_in.pdf_neutral, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, geometry, composition)
+        # interpolate ion particle dfn and place it in the neutral.buffer array for use in the collisions step
         vpavperp_to_vzvrvzeta!(pdf.neutral.buffer, fvec_in.pdf, vz, vr, vzeta, vpa, vperp, z, r, geometry, composition)
     end
 
@@ -1548,7 +1548,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                        collisions.charge_exchange, vpa_spectral,
                                        vz_spectral, dt)
     elseif advance.cx_collisions
-        charge_exchange_collisions_3V!(fvec_out.pdf, fvec_out.pdf_neutral, pdf.charged.buffer, pdf.neutral.buffer, fvec_in, composition,
+        charge_exchange_collisions_3V!(fvec_out.pdf, fvec_out.pdf_neutral, pdf.ion.buffer, pdf.neutral.buffer, fvec_in, composition,
                                         vz, vr, vzeta, vpa, vperp, z, r, collisions.charge_exchange, dt)
     end
     # account for ionization collisions between ions and neutrals
@@ -1557,7 +1557,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                   vperp, z, r, vz_spectral, moments, composition,
                                   collisions, dt)
     elseif advance.ionization_collisions
-        ionization_collisions_3V!(fvec_out.pdf, fvec_out.pdf_neutral, pdf.charged.buffer, fvec_in, composition,
+        ionization_collisions_3V!(fvec_out.pdf, fvec_out.pdf_neutral, pdf.ion.buffer, fvec_in, composition,
                                         vz, vr, vzeta, vpa, vperp, z, r, collisions, dt)
     end
     if advance.ionization_source
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index ae3fe8c29..b47deed70 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -12,7 +12,7 @@ export update_upar!
 export update_ppar!
 export update_qpar!
 export reset_moments_status!
-export moments_chrg_substruct, moments_ntrl_substruct
+export moments_ion_substruct, moments_neutral_substruct
 export update_neutral_density!
 export update_neutral_uz!
 export update_neutral_ur!
@@ -36,7 +36,7 @@ using ..looping
 
 """
 """
-struct moments_charged_substruct
+struct moments_ion_substruct
     # this is the particle density
     dens::MPISharedArray{mk_float,3}
     # flag that keeps track of if the density needs updating before use
@@ -165,7 +165,7 @@ end
 
 """
 """
-function create_moments_charged(nz, nr, n_species, evolve_density, evolve_upar,
+function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
                                 evolve_ppar, numerical_dissipation)
     # allocate array used for the particle density
     density = allocate_shared_float(nz, nr, n_species)
@@ -246,7 +246,7 @@ function create_moments_charged(nz, nr, n_species, evolve_density, evolve_upar,
     end
     
     # return struct containing arrays needed to update moments
-    return moments_charged_substruct(density, density_updated, parallel_flow,
+    return moments_ion_substruct(density, density_updated, parallel_flow,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, v_norm_fac,
         ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2, dppar_dz,
@@ -354,45 +354,45 @@ calculate the updated density (dens) and parallel pressure (ppar) for all specie
 function update_moments!(moments, ff, vpa, vperp, z, r, composition)
     begin_s_r_z_region()
     n_species = size(ff,5)
-    @boundscheck n_species == size(moments.charged.dens,3) || throw(BoundsError(moments))
+    @boundscheck n_species == size(moments.ion.dens,3) || throw(BoundsError(moments))
     @loop_s is begin
-        if moments.charged.dens_updated[is] == false
-            @views update_density_species!(moments.charged.dens[:,:,is], ff[:,:,:,:,is],
+        if moments.ion.dens_updated[is] == false
+            @views update_density_species!(moments.ion.dens[:,:,is], ff[:,:,:,:,is],
                                            vpa, vperp, z, r)
-            moments.charged.dens_updated[is] = true
+            moments.ion.dens_updated[is] = true
         end
-        if moments.charged.upar_updated[is] == false
+        if moments.ion.upar_updated[is] == false
             # Can pass moments.ppar here even though it has not been updated yet,
             # because moments.ppar is only needed if evolve_ppar=true, in which case it
             # will not be updated because it is not calculated from the distribution
             # function
-            @views update_upar_species!(moments.charged.upar[:,:,is],
-                                        moments.charged.dens[:,:,is],
-                                        moments.charged.ppar[:,:,is], ff[:,:,:,:,is], vpa,
+            @views update_upar_species!(moments.ion.upar[:,:,is],
+                                        moments.ion.dens[:,:,is],
+                                        moments.ion.ppar[:,:,is], ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
                                         moments.evolve_ppar)
-            moments.charged.upar_updated[is] = true
+            moments.ion.upar_updated[is] = true
         end
-        if moments.charged.ppar_updated[is] == false
-            @views update_ppar_species!(moments.charged.ppar[:,:,is],
-                                        moments.charged.dens[:,:,is],
-                                        moments.charged.upar[:,:,is], ff[:,:,:,:,is], vpa,
+        if moments.ion.ppar_updated[is] == false
+            @views update_ppar_species!(moments.ion.ppar[:,:,is],
+                                        moments.ion.dens[:,:,is],
+                                        moments.ion.upar[:,:,is], ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
                                         moments.evolve_upar)
-            moments.charged.ppar_updated[is] = true
+            moments.ion.ppar_updated[is] = true
         end
         @loop_r_z ir iz begin
-            moments.charged.vth[iz,ir,is] =
-                sqrt(2*moments.charged.ppar[iz,ir,is]/moments.charged.dens[iz,ir,is])
-        end
-        if moments.charged.qpar_updated[is] == false
-            @views update_qpar_species!(moments.charged.qpar[:,:,is],
-                                        moments.charged.dens[:,:,is],
-                                        moments.charged.upar[:,:,is],
-                                        moments.charged.vth[:,:,is], ff[:,:,:,:,is], vpa,
+            moments.ion.vth[iz,ir,is] =
+                sqrt(2*moments.ion.ppar[iz,ir,is]/moments.ion.dens[iz,ir,is])
+        end
+        if moments.ion.qpar_updated[is] == false
+            @views update_qpar_species!(moments.ion.qpar[:,:,is],
+                                        moments.ion.dens[:,:,is],
+                                        moments.ion.upar[:,:,is],
+                                        moments.ion.vth[:,:,is], ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
                                         moments.evolve_upar, moments.evolve_ppar)
-            moments.charged.qpar_updated[is] = true
+            moments.ion.qpar_updated[is] = true
         end
     end
     return nothing
@@ -645,8 +645,8 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
         density = @view scratch.density[:,:,is]
         upar = @view scratch.upar[:,:,is]
         ppar = @view scratch.ppar[:,:,is]
-        qpar = @view moments.charged.qpar[:,:,is]
-        vth = @view moments.charged.vth[:,:,is]
+        qpar = @view moments.ion.qpar[:,:,is]
+        vth = @view moments.ion.vth[:,:,is]
         dummy_zr = @view scratch_dummy.dummy_zrs[:,:,is]
         buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,is]
         buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,is]
@@ -659,7 +659,7 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -upar[iz,ir]
             end
-            @views derivative_z!(moments.charged.ddens_dz_upwind[:,:,is], density,
+            @views derivative_z!(moments.ion.ddens_dz_upwind[:,:,is], density,
                                  dummy_zr, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                                  buffer_r_5, buffer_r_6, z_spectral, z)
         end
@@ -669,11 +669,11 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             # centred second derivative for dissipation
             @views derivative_z!(dummy_zr, density, buffer_r_1, buffer_r_2, buffer_r_3,
                                  buffer_r_4, z_spectral, z)
-            @views derivative_z!(moments.charged.d2dens_dz2[:,:,is], dummy_zr, buffer_r_1,
+            @views derivative_z!(moments.ion.d2dens_dz2[:,:,is], dummy_zr, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
-            @views derivative_z!(moments.charged.dupar_dz[:,:,is], upar, buffer_r_1,
+            @views derivative_z!(moments.ion.dupar_dz[:,:,is], upar, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_upar
@@ -682,7 +682,7 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -upar[iz,ir]
             end
-            @views derivative_z!(moments.charged.dupar_dz_upwind[:,:,is], upar, dummy_zr,
+            @views derivative_z!(moments.ion.dupar_dz_upwind[:,:,is], upar, dummy_zr,
                                  buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                                  buffer_r_5, buffer_r_6, z_spectral, z)
         end
@@ -692,11 +692,11 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             # centred second derivative for dissipation
             @views derivative_z!(dummy_zr, upar, buffer_r_1, buffer_r_2, buffer_r_3,
                                  buffer_r_4, z_spectral, z)
-            @views derivative_z!(moments.charged.d2upar_dz2[:,:,is], dummy_zr, buffer_r_1,
+            @views derivative_z!(moments.ion.d2upar_dz2[:,:,is], dummy_zr, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_upar
-            @views derivative_z!(moments.charged.dppar_dz[:,:,is], ppar, buffer_r_1,
+            @views derivative_z!(moments.ion.dppar_dz[:,:,is], ppar, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_ppar
@@ -704,19 +704,19 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -upar[iz,ir]
             end
-            @views derivative_z!(moments.charged.dppar_dz_upwind[:,:,is], ppar, dummy_zr,
+            @views derivative_z!(moments.ion.dppar_dz_upwind[:,:,is], ppar, dummy_zr,
                                  buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                                  buffer_r_5, buffer_r_6, z_spectral, z)
 
             # centred second derivative for dissipation
             @views derivative_z!(dummy_zr, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
                                  buffer_r_4, z_spectral, z)
-            @views derivative_z!(moments.charged.d2ppar_dz2[:,:,is], dummy_zr, buffer_r_1,
+            @views derivative_z!(moments.ion.d2ppar_dz2[:,:,is], dummy_zr, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
-            @views derivative_z!(moments.charged.dqpar_dz[:,:,is], qpar, buffer_r_1,
+            @views derivative_z!(moments.ion.dqpar_dz[:,:,is], qpar, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-            @views derivative_z!(moments.charged.dvth_dz[:,:,is], vth, buffer_r_1,
+            @views derivative_z!(moments.ion.dvth_dz[:,:,is], vth, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
     end
@@ -1458,18 +1458,18 @@ end
 """
 function reset_moments_status!(moments)
     if moments.evolve_density == false
-        moments.charged.dens_updated .= false
+        moments.ion.dens_updated .= false
         moments.neutral.dens_updated .= false
     end
     if moments.evolve_upar == false
-        moments.charged.upar_updated .= false
+        moments.ion.upar_updated .= false
         moments.neutral.uz_updated .= false
     end
     if moments.evolve_ppar == false
-        moments.charged.ppar_updated .= false
+        moments.ion.ppar_updated .= false
         moments.neutral.pz_updated .= false
     end
-    moments.charged.qpar_updated .= false
+    moments.ion.qpar_updated .= false
     moments.neutral.uzeta_updated .= false
     moments.neutral.ur_updated .= false
     moments.neutral.pzeta_updated .= false
diff --git a/src/vpa_advection.jl b/src/vpa_advection.jl
index 00c0726d5..052b77d21 100644
--- a/src/vpa_advection.jl
+++ b/src/vpa_advection.jl
@@ -99,9 +99,9 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
             # • -wpar^2 * d(vth)/dz term
             @loop_z_vperp iz ivperp begin
                 @views @. advect[is].speed[:,ivperp,iz,ir] =
-                    moments.charged.dppar_dz[iz,ir,is]/(fvec.density[iz,ir,is]*moments.charged.vth[iz,ir,is]) +
-                    0.5*vpa.grid*moments.charged.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] -
-                    vpa.grid^2*moments.charged.dvth_dz[iz,ir,is]
+                    moments.ion.dppar_dz[iz,ir,is]/(fvec.density[iz,ir,is]*moments.ion.vth[iz,ir,is]) +
+                    0.5*vpa.grid*moments.ion.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] -
+                    vpa.grid^2*moments.ion.dvth_dz[iz,ir,is]
             end
         end
     end
@@ -120,7 +120,7 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
                           * (fvec.upar[iz,ir,is]-fvec.uz_neutral[iz,ir,is])^2)
                      - fvec.density_neutral[iz,ir,is]
                        * (fvec.uz_neutral[iz,ir,is]-fvec.upar[iz,ir,is])
-                       / moments.charged.vth[iz,ir,is]) +
+                       / moments.ion.vth[iz,ir,is]) +
                     collisions.ionization *
                     (0.5*vpa.grid
                        * (fvec.density_neutral[iz,ir,is]
@@ -131,7 +131,7 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
                             / fvec.ppar[iz,ir,is])
                      - fvec.density_neutral[iz,ir,is]
                        * (fvec.uz_neutral[iz,ir,is] - fvec.upar[iz,ir,is])
-                       / moments.charged.vth[iz,ir,is])
+                       / moments.ion.vth[iz,ir,is])
             end
         end
     end
@@ -153,9 +153,9 @@ function update_speed_n_p_evolution!(advect, fields, fvec, moments, vpa, z, r, c
             # • vpahat*d(upar)/dz
             # • -(1/2)*(dphi/dz)/vthi
             @loop_z_vperp iz ivperp begin
-                @views @. advect[is].speed[:,ivperp,iz,ir] = 0.5*vpa.grid*moments.charged.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] +
-                                                             vpa.grid*moments.charged.dvth_dz[iz] * (fvec.upar[iz,ir,is]/moments.vth[iz,ir,is] - vpa.grid) +
-                                                             vpa.grid*moments.charged.dupar_dz[iz,ir,is] +
+                @views @. advect[is].speed[:,ivperp,iz,ir] = 0.5*vpa.grid*moments.ion.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] +
+                                                             vpa.grid*moments.ion.dvth_dz[iz] * (fvec.upar[iz,ir,is]/moments.vth[iz,ir,is] - vpa.grid) +
+                                                             vpa.grid*moments.ion.dupar_dz[iz,ir,is] +
                                                              0.5*fields.Ez[iz,ir]/moments.vth[iz,ir,is]
             end
         end
@@ -190,8 +190,8 @@ function update_speed_n_u_evolution!(advect, fvec, moments, vpa, z, r, compositi
             # • -wpar*dupar/dz
             @loop_z_vperp iz ivperp begin
                 @views @. advect[is].speed[:,ivperp,iz,ir] =
-                    moments.charged.dppar_dz[iz,ir,is]/fvec.density[iz,ir,is] -
-                    vpa.grid*moments.charged.dupar_dz[iz,ir,is]
+                    moments.ion.dppar_dz[iz,ir,is]/fvec.density[iz,ir,is] -
+                    vpa.grid*moments.ion.dupar_dz[iz,ir,is]
             end
         end
     end
diff --git a/src/z_advection.jl b/src/z_advection.jl
index f67996009..8a367478d 100644
--- a/src/z_advection.jl
+++ b/src/z_advection.jl
@@ -21,13 +21,13 @@ function z_advection!(f_out, fvec_in, moments, fields, advect, z, vpa, vperp, r,
     @loop_s is begin
         # get the updated speed along the z direction using the current f
         @views update_speed_z!(advect[is], fvec_in.upar[:,:,is],
-                               moments.charged.vth[:,:,is], moments.evolve_upar,
+                               moments.ion.vth[:,:,is], moments.evolve_upar,
                                moments.evolve_ppar, fields, vpa, vperp, z, r, t, geometry)
         # update adv_fac
         @loop_r_vperp_vpa ir ivperp ivpa begin
             @views adjust_advection_speed!(advect[is].speed[:,ivpa,ivperp,ir],
                                            fvec_in.density[:,ir,is],
-                                           moments.charged.vth[:,ir,is],
+                                           moments.ion.vth[:,ir,is],
                                            moments.evolve_density, moments.evolve_ppar)
             @views @. advect[is].adv_fac[:,ivpa,ivperp,ir] = -dt*advect[is].speed[:,ivpa,ivperp,ir]
             # take the normalized pdf contained in fvec_in.pdf and remove the normalization,
@@ -35,7 +35,7 @@ function z_advection!(f_out, fvec_in, moments, fields, advect, z, vpa, vperp, r,
             @views unnormalize_pdf!(
                 scratch_dummy.buffer_vpavperpzrs_2[ivpa,ivperp,:,ir,is],
                 fvec_in.pdf[ivpa,ivperp,:,ir,is], fvec_in.density[:,ir,is],
-                moments.charged.vth[:,ir,is], moments.evolve_density, moments.evolve_ppar)
+                moments.ion.vth[:,ir,is], moments.evolve_density, moments.evolve_ppar)
         end
     end
     #calculate the upwind derivative
diff --git a/test/nonlinear_sound_wave_tests.jl b/test/nonlinear_sound_wave_tests.jl
index 2742760cd..9a926fbcf 100644
--- a/test/nonlinear_sound_wave_tests.jl
+++ b/test/nonlinear_sound_wave_tests.jl
@@ -9,7 +9,7 @@ using TimerOutputs
 using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data, load_species_data,
-                                 load_fields_data, load_charged_particle_moments_data, load_pdf_data,
+                                 load_fields_data, load_ion_particle_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data, load_neutral_pdf_data, load_time_data,
                                  load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -31,13 +31,13 @@ struct expected_data
     z::Array{mk_float, 1}
     vpa::Array{mk_float, 1}
     phi::Array{mk_float, 2}
-    n_charged::Array{mk_float, 2}
+    n_ion::Array{mk_float, 2}
     n_neutral::Array{mk_float, 2}
-    upar_charged::Array{mk_float, 2}
+    upar_ion::Array{mk_float, 2}
     upar_neutral::Array{mk_float, 2}
-    ppar_charged::Array{mk_float, 2}
+    ppar_ion::Array{mk_float, 2}
     ppar_neutral::Array{mk_float, 2}
-    f_charged::Array{mk_float, 3}
+    f_ion::Array{mk_float, 3}
     f_neutral::Array{mk_float, 3}
 end
 
@@ -57,7 +57,7 @@ const expected =
     -0.35341444947230544 -0.37562068472777577; -0.5495322983936355 -0.5902920278548919;
     -0.8609384185348539 -0.8726873701297669; -1.2115106026686981 -1.130685855316896;
     -1.3862820803244256 -1.2383045504753758],
-   # Expected n_charged:
+   # Expected n_ion:
    [0.2500030702177184 0.2898752704335188; 0.2977471383217195 0.3227988953227183;
     0.42274614626845974 0.417842206578383; 0.5772539714051019 0.5541536351162784;
     0.702254450621661 0.686868306132489; 0.7499999999999392 0.7467716863438243;
@@ -71,7 +71,7 @@ const expected =
     0.2977471383217197 0.30537567265010435; 0.4227461462684595 0.4096819689829924;
     0.5772539714051017 0.5583199869826102; 0.7022544506216611 0.7056147469533546;
     0.7499999999999394 0.7737996616909211],
-   # Expected upar_charged:
+   # Expected upar_ion:
    [1.1971912119126474e-17 -2.0968470015869656e-16;
     -5.818706134342973e-17 -0.18232119131671534;
     9.895531571141618e-17 -0.1967239995126128;
@@ -93,7 +93,7 @@ const expected =
     2.2213638810498713e-17 0.009226347310827925;
     -2.7413075842225616e-17 0.036181844730955835;
     -3.143783114880993e-17 -2.810175715538666e-17],
-   # Expected ppar_charged:
+   # Expected ppar_ion:
    [0.18749999999999997 0.23278940073547755; 0.20909100943488423 0.21912527958959363;
     0.24403280042122125 0.20817795270356831; 0.2440328004212212 0.21516422119834766;
     0.20909100943488412 0.2208129180125869; 0.18750000000000003 0.2213757117801786;
@@ -107,7 +107,7 @@ const expected =
     0.2090910094348842 0.19265693636741688; 0.24403280042122116 0.20584407392704462;
     0.2440328004212211 0.228696297221881; 0.2090910094348841 0.2438447730062432;
     0.18750000000000003 0.24825654204434672],
-   # Expected f_charged:
+   # Expected f_ion:
    [0.03704623609948259 0.04056128509273146 0.04289169811317835 0.030368915327672292 0.01235362235033934 0.0063385294703834204 0.012353622350339327 0.030368915327672247 0.04289169811317828 0.04056128509273145 0.0370462360994826;
     0.20411991941198782 0.251156132910555 0.3935556226209418 0.6276758497903185 0.9100827333021343 1.06066017177965 0.9100827333021342 0.6276758497903192 0.3935556226209421 0.25115613291055494 0.2041199194119877;
     0.03704623609948259 0.04056128509273146 0.04289169811317835 0.030368915327672292 0.01235362235033934 0.0063385294703834204 0.012353622350339327 0.030368915327672247 0.04289169811317828 0.04056128509273145 0.0370462360994826;;;
@@ -260,10 +260,10 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
 
     # Suppress console output while running
     phi = undef
-    n_charged = undef
-    upar_charged = undef
-    ppar_charged = undef
-    f_charged = undef
+    n_ion = undef
+    upar_ion = undef
+    ppar_ion = undef
+    f_ion = undef
     n_neutral = undef
     upar_neutral = undef
     ppar_neutral = undef
@@ -293,7 +293,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, qpar_charged_zrst, v_t_charged_zrst = load_charged_particle_moments_data(fid)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_particle_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
             z, z_spectral = load_coordinate_data(fid, "z")
 
@@ -303,18 +303,18 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             fid = open_readonly_output_file(path, "dfns")
             
             # load particle distribution function (pdf) data
-            f_charged_vpavperpzrst = load_pdf_data(fid)
+            f_ion_vpavperpzrst = load_pdf_data(fid)
             f_neutral_vzvrvzetazrst = load_neutral_pdf_data(fid)
 
             close(fid)
             
             phi = phi_zrt[:,1,:]
-            n_charged = n_charged_zrst[:,1,:,:]
-            upar_charged = upar_charged_zrst[:,1,:,:]
-            ppar_charged = ppar_charged_zrst[:,1,:,:]
-            qpar_charged = qpar_charged_zrst[:,1,:,:]
-            v_t_charged = v_t_charged_zrst[:,1,:,:]
-            f_charged = f_charged_vpavperpzrst[:,1,:,1,:,:]
+            n_ion = n_ion_zrst[:,1,:,:]
+            upar_ion = upar_ion_zrst[:,1,:,:]
+            ppar_ion = ppar_ion_zrst[:,1,:,:]
+            qpar_ion = qpar_ion_zrst[:,1,:,:]
+            v_t_ion = v_t_ion_zrst[:,1,:,:]
+            f_ion = f_ion_vpavperpzrst[:,1,:,1,:,:]
             n_neutral = n_neutral_zrst[:,1,:,:]
             upar_neutral = upar_neutral_zrst[:,1,:,:]
             ppar_neutral = ppar_neutral_zrst[:,1,:,:]
@@ -325,7 +325,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # Unnormalize f
             if input["evolve_moments_density"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] .*= n_charged[iz,is,it]
+                    f_ion[:,iz,is,it] .*= n_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] .*= n_neutral[iz,isn,it]
@@ -333,7 +333,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             end
             if input["evolve_moments_parallel_pressure"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] ./= v_t_charged[iz,is,it]
+                    f_ion[:,iz,is,it] ./= v_t_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] ./= v_t_neutral[iz,isn,it]
@@ -371,11 +371,11 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
         #println("phi ", size(newgrid_phi))
         #println(newgrid_phi)
         #println()
-        #newgrid_n_charged = cat(interpolate_to_grid_z(expected.z, n_charged[:, :, 1], z, z_spectral)[:,1],
-        #                        interpolate_to_grid_z(expected.z, n_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_n_ion = cat(interpolate_to_grid_z(expected.z, n_ion[:, :, 1], z, z_spectral)[:,1],
+        #                        interpolate_to_grid_z(expected.z, n_ion[:, :, 2], z, z_spectral)[:,1];
         #                        dims=2)
-        #println("n_charged ", size(newgrid_n_charged))
-        #println(newgrid_n_charged)
+        #println("n_ion ", size(newgrid_n_ion))
+        #println(newgrid_n_ion)
         #println()
         #newgrid_n_neutral = cat(interpolate_to_grid_z(expected.z, n_neutral[:, :, 1], z, z_spectral)[:,1],
         #                        interpolate_to_grid_z(expected.z, n_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -383,11 +383,11 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
         #println("n_neutral ", size(newgrid_n_neutral))
         #println(newgrid_n_neutral)
         #println()
-        #newgrid_upar_charged = cat(interpolate_to_grid_z(expected.z, upar_charged[:, :, 1], z, z_spectral)[:,1],
-        #                           interpolate_to_grid_z(expected.z, upar_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_upar_ion = cat(interpolate_to_grid_z(expected.z, upar_ion[:, :, 1], z, z_spectral)[:,1],
+        #                           interpolate_to_grid_z(expected.z, upar_ion[:, :, 2], z, z_spectral)[:,1];
         #                           dims=2)
-        #println("upar_charged ", size(newgrid_upar_charged))
-        #println(newgrid_upar_charged)
+        #println("upar_ion ", size(newgrid_upar_ion))
+        #println(newgrid_upar_ion)
         #println()
         #newgrid_upar_neutral = cat(interpolate_to_grid_z(expected.z, upar_neutral[:, :, 1], z, z_spectral)[:,1],
         #                           interpolate_to_grid_z(expected.z, upar_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -395,11 +395,11 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
         #println("upar_neutral ", size(newgrid_upar_neutral))
         #println(newgrid_upar_neutral)
         #println()
-        #newgrid_ppar_charged = cat(interpolate_to_grid_z(expected.z, ppar_charged[:, :, 1], z, z_spectral)[:,1],
-        #                           interpolate_to_grid_z(expected.z, ppar_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_ppar_ion = cat(interpolate_to_grid_z(expected.z, ppar_ion[:, :, 1], z, z_spectral)[:,1],
+        #                           interpolate_to_grid_z(expected.z, ppar_ion[:, :, 2], z, z_spectral)[:,1];
         #                           dims=2)
-        #println("ppar_charged ", size(newgrid_ppar_charged))
-        #println(newgrid_ppar_charged)
+        #println("ppar_ion ", size(newgrid_ppar_ion))
+        #println(newgrid_ppar_ion)
         #println()
         #newgrid_ppar_neutral = cat(interpolate_to_grid_z(expected.z, ppar_neutral[:, :, 1], z, z_spectral)[:,1],
         #                           interpolate_to_grid_z(expected.z, ppar_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -407,11 +407,11 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
         #println("ppar_neutral ", size(newgrid_ppar_neutral))
         #println(newgrid_ppar_neutral)
         #println()
-        #newgrid_f_charged = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_charged[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
-        #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_charged[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
+        #newgrid_f_ion = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_ion[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
+        #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_ion[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
         #                        dims=4)
-        #println("f_charged ", size(newgrid_f_charged))
-        #println(newgrid_f_charged)
+        #println("f_ion ", size(newgrid_f_ion))
+        #println(newgrid_f_ion)
         #println()
         #newgrid_f_neutral = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_neutral[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
         #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_neutral[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
@@ -424,21 +424,21 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
                 newgrid_phi = interpolate_to_grid_z(expected.z, phi[:, tind], z, z_spectral)
                 @test isapprox(expected.phi[:, tind], newgrid_phi, rtol=rtol)
 
-                # Check charged particle moments and f
+                # Check ion particle moments and f
                 ######################################
 
-                newgrid_n_charged = interpolate_to_grid_z(expected.z, n_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.n_charged[:, tind], newgrid_n_charged[:,1], rtol=rtol)
+                newgrid_n_ion = interpolate_to_grid_z(expected.z, n_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.n_ion[:, tind], newgrid_n_ion[:,1], rtol=rtol)
 
-                newgrid_upar_charged = interpolate_to_grid_z(expected.z, upar_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.upar_charged[:, tind], newgrid_upar_charged[:,1], rtol=upar_rtol, atol=atol)
+                newgrid_upar_ion = interpolate_to_grid_z(expected.z, upar_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.upar_ion[:, tind], newgrid_upar_ion[:,1], rtol=upar_rtol, atol=atol)
 
-                newgrid_ppar_charged = interpolate_to_grid_z(expected.z, ppar_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.ppar_charged[:, tind], newgrid_ppar_charged[:,1], rtol=rtol)
+                newgrid_ppar_ion = interpolate_to_grid_z(expected.z, ppar_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.ppar_ion[:, tind], newgrid_ppar_ion[:,1], rtol=rtol)
 
-                newgrid_f_charged = interpolate_to_grid_z(expected.z, f_charged[:, :, :, tind], z, z_spectral)
-                newgrid_f_charged = interpolate_to_grid_vpa(expected.vpa, newgrid_f_charged, vpa, vpa_spectral)
-                @test isapprox(expected.f_charged[:, :, tind], newgrid_f_charged[:,:,1], rtol=rtol)
+                newgrid_f_ion = interpolate_to_grid_z(expected.z, f_ion[:, :, :, tind], z, z_spectral)
+                newgrid_f_ion = interpolate_to_grid_vpa(expected.vpa, newgrid_f_ion, vpa, vpa_spectral)
+                @test isapprox(expected.f_ion[:, :, tind], newgrid_f_ion[:,:,1], rtol=rtol)
 
                 # Check neutral particle moments and f
                 ######################################

From b84283e9383590cd8c606fed40f3a2fcb0876fe2 Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Wed, 7 Jun 2023 16:33:22 +0100
Subject: [PATCH 032/428] 1) Added electrons to the moments struct and added
 calculation of the initial electron density, parallel flow, thermal speed and
 parallel pressure. 2) Density is obtained via summing over ion densities,
 flow by using charge conservation (which I have added a module for), initial
 thermal speed is set to one (I have not added the feature to evolve the
 electron energy equation), and parallel pressure is calculated
 self-consistently from these other quantities. 3) All of the electron moments
 info has also been added to the binary output, though I have yet to modify
 the post-processing module to plot it.

All of the test from the standard test suite pass.
---
 src/charge_conservation.jl   |  38 ++++++++++
 src/file_io.jl               |  48 +++++++++++-
 src/initial_conditions.jl    |  87 +++++++++++++++++++++-
 src/moment_kinetics.jl       |  11 ++-
 src/moment_kinetics_input.jl |  20 ++++-
 src/time_advance.jl          |   9 +--
 src/velocity_moments.jl      | 139 ++++++++++++++++++++++++++++++++++-
 7 files changed, 333 insertions(+), 19 deletions(-)
 create mode 100644 src/charge_conservation.jl

diff --git a/src/charge_conservation.jl b/src/charge_conservation.jl
new file mode 100644
index 000000000..d91bb45d5
--- /dev/null
+++ b/src/charge_conservation.jl
@@ -0,0 +1,38 @@
+module charge_conservation
+
+export calculate_electron_upar_from_charge_conservation!
+
+"""
+use charge conservation equation to solve for the electron parallel flow density:
+    d/dz(sum_i n_i upar_i - n_e upar_e) = 0
+    ==> [sum_i n_i upar_i](z) - [sum_i n_i upar_i](zbound) = [n_e upar_e](z) - [n_e upar_e](zbound)
+inputs: 
+    upar_e - should contain updated electron parallel flow density at boundaries in zed
+    dens_e - electron particle density
+    upar_i - ion parallel flow density
+    dens_i - ion particle density
+output:
+    upar_e - contains the updated electron parallel flow density
+"""
+function calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i)
+    nr = size(upar_e, 2)
+    nz = size(upar_e, 1)
+    ns = size(upar_i, 3)
+    for ir in 1:nr
+        boundary_flux = dens_e[1,ir] * upar_e[1,ir]
+        for iz in 2:nz-1
+            # calculate the boundary value for the particle flux, and initialise 
+            # the electron particle flux to it
+            upar_e[iz,ir] = boundary_flux
+            # add the contributions to the electron particle flux from the various ion species
+            # particle fluxes
+            for is in 1:ns
+                upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
+            end
+            # convert from parallel particle flux to parallel particle density
+            upar_e[iz,ir] /= dens_e[iz,ir]
+        end
+    end
+end
+
+end
\ No newline at end of file
diff --git a/src/file_io.jl b/src/file_io.jl
index 10b212f90..a4065e435 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -5,6 +5,7 @@ module file_io
 export input_option_error
 export open_output_file, open_ascii_output_file
 export setup_file_io, finish_file_io
+export write_moments_data_to_binary
 export write_data_to_ascii
 export write_data_to_netcdf, write_data_to_hdf5
 
@@ -59,6 +60,15 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn}
     # handle for the ion species thermal speed
     thermal_speed::Tmomi
 
+    # handle for the electron species density
+    electron_density::Tphi
+    # handle for the electron species parallel flow
+    electron_parallel_flow::Tphi
+    # handle for the electron species parallel pressure
+    electron_parallel_pressure::Tphi
+    # handle for the electron species thermal speed
+    electron_thermal_speed::Tphi
+
     # handle for the neutral species density
     density_neutral::Tmomn
     uz_neutral::Tmomn
@@ -116,6 +126,7 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
         if io_input.ascii_output
             #ff_io = open_ascii_output_file(out_prefix, "f_vs_t")
             mom_ion_io = open_ascii_output_file(out_prefix, "moments_ion_vs_t")
+            mom_eon_io = open_ascii_output_file(out_prefix, "moments_electron_vs_t")
             mom_ntrl_io = open_ascii_output_file(out_prefix, "moments_neutral_vs_t")
             fields_io = open_ascii_output_file(out_prefix, "fields_vs_t")
             ascii = ascii_ios(mom_ion_io, mom_ntrl_io, fields_io)
@@ -483,6 +494,30 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                           description="ion species thermal speed",
                                           units="c_ref")
 
+        # io_density is the handle for the ion particle density
+        io_electron_density = create_dynamic_variable!(dynamic, "electron_density", mk_float, z, r;
+                                              parallel_io=parallel_io,
+                                              description="electron species density",
+                                              units="n_ref")
+
+        # io_electron_upar is the handle for the electron parallel flow density
+        io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
+                                           parallel_io=parallel_io,
+                                           description="electron species parallel flow",
+                                           units="c_ref = sqrt(2*T_ref/mi)")
+
+        # io_electron_ppar is the handle for the electron parallel pressure
+        io_electron_ppar = create_dynamic_variable!(dynamic, "electron_parallel_pressure", mk_float, z, r;
+                                           parallel_io=parallel_io,
+                                           description="electron species parallel pressure",
+                                           units="n_ref*T_ref")
+
+        # io_electron_vth is the handle for the electron thermal speed
+        io_electron_vth = create_dynamic_variable!(dynamic, "electron_thermal_speed", mk_float, z, r;
+                                          parallel_io=parallel_io,
+                                          description="electron species thermal speed",
+                                          units="c_ref")
+
         # io_density_neutral is the handle for the neutral particle density
         io_density_neutral = create_dynamic_variable!(dynamic, "density_neutral", mk_float, z, r;
                                                       n_neutral_species=n_neutral_species,
@@ -518,9 +553,10 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
             parallel_io=parallel_io, description="neutral species thermal speed",
             units="c_ref")
 
-        return io_moments_info(fid, io_time, io_phi, io_Er, io_Ez, io_density, io_upar,
-                               io_ppar, io_qpar, io_vth, io_density_neutral, io_uz_neutral,
-                               io_pz_neutral, io_qz_neutral, io_thermal_speed_neutral,
+        return io_moments_info(fid, io_time, io_phi, io_Er, io_Ez, 
+                               io_density, io_upar, io_ppar, io_qpar, io_vth, 
+                               io_electron_density, io_electron_upar, io_electron_ppar, io_electron_vth,
+                               io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral, io_thermal_speed_neutral,
                                parallel_io)
     end
 
@@ -701,6 +737,12 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
                               z, r, n_ion_species)
         append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth, t_idx, z, r,
                               n_ion_species)
+        # add the electron velocity-moments data at this time slice to the output file
+        append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar, t_idx, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_pressure, moments.electron.ppar, t_idx, z, r)
+        append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth, t_idx, z, r)
+        # add the neutral velocity-moments data at this time slice to the output file
         if n_neutral_species > 0
             append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
                                   z, r, n_neutral_species)
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index bf1145ba5..f8d7cd06c 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -27,11 +27,13 @@ using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
-using ..velocity_moments: create_moments_ion, create_moments_neutral, update_qpar!
-using ..velocity_moments: moments_ion_substruct, moments_neutral_substruct
+using ..velocity_moments: create_moments_ion, create_moments_electron, create_moments_neutral
+using ..velocity_moments: update_qpar!
+using ..velocity_moments: moments_ion_substruct, moments_electron_substruct, moments_neutral_substruct
 using ..velocity_moments: update_neutral_density!, update_neutral_pz!, update_neutral_pr!, update_neutral_pzeta!
 using ..velocity_moments: update_neutral_uz!, update_neutral_ur!, update_neutral_uzeta!, update_neutral_qz!
 using ..velocity_moments: update_ppar!, update_upar!, update_density!
+using ..charge_conservation: calculate_electron_upar_from_charge_conservation!
 
 using ..manufactured_solns: manufactured_solutions
 
@@ -48,12 +50,15 @@ end
 struct pdf_struct
     #ion particles: s + r + z + vperp + vpa
     ion::pdf_substruct{5}
+    # electron particles: 1 + r + z + vperp + vpa
+    electron::pdf_substruct{5}
     #neutral particles: s + r + z + vzeta + vr + vz
     neutral::pdf_substruct{6}
 end
 
 struct moments_struct
     ion::moments_ion_substruct
+    electron::moments_electron_substruct
     neutral::moments_neutral_substruct
     # flag that indicates if the density should be evolved via continuity equation
     evolve_density::Bool
@@ -88,11 +93,14 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
     # create the 'moments' struct that contains various v-space moments and other
     # information related to these moments.
     # the time-dependent entries are not initialised.
-    # moments arrays have same r and z grids for both ion and neutral species
+    # moments arrays have same r and z grids for ion, electron and neutral species
     # and so are included in the same struct
     ion = create_moments_ion(z.n, r.n, composition.n_ion_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, numerical_dissipation)
+    electron = create_moments_electron(z.n, r.n,
+        evolve_moments.density, evolve_moments.parallel_flow,
+        evolve_moments.parallel_pressure, numerical_dissipation)
     neutral = create_moments_neutral(z.n, r.n, composition.n_neutral_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, numerical_dissipation)
@@ -106,7 +114,7 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
         particle_number_conserved = true
     end
 
-    moments = moments_struct(ion, neutral, evolve_moments.density,
+    moments = moments_struct(ion, electron, neutral, evolve_moments.density,
                              particle_number_conserved,
                              evolve_moments.conservation,
                              evolve_moments.parallel_flow,
@@ -127,8 +135,13 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     pdf_ion_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
+    pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, 1)
+    # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
+    # to emulate ion and neutral behaviour
+    pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, 1)
 
     return pdf_struct(pdf_substruct(pdf_ion_norm, pdf_ion_buffer),
+                      pdf_substruct(pdf_electron_norm, pdf_electron_buffer),
                       pdf_substruct(pdf_neutral_norm, pdf_neutral_buffer))
 
 end
@@ -166,10 +179,23 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                 @. moments.neutral.pz = 0.5 * moments.neutral.dens * moments.neutral.vth^2
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
             end
+
+            # initialise the electron density profile
+            init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
+            # initialise the electron parallel flow profile
+            init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
+                moments.ion.upar, n_ion_species, r.n, composition.electron_physics)
+            # initialise the electron thermal speed profile
+            init_electron_vth!(moments.electron.vth)
+            # calculate the electron parallel pressure from the density and thermal speed
+            @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
         end
         moments.ion.dens_updated .= true
         moments.ion.upar_updated .= true
         moments.ion.ppar_updated .= true
+        moments.electron.dens_updated = true
+        moments.electron.upar_updated = true
+        moments.electron.ppar_updated = true
         moments.neutral.dens_updated .= true
         moments.neutral.uz_updated .= true
         moments.neutral.pz_updated .= true
@@ -393,6 +419,59 @@ function init_uzeta!(uzeta, z, r, spec, n_species)
     return nothing
 end
 
+"""
+initialise the electron density
+"""
+function init_electron_density!(electron_density, ion_density, n_ion_species)
+    # use quasineutrality to obtain the electron density from the initial
+    # densities of the various ion species
+    for is ∈ 1:n_ion_species
+        @loop_r_z ir iz begin
+             electron_density[iz,ir] += ion_density[iz,ir,is]
+        end
+    end
+    return nothing
+end
+
+"""
+initialise the electron parallel flow density
+"""
+function init_electron_upar!(electron_upar, electron_density, ion_upar, ion_density, n_ion_species, nr, electron_model)
+    # initialise the electron parallel flow density
+    electron_upar .= 0.0
+    # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
+    # is equal and opposite to the ion parallel current
+    if electron_model == "boltzmann_electron_response_with_simple_sheath"
+        # loop over ion species, adding each species contribution to the
+        # ion parallel particle flux at the boundaries in zed
+        for is ∈ 1:n_ion_species
+            for ir ∈ 1:nr
+                # electron_upar at this intermediate stage is actually
+                # the electron particle flux
+                electron_upar[1,ir] += ion_dens[1,ir,is] * ion_upar[1,ir,is]
+                electron_upar[end,ir] += ion_dens[end,ir,is] * ion_upar[end,ir,is]
+            end
+        end
+        # convert from electron particle flux to electron parallel flow density
+        electron_upar[1,ir] /= electron_dens[1,ir]
+        electron_upar[end,ir] /= electron_dens[end,ir]
+        # use charge conservation to solve for the electron upar for the rest of the zed domain
+        calculate_electron_upar_from_charge_conservation!(electron_upar, electron_dens, ion_upar, ion_dens)
+    end
+    return nothing
+end
+
+"""
+initialise the electron thermal speed profile.
+for now the only initialisation option for the temperature is constant in z.
+returns vth0 = sqrt(Ts/Te) = 1.0
+"""
+function init_electron_vth!(vth)
+    @loop_s_r_z is ir iz begin
+        vth[iz,ir,is] = 1.0
+    end
+end
+
 """
 """
 function init_ion_pdf_over_density!(pdf, spec, composition, vpa, vperp, z,
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index dfaad2a3c..6ce681ad9 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -29,6 +29,7 @@ include("coordinates.jl")
 include("file_io.jl")
 include("velocity_moments.jl")
 include("velocity_grid_transforms.jl")
+include("charge_conservation.jl")
 include("em_fields.jl")
 include("bgk.jl")
 include("manufactured_solns.jl") # MRH Here?
@@ -79,6 +80,7 @@ using .looping: debug_setup_loop_ranges_split_one_combination!
 using .moment_kinetics_input: mk_input, read_input_file, run_type, performance_test
 using .time_advance: setup_time_advance!, time_advance!
 using .type_definitions: mk_int
+using .em_fields: setup_em_fields
 
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
@@ -356,6 +358,11 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
         allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
                                  evolve_moments, collisions, num_diss_params)
 
+    # create the "fields" structure that contains arrays
+    # for the electrostatic potential phi (and eventually the electromagnetic fields)
+    fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude, 
+                             drive_input.frequency, drive_input.force_Er_zero_at_wall)
+
     if restart_prefix_iblock === nothing
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
@@ -377,11 +384,11 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
-    moments, fields, spectral_objects, advect_objects,
+    moments, spectral_objects, advect_objects,
     scratch, advance, scratch_dummy, manufactured_source_list =
         setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral,
-            r_spectral, composition, drive_input, moments, t_input, collisions, species,
+            r_spectral, composition, drive_input, moments, fields, t_input, collisions, species,
             geometry, boundary_distributions, num_diss_params, restarting)
     # setup i/o
     ascii_io, io_moments, io_dfns = setup_file_io(io_input, boundary_distributions, vz,
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 7284c04b2..f4ef39521 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -414,7 +414,21 @@ function mk_input(scan_input=Dict())
                 species.neutral[is].initial_density, z_IC, vpa_IC)
         end
     end 
-    species_immutable = (ion = species_ion_immutable, neutral = species_neutral_immutable)
+    z_IC = initial_condition_input(species.electron.z_IC.initialization_option,
+        species.electron.z_IC.width, species.electron.z_IC.wavenumber,
+        species.electron.z_IC.density_amplitude, species.electron.z_IC.density_phase,
+        species.electron.z_IC.upar_amplitude, species.electron.z_IC.upar_phase,
+        species.electron.z_IC.temperature_amplitude, species.electron.z_IC.temperature_phase,
+        species.electron.z_IC.monomial_degree)
+    vpa_IC = initial_condition_input(species.electron.vpa_IC.initialization_option,
+        species.electron.vpa_IC.width, species.electron.vpa_IC.wavenumber,
+        species.electron.vpa_IC.density_amplitude, species.electron.vpa_IC.density_phase,
+        species.electron.vpa_IC.upar_amplitude, species.electron.vpa_IC.upar_phase,
+        species.electron.vpa_IC.temperature_amplitude,
+        species.electron.vpa_IC.temperature_phase, species.electron.vpa_IC.monomial_degree)
+    species_electron_immutable = species_parameters("electron", species.electron.initial_temperature,
+        species.electron.initial_density, z_IC, vpa_IC)
+    species_immutable = (ion = species_ion_immutable, electron = species_electron_immutable, neutral = species_neutral_immutable)
     
     force_Er_zero = get(scan_input, "force_Er_zero_at_wall", false)
     drive_immutable = drive_input(drive.force_phi, drive.amplitude, drive.frequency, force_Er_zero)
@@ -840,7 +854,9 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
                 initial_density, deepcopy(z_initial_conditions), deepcopy(vpa_initial_conditions))
         end
     end
-    species = (ion = species_ion, neutral = species_neutral)
+    species_electron = species_parameters_mutable("electron", T_e, 1.0, deepcopy(z_initial_conditions),
+        deepcopy(vpa_initial_conditions))
+    species = (ion = species_ion, electron = species_electron, neutral = species_neutral)
     
     # if drive_phi = true, include external electrostatic potential of form
     # phi(z,t=0)*drive_amplitude*sinpi(time*drive_frequency)
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 501d3519f..66f6c609f 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -159,7 +159,7 @@ EM fields, and advection terms
 function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
                              vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
                              z_spectral, r_spectral, composition, drive_input, moments,
-                             t_input, collisions, species, geometry,
+                             fields, t_input, collisions, species, geometry,
                              boundary_distributions, num_diss_params, restarting)
     # define some local variables for convenience/tidiness
     n_species = composition.n_species
@@ -184,9 +184,6 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
                                    composition.n_ion_species,n_neutral_species_alloc)
-    # create the "fields" structure that contains arrays
-    # for the electrostatic potential phi and eventually the electromagnetic fields
-    fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude, drive_input.frequency, drive_input.force_Er_zero_at_wall)
     # initialize the electrostatic potential
     begin_serial_region()
     update_phi!(fields, scratch[1], z, r, composition, z_spectral, r_spectral, scratch_dummy)
@@ -367,8 +364,8 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     # Ensure all processes are synchronized at the end of the setup
     _block_synchronize()
 
-    return moments, fields, spectral_objects, advect_objects,
-    scratch, advance, scratch_dummy, manufactured_source_list
+    return moments, spectral_objects, advect_objects, scratch, advance, 
+        scratch_dummy, manufactured_source_list
 end
 
 """
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index b47deed70..381bb08a6 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -5,14 +5,14 @@ module velocity_moments
 export integrate_over_vspace
 export integrate_over_positive_vpa, integrate_over_negative_vpa
 export integrate_over_positive_vz, integrate_over_negative_vz
-export create_moments_chrg, create_moments_ntrl
+export create_moments_ion, create_moments_electron, create_moments_neutral
 export update_moments!
 export update_density!
 export update_upar!
 export update_ppar!
 export update_qpar!
 export reset_moments_status!
-export moments_ion_substruct, moments_neutral_substruct
+export moments_ion_substruct, moments_electron_substruct, moments_neutral_substruct
 export update_neutral_density!
 export update_neutral_uz!
 export update_neutral_ur!
@@ -94,6 +94,55 @@ struct moments_ion_substruct
     dvth_dz::Union{MPISharedArray{mk_float,3},Nothing}
 end
 
+"""
+moments_electron_substruct is a struct that contains moment information for electrons
+"""
+mutable struct moments_electron_substruct
+    # this is the particle density
+    dens::MPISharedArray{mk_float,2}
+    # flag that keeps track of if the density needs updating before use
+    dens_updated::Bool
+    # this is the parallel flow
+    upar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not upar needs updating before use
+    upar_updated::Bool
+    # this is the parallel pressure
+    ppar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not ppar needs updating before use
+    ppar_updated::Bool
+    # this is the parallel heat flux
+    qpar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not qpar needs updating before use
+    qpar_updated::Bool
+    # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
+    vth::MPISharedArray{mk_float,2}
+    # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
+    # a factor of vth for each power of wpa in velocity space integrals.
+    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
+    # and it is one otherwise
+    v_norm_fac::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the particle density
+    ddens_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the second-z-derivative of the particle density
+    d2dens_dz2::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel flow
+    dupar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the parallel flow
+    dupar_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the second-z-derivative of the parallel flow
+    d2upar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel pressure
+    dppar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the parallel pressure
+    dppar_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the second-z-derivative of the parallel pressure
+    d2ppar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel heat flux
+    dqpar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
+    dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
+end
+
 """
 """
 struct moments_neutral_substruct
@@ -253,6 +302,92 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
         dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz)
 end
 
+"""
+create a moment struct containing information about the electron moments
+"""
+function create_moments_electron(nz, nr, evolve_density, evolve_upar,
+                                evolve_ppar, numerical_dissipation)
+    # allocate array used for the particle density
+    density = allocate_shared_float(nz, nr)
+    # initialise Bool variable that indicates if the density is updated for each species
+    density_updated = false
+    # allocate array used for the parallel flow
+    parallel_flow = allocate_shared_float(nz, nr)
+    # allocate Bool variable that indicates if the parallel flow is updated for each species
+    parallel_flow_updated = false
+    # allocate array used for the parallel pressure
+    parallel_pressure = allocate_shared_float(nz, nr)
+    # allocate Bool variable that indicates if the parallel pressure is updated for each species
+    parallel_pressure_updated = false
+    # allocate array used for the parallel flow
+    parallel_heat_flux = allocate_shared_float(nz, nr)
+    # allocate Bool variables that indicates if the parallel flow is updated for each species
+    parallel_heat_flux_updated = false
+    # allocate array used for the thermal speed
+    thermal_speed = allocate_shared_float(nz, nr)
+    if evolve_ppar
+        v_norm_fac = thermal_speed
+    else
+        v_norm_fac = allocate_shared_float(nz, nr)
+        @serial_region begin
+            v_norm_fac .= 1.0
+        end
+    end
+
+    if evolve_density
+        ddens_dz_upwind = allocate_shared_float(nz, nr)
+    else
+        ddens_dz_upwind = nothing
+    end
+    if evolve_density &&
+            numerical_dissipation.moment_dissipation_coefficient > 0.0
+
+        d2dens_dz2 = allocate_shared_float(nz, nr)
+    else
+        d2dens_dz2 = nothing
+    end
+    if evolve_density || evolve_upar || evolve_ppar
+        dupar_dz = allocate_shared_float(nz, nr)
+    else
+        dupar_dz = nothing
+    end
+    if evolve_upar
+        dupar_dz_upwind = allocate_shared_float(nz, nr)
+    else
+        dupar_dz_upwind = nothing
+    end
+    if evolve_upar &&
+            numerical_dissipation.moment_dissipation_coefficient > 0.0
+
+        d2upar_dz2 = allocate_shared_float(nz, nr)
+    else
+        d2upar_dz2 = nothing
+    end
+    if evolve_upar
+        dppar_dz = allocate_shared_float(nz, nr)
+    else
+        dppar_dz = nothing
+    end
+    if evolve_ppar
+        dppar_dz_upwind = allocate_shared_float(nz, nr)
+        d2ppar_dz2 = allocate_shared_float(nz, nr)
+        dqpar_dz = allocate_shared_float(nz, nr)
+        dvth_dz = allocate_shared_float(nz, nr)
+    else
+        dppar_dz_upwind = nothing
+        d2ppar_dz2 = nothing
+        dqpar_dz = nothing
+        dvth_dz = nothing
+    end
+    
+    # return struct containing arrays needed to update moments
+    return moments_electron_substruct(density, density_updated, parallel_flow,
+        parallel_flow_updated, parallel_pressure, parallel_pressure_updated,
+        parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, v_norm_fac,
+        ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2, dppar_dz,
+        dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz)
+end
+
 # neutral particles have natural mean velocities 
 # uz, ur, uzeta =/= upar 
 # and similarly for heat fluxes

From 228573575facd8e659928eeb3e972e8ce36db41a Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Thu, 8 Jun 2023 16:45:37 +0100
Subject: [PATCH 033/428] added electrons to the time advance loop in a passive
 way, where the electron density, upar, ppar, vth and qpar are kept updated
 without influencing the fields and thus the ion/neutral time advance. the
 usual test suite passes + an example with the simple sheath model behaves
 correctly (n_i = n_e and upar_i = upar_e)

---
 src/charge_conservation.jl      |  38 -----------
 src/electron_fluid_equations.jl | 101 +++++++++++++++++++++++++++
 src/file_io.jl                  |  44 +++++++++++-
 src/initial_conditions.jl       | 116 ++++++++++++++++++++------------
 src/moment_kinetics.jl          |   4 +-
 src/moment_kinetics_structs.jl  |   8 ++-
 src/post_processing.jl          |   2 +-
 src/time_advance.jl             |  70 +++++++++++++++++--
 8 files changed, 287 insertions(+), 96 deletions(-)
 delete mode 100644 src/charge_conservation.jl
 create mode 100644 src/electron_fluid_equations.jl

diff --git a/src/charge_conservation.jl b/src/charge_conservation.jl
deleted file mode 100644
index d91bb45d5..000000000
--- a/src/charge_conservation.jl
+++ /dev/null
@@ -1,38 +0,0 @@
-module charge_conservation
-
-export calculate_electron_upar_from_charge_conservation!
-
-"""
-use charge conservation equation to solve for the electron parallel flow density:
-    d/dz(sum_i n_i upar_i - n_e upar_e) = 0
-    ==> [sum_i n_i upar_i](z) - [sum_i n_i upar_i](zbound) = [n_e upar_e](z) - [n_e upar_e](zbound)
-inputs: 
-    upar_e - should contain updated electron parallel flow density at boundaries in zed
-    dens_e - electron particle density
-    upar_i - ion parallel flow density
-    dens_i - ion particle density
-output:
-    upar_e - contains the updated electron parallel flow density
-"""
-function calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i)
-    nr = size(upar_e, 2)
-    nz = size(upar_e, 1)
-    ns = size(upar_i, 3)
-    for ir in 1:nr
-        boundary_flux = dens_e[1,ir] * upar_e[1,ir]
-        for iz in 2:nz-1
-            # calculate the boundary value for the particle flux, and initialise 
-            # the electron particle flux to it
-            upar_e[iz,ir] = boundary_flux
-            # add the contributions to the electron particle flux from the various ion species
-            # particle fluxes
-            for is in 1:ns
-                upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
-            end
-            # convert from parallel particle flux to parallel particle density
-            upar_e[iz,ir] /= dens_e[iz,ir]
-        end
-    end
-end
-
-end
\ No newline at end of file
diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
new file mode 100644
index 000000000..cb61b8f11
--- /dev/null
+++ b/src/electron_fluid_equations.jl
@@ -0,0 +1,101 @@
+module electron_fluid_equations
+
+export calculate_electron_density!
+export calculate_electron_upar_from_charge_conservation!
+export calculate_electron_qpar!
+
+using ..looping
+using ..input_structs: boltzmann_electron_response_with_simple_sheath
+
+"""
+use quasineutrality to obtain the electron density from the initial
+densities of the various ion species:
+    sum_i dens_i = dens_e
+inputs:
+    dens_e - electron density at previous time level
+    dens_i - updated ion density
+output:
+    dens_e - updated electron density
+"""
+function calculate_electron_density!(dens_e, dens_i)
+    dens_e .= 0.0
+    @loop_s_r_z is ir iz begin
+        dens_e[iz,ir] += dens_i[iz,ir,is]
+    end
+    return nothing
+end
+
+"""
+use charge conservation equation to solve for the electron parallel flow density:
+    d/dz(sum_i n_i upar_i - n_e upar_e) = 0
+    ==> [sum_i n_i upar_i](z) - [sum_i n_i upar_i](zbound) = [n_e upar_e](z) - [n_e upar_e](zbound)
+inputs: 
+    upar_e - should contain updated electron parallel flow density at boundaries in zed
+    dens_e - electron particle density
+    upar_i - ion parallel flow density
+    dens_i - ion particle density
+output:
+    upar_e - contains the updated electron parallel flow density
+"""
+function calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i, electron_model)
+    # get the number of zed grid points, nz
+    nz = size(upar_e,1)
+    # initialise the electron parallel flow density to zero
+    upar_e .= 0.0
+    # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
+    # is equal and opposite to the ion parallel current
+    if electron_model == boltzmann_electron_response_with_simple_sheath
+        # loop over ion species, adding each species contribution to the
+        # ion parallel particle flux at the boundaries in zed
+        @loop_s_r is ir begin
+            # electron_upar at this intermediate stage is actually
+            # the electron particle flux
+            upar_e[1,ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
+            upar_e[end,ir] += dens_i[end,ir,is] * upar_i[end,ir,is]
+        end
+        @loop_r ir begin
+            # fix the boundary flux
+            boundary_flux = upar_e[1,ir]
+            for iz in 2:nz-1
+                # initialise the electron particle flux to its value at the boundary in zed
+                upar_e[iz,ir] = boundary_flux
+                # add the contributions to the electron particle flux from the various ion species
+                # particle fluxes
+                @loop_s is begin
+                    upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
+                end
+                # convert from parallel particle flux to parallel particle density
+                upar_e[iz,ir] /= dens_e[iz,ir]
+            end
+            # convert from parallel particle flux to parallel particle density for boundary points
+            upar_e[1,ir] /= dens_e[1,ir]
+            upar_e[end,ir] /= dens_e[end,ir]
+        end
+    end
+    return nothing
+end
+
+"""
+calculate the parallel component of the electron heat flux.
+there are currently two supported options for the parallel heat flux:
+    Braginskii collisional closure - qpar_e = -3.16*ppar_e/(m_e*nu_ei)*dT/dz - 0.71*ppar_e*(upar_i-upar_e)
+    collisionless closure - d(qpar_e)/dz = 0 ==> qpar_e = 0
+inputs:
+    qpar_e - parallel electron heat flux at the previous time level
+output:
+    qpar_e - updated parallel electron heat flux
+"""
+function calculate_electron_qpar!(qpar_e, electron_model)
+    if electron_model == "braginskii_fluid"
+        # NEED TO FIX!
+        qpar_e .= 0.0 
+    else
+        # if not using Braginskii fluid model, then assume for now
+        # that we are in the collisionless limit, for which d/dz(qpar_e) = 0.
+        # take qpar_e = 0
+        qpar_e .= 0.0
+    end
+    return nothing
+end
+
+end
\ No newline at end of file
diff --git a/src/file_io.jl b/src/file_io.jl
index a4065e435..f599a35dc 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -28,6 +28,7 @@ struct ascii_ios{T <: Union{IOStream,Nothing}}
     # corresponds to the ascii file to which velocity space moments of the
     # distribution function such as density and pressure are written
     moments_ion::T
+    moments_electron::T
     moments_neutral::T
     # corresponds to the ascii file to which electromagnetic fields
     # such as the electrostatic potential are written
@@ -66,6 +67,8 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn}
     electron_parallel_flow::Tphi
     # handle for the electron species parallel pressure
     electron_parallel_pressure::Tphi
+    # handle for the electron species parallel heat flux
+    electron_parallel_heat_flux::Tphi
     # handle for the electron species thermal speed
     electron_thermal_speed::Tphi
 
@@ -129,9 +132,9 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
             mom_eon_io = open_ascii_output_file(out_prefix, "moments_electron_vs_t")
             mom_ntrl_io = open_ascii_output_file(out_prefix, "moments_neutral_vs_t")
             fields_io = open_ascii_output_file(out_prefix, "fields_vs_t")
-            ascii = ascii_ios(mom_ion_io, mom_ntrl_io, fields_io)
+            ascii = ascii_ios(mom_ion_io, mom_eon_io, mom_ntrl_io, fields_io)
         else
-            ascii = ascii_ios(nothing, nothing, nothing)
+            ascii = ascii_ios(nothing, nothing, nothing, nothing)
         end
 
         io_moments = setup_moments_io(out_prefix, io_input.binary_format, r, z,
@@ -512,6 +515,12 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                            description="electron species parallel pressure",
                                            units="n_ref*T_ref")
 
+        # io_electron_qpar is the handle for the electron parallel heat flux
+        io_electron_qpar = create_dynamic_variable!(dynamic, "electron_parallel_heat_flux", mk_float, z, r;
+                                           parallel_io=parallel_io,
+                                           description="electron species parallel heat flux",
+                                           units="n_ref*T_ref*c_ref")
+
         # io_electron_vth is the handle for the electron thermal speed
         io_electron_vth = create_dynamic_variable!(dynamic, "electron_thermal_speed", mk_float, z, r;
                                           parallel_io=parallel_io,
@@ -555,7 +564,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
 
         return io_moments_info(fid, io_time, io_phi, io_Er, io_Ez, 
                                io_density, io_upar, io_ppar, io_qpar, io_vth, 
-                               io_electron_density, io_electron_upar, io_electron_ppar, io_electron_vth,
+                               io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar, io_electron_vth,
                                io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral, io_thermal_speed_neutral,
                                parallel_io)
     end
@@ -741,6 +750,7 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
         append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx, z, r)
         append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar, t_idx, z, r)
         append_to_dynamic_var(io_moments.electron_parallel_pressure, moments.electron.ppar, t_idx, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_heat_flux, moments.electron.qpar, t_idx, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth, t_idx, z, r)
         # add the neutral velocity-moments data at this time slice to the output file
         if n_neutral_species > 0
@@ -810,6 +820,12 @@ end
                                   moments.ion.qpar.data, t_idx, z, r, n_ion_species)
             append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth.data,
                                   t_idx, z, r, n_ion_species)
+            # add the electron velocity-moments data at this time slice to the output file
+            append_to_dynamic_var(io_moments.electron_density, moments.electron.dens.data, t_idx, z, r)
+            append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar.data, t_idx, z, r)
+            append_to_dynamic_var(io_moments.electron_parallel_pressure, moments.electron.ppar.data, t_idx, z, r)
+            append_to_dynamic_var(io_moments.electron_parallel_heat_flux, moments.electron.qpar.data, t_idx, z, r)
+            append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth.data, t_idx, z, r)
             if n_neutral_species > 0
                 append_to_dynamic_var(io_moments.density_neutral,
                                       moments.neutral.dens.data, t_idx, z, r,
@@ -901,6 +917,7 @@ function write_data_to_ascii(moments, fields, vpa, vperp, z, r, t, n_ion_species
 
         #write_f_ascii(ff, z, vpa, t, ascii_io.ff)
         write_moments_ion_ascii(moments.ion, z, r, t, n_ion_species, ascii_io.moments_ion)
+        write_moments_electron_ascii(moments.electron, z, r, t, ascii_io.moments_electron)
         if n_neutral_species > 0
             write_moments_neutral_ascii(moments.neutral, z, r, t, n_neutral_species, ascii_io.moments_neutral)
         end
@@ -957,6 +974,27 @@ function write_moments_ion_ascii(mom, z, r, t, n_species, ascii_io)
     return nothing
 end
 
+"""
+write moments of the ion species distribution function f at this time slice
+"""
+function write_moments_electron_ascii(mom, z, r, t, ascii_io)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+    
+        @inbounds begin
+            for ir ∈ 1:r.n
+                for iz ∈ 1:z.n
+                    println(ascii_io,"t: ", t, "   r: ", r.grid[ir], "   z: ", z.grid[iz],
+                            "  dens: ", mom.dens[iz,ir], "   upar: ", mom.upar[iz,ir],
+                            "   ppar: ", mom.ppar[iz,ir], "   qpar: ", mom.qpar[iz,ir])
+                end
+            end
+        end
+        println(ascii_io,"")
+    end
+    return nothing
+end
+
 """
 write moments of the neutral species distribution function f_neutral at this time slice
 """
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index f8d7cd06c..507a7dc6c 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -33,7 +33,10 @@ using ..velocity_moments: moments_ion_substruct, moments_electron_substruct, mom
 using ..velocity_moments: update_neutral_density!, update_neutral_pz!, update_neutral_pr!, update_neutral_pzeta!
 using ..velocity_moments: update_neutral_uz!, update_neutral_ur!, update_neutral_uzeta!, update_neutral_qz!
 using ..velocity_moments: update_ppar!, update_upar!, update_density!
-using ..charge_conservation: calculate_electron_upar_from_charge_conservation!
+using ..electron_fluid_equations: calculate_electron_density!
+using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
+using ..electron_fluid_equations: calculate_electron_qpar!
+using ..input_structs: boltzmann_electron_response_with_simple_sheath
 
 using ..manufactured_solns: manufactured_solutions
 
@@ -180,22 +183,24 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
             end
 
-            # initialise the electron density profile
-            init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
-            # initialise the electron parallel flow profile
-            init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
-                moments.ion.upar, n_ion_species, r.n, composition.electron_physics)
-            # initialise the electron thermal speed profile
-            init_electron_vth!(moments.electron.vth)
-            # calculate the electron parallel pressure from the density and thermal speed
-            @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
+            # # initialise the electron density profile
+            # init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
+
+            # println("ni: ", moments.ion.dens[1,1,1], " ne: ", moments.electron.dens[1,1])
+            # # initialise the electron parallel flow profile
+            # init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
+            #     moments.ion.upar, n_ion_species, r.n, composition.electron_physics)
+            # # initialise the electron thermal speed profile
+            # init_electron_vth!(moments.electron.vth)
+            # # calculate the electron parallel pressure from the density and thermal speed
+            # @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
         end
         moments.ion.dens_updated .= true
         moments.ion.upar_updated .= true
         moments.ion.ppar_updated .= true
-        moments.electron.dens_updated = true
-        moments.electron.upar_updated = true
-        moments.electron.ppar_updated = true
+        # moments.electron.dens_updated = true
+        # moments.electron.upar_updated = true
+        # moments.electron.ppar_updated = true
         moments.neutral.dens_updated .= true
         moments.neutral.uz_updated .= true
         moments.neutral.pz_updated .= true
@@ -206,12 +211,15 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         # when evolve_ppar = true.
         initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
                         vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
+
+        println("post_init_pdf_ni: ", moments.ion.dens[1,1,1])
         begin_s_r_z_region()
         # calculate the initial parallel heat flux from the initial un-normalised pdf
         update_qpar!(moments.ion.qpar, moments.ion.qpar_updated,
                      moments.ion.dens, moments.ion.upar, moments.ion.vth,
                      pdf.ion.norm, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
+        # calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
 
         if n_neutral_species > 0
             update_neutral_qz!(moments.neutral.qz, moments.neutral.qz_updated,
@@ -233,6 +241,26 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
+    @serial_region begin 
+        # initialise the electron density profile
+        init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
+        # initialise the electron parallel flow profile
+        init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
+            moments.ion.dens, composition.electron_physics)
+        # initialise the electron thermal speed profile
+        init_electron_vth!(moments.electron.vth, composition.T_e)
+        # calculate the electron parallel pressure from the density and thermal speed
+        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
+        # calculate the electron parallel heat flux
+        calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
+        # the electron moments have now been updated
+        moments.electron.dens_updated = true
+        moments.electron.upar_updated = true
+        moments.electron.ppar_updated = true
+    end
+    
+    println("upar_i: ", moments.ion.upar[1,1,1], " upar_e: ", moments.electron.upar[1,1])
+
     return nothing
 end
 
@@ -425,40 +453,40 @@ initialise the electron density
 function init_electron_density!(electron_density, ion_density, n_ion_species)
     # use quasineutrality to obtain the electron density from the initial
     # densities of the various ion species
-    for is ∈ 1:n_ion_species
-        @loop_r_z ir iz begin
-             electron_density[iz,ir] += ion_density[iz,ir,is]
-        end
-    end
+    calculate_electron_density!(electron_density, ion_density)
     return nothing
 end
 
 """
 initialise the electron parallel flow density
 """
-function init_electron_upar!(electron_upar, electron_density, ion_upar, ion_density, n_ion_species, nr, electron_model)
-    # initialise the electron parallel flow density
-    electron_upar .= 0.0
-    # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
-    # is equal and opposite to the ion parallel current
-    if electron_model == "boltzmann_electron_response_with_simple_sheath"
-        # loop over ion species, adding each species contribution to the
-        # ion parallel particle flux at the boundaries in zed
-        for is ∈ 1:n_ion_species
-            for ir ∈ 1:nr
-                # electron_upar at this intermediate stage is actually
-                # the electron particle flux
-                electron_upar[1,ir] += ion_dens[1,ir,is] * ion_upar[1,ir,is]
-                electron_upar[end,ir] += ion_dens[end,ir,is] * ion_upar[end,ir,is]
-            end
-        end
-        # convert from electron particle flux to electron parallel flow density
-        electron_upar[1,ir] /= electron_dens[1,ir]
-        electron_upar[end,ir] /= electron_dens[end,ir]
-        # use charge conservation to solve for the electron upar for the rest of the zed domain
-        calculate_electron_upar_from_charge_conservation!(electron_upar, electron_dens, ion_upar, ion_dens)
-    end
-    return nothing
+function init_electron_upar!(upar_e, dens_e, upar_i, dens_i, electron_model)
+    calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i, electron_model)
+    # # initialise the electron parallel flow density
+    # electron_upar .= 0.0
+    # println("electron_model: ", electron_model)
+    # # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
+    # # is equal and opposite to the ion parallel current
+    # if electron_model == boltzmann_electron_response_with_simple_sheath
+    #     # loop over ion species, adding each species contribution to the
+    #     # ion parallel particle flux at the boundaries in zed
+    #     for is ∈ 1:n_ion_species
+    #         for ir ∈ 1:nr
+    #             # electron_upar at this intermediate stage is actually
+    #             # the electron particle flux
+    #             electron_upar[1,ir] += ion_density[1,ir,is] * ion_upar[1,ir,is]
+    #             electron_upar[end,ir] += ion_density[end,ir,is] * ion_upar[end,ir,is]
+    #         end
+    #     end
+    #     # convert from electron particle flux to electron parallel flow density
+    #     for ir in 1:nr
+    #         electron_upar[1,ir] /= electron_density[1,ir]
+    #         electron_upar[end,ir] /= electron_density[end,ir]
+    #     end
+    #     # use charge conservation to solve for the electron upar for the rest of the zed domain
+    #     calculate_electron_upar_from_charge_conservation!(electron_upar, electron_density, ion_upar, ion_density)
+    # end
+    # return nothing
 end
 
 """
@@ -466,9 +494,9 @@ initialise the electron thermal speed profile.
 for now the only initialisation option for the temperature is constant in z.
 returns vth0 = sqrt(Ts/Te) = 1.0
 """
-function init_electron_vth!(vth)
-    @loop_s_r_z is ir iz begin
-        vth[iz,ir,is] = 1.0
+function init_electron_vth!(vth, T_e)
+    @loop_r_z ir iz begin
+        vth[iz,ir] = sqrt(T_e)
     end
 end
 
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 6ce681ad9..db0f3d017 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -29,7 +29,7 @@ include("coordinates.jl")
 include("file_io.jl")
 include("velocity_moments.jl")
 include("velocity_grid_transforms.jl")
-include("charge_conservation.jl")
+include("electron_fluid_equations.jl")
 include("em_fields.jl")
 include("bgk.jl")
 include("manufactured_solns.jl") # MRH Here?
@@ -381,6 +381,8 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
                                             composition, r, z, vpa, vperp, vzeta, vr, vz)
         _block_synchronize()
     end
+    println("dens0_ion: ", moments.ion.dens[1,1,1], " dens0_electron: ", moments.electron.dens[1,1,1])
+
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
diff --git a/src/moment_kinetics_structs.jl b/src/moment_kinetics_structs.jl
index 56829f815..17ea13e4d 100644
--- a/src/moment_kinetics_structs.jl
+++ b/src/moment_kinetics_structs.jl
@@ -10,13 +10,17 @@ using ..type_definitions: mk_float
 
 """
 """
-struct scratch_pdf{n_distribution_ion, n_moment, n_distribution_neutral,n_moment_neutral}
-    # ion particles
+struct scratch_pdf{n_distribution_ion, n_moment, n_moment_electron, n_distribution_neutral, n_moment_neutral}
+    # ions
     pdf::MPISharedArray{mk_float, n_distribution_ion}
     density::MPISharedArray{mk_float, n_moment}
     upar::MPISharedArray{mk_float, n_moment}
     ppar::MPISharedArray{mk_float, n_moment}
     temp_z_s::MPISharedArray{mk_float, n_moment}
+    # electrons
+    electron_density::MPISharedArray{mk_float, n_moment_electron}
+    electron_upar::MPISharedArray{mk_float, n_moment_electron}
+    electron_ppar::MPISharedArray{mk_float, n_moment_electron}
     # neutral particles 
     pdf_neutral::MPISharedArray{mk_float, n_distribution_neutral}
     density_neutral::MPISharedArray{mk_float, n_moment_neutral}
diff --git a/src/post_processing.jl b/src/post_processing.jl
index c793ac540..7a68df580 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -1164,7 +1164,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
     denstot_max = maximum(maximum(dtot) for dtot in denstot) + 0.1
     plot(legend=legend)
     for (t, dtot, run_label) ∈ zip(time, denstot, run_names)
-        @views plot!(t, dtot[1,:], ylims=(denstot_min,denstot_max), xlabel="time",
+        @views plot!(t, dtot, ylims=(denstot_min,denstot_max), xlabel="time",
                      ylabel="∑ⱼn̅ⱼ(t)/∑ⱼn̅ⱼ(0)", linewidth=2, label=run_label)
     end
     outfile = string(prefix, "_$(label)_denstot_vs_t.pdf")
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 66f6c609f..54d2a54e9 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -13,7 +13,6 @@ using ..debugging
 using ..file_io: write_data_to_ascii, write_moments_data_to_binary, write_dfns_data_to_binary, debug_dump
 using ..looping
 using ..moment_kinetics_structs: scratch_pdf
-using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: update_moments!, update_moments_neutral!, reset_moments_status!
 using ..velocity_moments: update_density!, update_upar!, update_ppar!, update_qpar!
 using ..velocity_moments: update_neutral_density!, update_neutral_qz!
@@ -52,6 +51,9 @@ using ..energy_equation: energy_equation!, neutral_energy_equation!
 using ..em_fields: setup_em_fields, update_phi!
 using ..manufactured_solns: manufactured_sources
 using ..advection: advection_info
+using ..electron_fluid_equations: calculate_electron_density!
+using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
+using ..electron_fluid_equations: calculate_electron_qpar!
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
 using Dates
@@ -346,6 +348,21 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             scratch[1].ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
         end
 
+        # update the electron, parallel flow density and parallel pressure
+        # in case the corresponding ion quantities have been changed by applying
+        # constraints to the ion pdf
+        calculate_electron_density!(moments.electron.dens, moments.ion.dens)
+        calculate_electron_upar_from_charge_conservation!(moments.electron.upar, moments.electron.dens,
+                                                          moments.ion.upar, moments.ion.dens, composition.electron_physics)
+        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
+        # update the electron moment entries in the scratch array
+        begin_r_z_region()
+        @loop_r_z ir iz begin
+            scratch[1].electron_density[iz,ir] = moments.electron.dens[iz,ir]
+            scratch[1].electron_upar[iz,ir] = moments.electron.upar[iz,ir]
+            scratch[1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+        end
+
         begin_sn_r_z_region(no_synchronize=true)
         @loop_sn_r_z isn ir iz begin
             scratch[1].pdf_neutral[:,:,:,iz,ir,isn] .= pdf.neutral.norm[:,:,:,iz,ir,isn]
@@ -619,9 +636,10 @@ that may be evolved separately via fluid equations
 function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
     # create n_rk_stages+1 structs, each of which will contain one pdf,
     # one density, and one parallel flow array
-    scratch = Vector{scratch_pdf{5,3,6,3}}(undef, n_rk_stages+1)
+    scratch = Vector{scratch_pdf{5,3,2,6,3}}(undef, n_rk_stages+1)
     pdf_dims = size(pdf_ion_in)
-    moment_dims = size(moments.ion.dens)
+    moment_ion_dims = size(moments.ion.dens)
+    moment_electron_dims = size(moments.electron.dens)
     pdf_neutral_dims = size(pdf_neutral_in)
     moment_neutral_dims = size(moments.neutral.dens)
     # populate each of the structs
@@ -629,10 +647,14 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
         # Allocate arrays in temporary variables so that we can identify them
         # by source line when using @debug_shared_array
         pdf_array = allocate_shared_float(pdf_dims...)
-        density_array = allocate_shared_float(moment_dims...)
-        upar_array = allocate_shared_float(moment_dims...)
-        ppar_array = allocate_shared_float(moment_dims...)
-        temp_z_s_array = allocate_shared_float(moment_dims...)
+        density_array = allocate_shared_float(moment_ion_dims...)
+        upar_array = allocate_shared_float(moment_ion_dims...)
+        ppar_array = allocate_shared_float(moment_ion_dims...)
+        temp_z_s_array = allocate_shared_float(moment_ion_dims...)
+
+        electron_density_array = allocate_shared_float(moment_electron_dims...)
+        electron_upar_array = allocate_shared_float(moment_electron_dims...)
+        electron_ppar_array = allocate_shared_float(moment_electron_dims...)
 
         pdf_neutral_array = allocate_shared_float(pdf_neutral_dims...)
         density_neutral_array = allocate_shared_float(moment_neutral_dims...)
@@ -642,6 +664,8 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
 
         scratch[istage] = scratch_pdf(pdf_array, density_array, upar_array,
                                       ppar_array, temp_z_s_array,
+                                      electron_density_array, electron_upar_array,
+                                      electron_ppar_array,
                                       pdf_neutral_array, density_neutral_array,
                                       uz_neutral_array, pz_neutral_array)
         @serial_region begin
@@ -650,6 +674,10 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
             scratch[istage].upar .= moments.ion.upar
             scratch[istage].ppar .= moments.ion.ppar
 
+            scratch[istage].electron_density .= moments.electron.dens
+            scratch[istage].electron_upar .= moments.electron.upar
+            scratch[istage].electron_ppar .= moments.electron.ppar
+
             scratch[istage].pdf_neutral .= pdf_neutral_in
             scratch[istage].density_neutral .= moments.neutral.dens
             scratch[istage].uz_neutral .= moments.neutral.uz
@@ -1160,6 +1188,13 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                   num_diss_params)
 
+    # update the electron moments
+    calculate_electron_density!(new_scratch.electron_density, new_scratch.density)
+    calculate_electron_upar_from_charge_conservation!(new_scratch.electron_upar, new_scratch.electron_density,
+        new_scratch.density, new_scratch.upar, composition.electron_physics)
+    @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
+    calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
+
     ##
     # update the neutral particle distribution and moments
     ##
@@ -1350,6 +1385,13 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         first_scratch.ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
     end
 
+    begin_r_z_region()
+    @loop_r_z ir iz begin
+        first_scratch.electron_density[iz,ir] = moments.electron.dens[iz,ir]
+        first_scratch.electron_upar[iz,ir] = moments.electron.upar[iz,ir]
+        first_scratch.electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+    end
+
     if composition.n_neutral_species > 0
         begin_sn_r_z_region()
         @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
@@ -1399,6 +1441,14 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         moments.ion.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
         moments.ion.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
     end
+    # No need to synchronize here as we only change electron quantities and previous
+    # region only changed ion quantities.
+    begin_r_z_region(no_synchronize=true)
+    @loop_r_z ir iz begin
+        moments.electron.dens[iz,ir] = final_scratch.electron_density[iz,ir]
+        moments.electron.upar[iz,ir] = final_scratch.electron_upar[iz,ir]
+        moments.electron.ppar[iz,ir] = final_scratch.electron_ppar[iz,ir]
+    end
     if composition.n_neutral_species > 0
         # No need to synchronize here as we only change neutral quantities and previous
         # region only changed plasma quantities.
@@ -1644,6 +1694,12 @@ function update_solution_vector!(evolved, moments, istage, composition, vpa, vpe
         new_evolved.upar[iz,ir,is] = old_evolved.upar[iz,ir,is]
         new_evolved.ppar[iz,ir,is] = old_evolved.ppar[iz,ir,is]
     end
+    begin_r_z_region()
+    @loop_r_z ir iz begin
+        new_evolved.electron_density[iz,ir] = old_evolved.electron_density[iz,ir]
+        new_evolved.electron_upar[iz,ir] = old_evolved.electron_upar[iz,ir]
+        new_evolved.electron_ppar[iz,ir] = old_evolved.electron_ppar[iz,ir]
+    end
     if composition.n_neutral_species > 0
         begin_sn_r_z_region()
         @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin

From ac94fdfb8421ea34b131a5556e84e780cc871d5e Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Fri, 9 Jun 2023 05:27:15 +0100
Subject: [PATCH 034/428] two changes in this commit: 1) added electron moment
 diagnostics to post_processing.jl 2) modified update_phi! to use the electron
 density stored in fvec when calculating the electron Boltzmann response,
 rather than re-computing.

---
 src/em_fields.jl       |  11 +---
 src/post_processing.jl | 126 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 127 insertions(+), 10 deletions(-)

diff --git a/src/em_fields.jl b/src/em_fields.jl
index 131979e9a..181c2af30 100644
--- a/src/em_fields.jl
+++ b/src/em_fields.jl
@@ -64,15 +64,8 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
         
     # TODO: parallelise this...
     @serial_region begin
-        ne = @view scratch_dummy.dummy_zrs[:,:,1]
         jpar_i = @view scratch_dummy.buffer_rs_1[:,:,1]
         N_e = @view scratch_dummy.buffer_rs_2[:,:,1]
-        ne .= 0.0
-        for is in 1:n_ion_species
-            @loop_r_z ir iz begin
-                 ne[iz,ir] += fvec.density[iz,ir,is]
-            end
-        end
         if composition.electron_physics == boltzmann_electron_response
             N_e .= 1.0
         elseif composition.electron_physics == boltzmann_electron_response_with_simple_sheath
@@ -89,7 +82,7 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
             # where positive sign above (and negative sign below)
             # is due to the fact that electrons reaching the wall flow towards more negative z.
             # Using J_||e + J_||i = 0, and rearranging for N_e, we have
-            #N_e = - 2.0 * sqrt( pi * composition.me_over_mi) * jpar_i * exp( - composition.phi_wall / composition.T_e)
+            # N_e = - 2.0 * sqrt( pi * composition.me_over_mi) * jpar_i * exp( - composition.phi_wall / composition.T_e)
             @loop_r ir begin
                  N_e[ir] = - 2.0 * sqrt( pi * composition.me_over_mi) * jpar_i[ir] * exp( - composition.phi_wall / composition.T_e)
                  # N_e must be positive, so force this in case a numerical error or something
@@ -100,7 +93,7 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
         end
         if composition.electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
             @loop_r_z ir iz begin
-                 fields.phi[iz,ir] = composition.T_e * log(ne[iz,ir] / N_e[ir])
+                fields.phi[iz,ir] = composition.T_e * log(fvec.electron_density[iz,ir] / N_e[ir])
             end
         end
     end
diff --git a/src/post_processing.jl b/src/post_processing.jl
index 7a68df580..a507638a6 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -9,6 +9,7 @@ export compare_neutral_pdf_symbolic_test
 export compare_fields_symbolic_test
 export construct_global_zr_coords
 export allocate_global_zr_ion_moments
+export allocate_global_zr_electron_moments
 export allocate_global_zr_neutral_moments
 export allocate_global_zr_fields
 export get_coords_nelement
@@ -236,6 +237,15 @@ function allocate_global_zr_ion_dfns(nvpa_global, nvperp_global, nz_global, nr_g
     return f
 end
 
+function allocate_global_zr_electron_moments(nz_global,nr_global,ntime)
+    density = allocate_float(nz_global,nr_global,ntime)
+    parallel_flow = allocate_float(nz_global,nr_global,ntime)
+    parallel_pressure = allocate_float(nz_global,nr_global,ntime)
+    parallel_heat_flux = allocate_float(nz_global,nr_global,ntime)
+    thermal_speed = allocate_float(nz_global,nr_global,ntime)
+    return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed
+end
+
 function allocate_global_zr_neutral_dfns(nvz_global, nvr_global, nvzeta_global, nz_global, nr_global,
                                          n_ion_species, ntime)
     f = allocate_float(nvz_global, nvr_global, nvzeta_global, nz_global, nr_global,
@@ -410,6 +420,12 @@ function analyze_and_plot_data(prefix...)
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r),
                                    n_ion_species, ntime)
+    electron_density, electron_parallel_flow, electron_parallel_pressure, 
+        electron_parallel_heat_flux, electron_thermal_speed =
+        get_tuple_of_return_values(allocate_global_zr_electron_moments,
+                                   Tuple(this_z.n_global for this_z ∈ z),
+                                   Tuple(this_r.n_global for this_r ∈ r),
+                                   ntime)
     if any(n_neutral_species .> 0)
         neutral_density, neutral_uz, neutral_pz, neutral_qz, neutral_thermal_speed =
             get_tuple_of_return_values(allocate_global_zr_neutral_moments,
@@ -456,6 +472,27 @@ function analyze_and_plot_data(prefix...)
                                run_names, "moments", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r))
+    # electron particle moments
+    get_tuple_of_return_values(read_distributed_zr_data!, electron_density, "electron_density", run_names,
+                               "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
+    get_tuple_of_return_values(read_distributed_zr_data!, electron_parallel_flow, "electron_parallel_flow",
+                               run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
+    get_tuple_of_return_values(read_distributed_zr_data!, electron_parallel_pressure,
+                               "electron_parallel_pressure", run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
+    get_tuple_of_return_values(read_distributed_zr_data!, electron_parallel_heat_flux,
+                               "electron_parallel_heat_flux", run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
+    get_tuple_of_return_values(read_distributed_zr_data!, electron_thermal_speed, "electron_thermal_speed",
+                               run_names, "moments", nblocks,
+                               Tuple(this_z.n for this_z ∈ z),
+                               Tuple(this_r.n for this_r ∈ r))
     # neutral particle moments
     if any(n_neutral_species .> 0)
         get_tuple_of_return_values(read_distributed_zr_data!, neutral_density,
@@ -644,6 +681,11 @@ function analyze_and_plot_data(prefix...)
             Tuple(qpar[:,ir0,:,:] for qpar ∈ parallel_heat_flux_at_pdf_times),
             Tuple(vth[:,ir0,:,:] for vth ∈ thermal_speed_at_pdf_times),
             Tuple(f[:,ivperp0,:,ir0,:,:] for f ∈ ff),
+            Tuple(n[:,ir0,:,:] for n ∈ electron_density),
+            Tuple(upar[:,ir0,:,:] for upar ∈ electron_parallel_flow),
+            Tuple(ppar[:,ir0,:,:] for ppar ∈ electron_parallel_pressure),
+            Tuple(qpar[:,ir0,:,:] for qpar ∈ electron_parallel_heat_flux),
+            Tuple(vth[:,ir0,:,:] for vth ∈ electron_thermal_speed),
             Tuple(neutral_n[:,ir0,:,:] for neutral_n ∈ neutral_density),
             Tuple(uz[:,ir0,:,:] for uz ∈ neutral_uz),
             Tuple(pz[:,ir0,:,:] for pz ∈ neutral_pz),
@@ -926,7 +968,10 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed,
         phi_at_pdf_times, density_at_pdf_times, parallel_flow_at_pdf_times,
         parallel_pressure_at_pdf_times, parallel_heat_flux_at_pdf_times,
-        thermal_speed_at_pdf_times, ff, neutral_density, neutral_uz, neutral_pz,
+        thermal_speed_at_pdf_times, ff, 
+        electron_density, electron_parallel_flow, electron_parallel_pressure, 
+        electron_parallel_heat_flux, electron_thermal_speed,
+        neutral_density, neutral_uz, neutral_pz,
         neutral_qz, neutral_thermal_speed, neutral_density_at_pdf_times,
         neutral_uz_at_pdf_times, neutral_pz_at_pdf_times, neutral_qz_at_pdf_times,
         neutral_thermal_speed_at_pdf_times, neutral_ff, n_ion_species, n_neutral_species,
@@ -964,6 +1009,12 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
         parallel_heat_flux, delta_qpar, qpar_fldline_avg,
         pp, run_names, time, itime_min, itime_max, nwrite_movie, z, iz0, n_ion_species,
         "ion")
+
+    # create the requested plots of the electron moments
+    plot_electron_moments(electron_density, electron_parallel_flow, electron_parallel_pressure,
+        electron_parallel_heat_flux, electron_thermal_speed, pp, run_names, time, itime_min,
+        itime_max, nwrite_movie, z, iz0)
+
     if maximum(n_neutral_species) > 0
         # analyze the velocity neutral moments data
         neutral_density_fldline_avg, neutral_uz_fldline_avg, neutral_pz_fldline_avg,
@@ -1134,6 +1185,78 @@ function plot_fields(phi, delta_phi, time, itime_min, itime_max, nwrite_movie,
     println("done.")
 end
 
+function plot_electron_moments(density, parallel_flow, parallel_pressure,
+    parallel_heat_flux, thermal_speed, pp, run_names, time, itime_min,
+    itime_max, nwrite_movie, z, iz0)
+
+    println("Plotting electron moments data...")
+
+    # determine what to use for the legend depending on how many runs are being plotted
+    n_runs = length(run_names)
+    if n_runs == 1
+        prefix = run_names[1]
+        legend = false
+    else
+        prefix = default_compare_prefix
+        legend = true
+    end
+
+    moment_string = "electron_density"
+    if pp.plot_dens0_vs_t
+        plot_single_moment_z0(density, time, iz0, prefix, legend, run_names, moment_string)
+    end
+    # if pp.plot_dens_vs_z_t
+    #     plot_single_moment_vs_z_t(density, time, z, run_names, prefix, moment_string)
+    # end
+    moment_string = "electron_upar"
+    if pp.plot_upar0_vs_t
+        plot_single_moment_z0(parallel_flow, time, iz0, prefix, legend, run_names, moment_string)
+    end
+    # if pp.plot_upar_vs_z_t
+    #     plot_single_moment_vs_z_t(parallel_flow, time, z, run_names, prefix, moment_string)
+    # end
+    moment_string = "electron_ppar"
+    if pp.plot_ppar0_vs_t
+        plot_single_moment_z0(parallel_pressure, time, iz0, prefix, legend, run_names, moment_string)
+    end
+    # if pp.plot_ppar_vs_z_t
+    #     plot_single_moment_vs_z_t(parallel_pressure, time, z, run_names, prefix, moment_string)
+    # end
+    moment_string = "electron_vth"
+    if pp.plot_vth0_vs_t
+        plot_single_moment_z0(thermal_speed, time, iz0, prefix, legend, run_names, moment_string)
+    end
+    moment_string = "electron_qpar"
+    if pp.plot_qpar0_vs_t
+        plot_single_moment_z0(parallel_heat_flux, time, iz0, prefix, legend, run_names, moment_string)
+    end
+    # if pp.plot_qpar_vs_z_t
+    #     plot_single_moment_vs_z_t(parallel_heat_flux, time, z, run_names, prefix, moment_string)
+    # end
+end
+
+function plot_single_moment_z0(mom, time, iz0, prefix, legend, run_names, mom_string)
+    # plot the time trace of the moment at z=z0
+    #plot(time, log.(phi[i,:]), yscale = :log10)
+    plot(legend=legend)
+    for (t, p, run_label) ∈ zip(time, mom, run_names)
+        @views plot!(t, p[iz0,:], label=run_label)
+    end
+    outfile = string(prefix, "_", mom_string, "0_vs_t.pdf")
+    savefig(outfile)
+end
+
+function plot_single_moment_vs_z_t(mom, time, z, run_names, prefix, mom_string)
+    # make a heatmap plot of moment(z,t)
+    n_runs = length(run_names)
+    subplots = (heatmap(t, this_z.grid, p, xlabel="time", ylabel="z", title=run_label,
+                       c = :deep)
+               for (t, this_z, p, run_label) ∈ zip(time, z, mom, run_names))
+    plot(subplots..., layout=(1,n_runs), size=(600*n_runs, 400))
+    outfile = string(prefix, "_", mom_string, "_vs_z_t.pdf")
+    savefig(outfile)
+end
+
 """
 """
 function plot_moments(density, delta_density, density_fldline_avg,
@@ -1169,6 +1292,7 @@ function plot_moments(density, delta_density, density_fldline_avg,
     end
     outfile = string(prefix, "_$(label)_denstot_vs_t.pdf")
     savefig(outfile)
+    
     for is ∈ 1:maximum(n_species)
         spec_string = string(is)
         dens_min = minimum(minimum(n[:,is,:]) for n ∈ density)

From d2c290d853d08f255f6328292b151001d134d32f Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Fri, 23 Jun 2023 09:48:49 +0100
Subject: [PATCH 035/428] implemented electron fluid equations with a
 braginskii closure.  successfully reproduces Boltzmann electron response
 results with an appropriate heating source in the electron energy equation.

---
 runs/wall+sheath-bc_braginskii.toml |  79 ++++++++
 src/electron_fluid_equations.jl     | 286 ++++++++++++++++++++++------
 src/em_fields.jl                    |  40 +++-
 src/file_io.jl                      |   3 +-
 src/initial_conditions.jl           |  85 ++++-----
 src/input_structs.jl                |  22 ++-
 src/moment_kinetics.jl              |   7 +-
 src/moment_kinetics_input.jl        |  21 +-
 src/moment_kinetics_structs.jl      |   1 +
 src/time_advance.jl                 | 132 ++++++++++---
 src/velocity_moments.jl             | 166 ++++++++++------
 11 files changed, 630 insertions(+), 212 deletions(-)
 create mode 100644 runs/wall+sheath-bc_braginskii.toml

diff --git a/runs/wall+sheath-bc_braginskii.toml b/runs/wall+sheath-bc_braginskii.toml
new file mode 100644
index 000000000..13a21e1b0
--- /dev/null
+++ b/runs/wall+sheath-bc_braginskii.toml
@@ -0,0 +1,79 @@
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+electron_physics = "braginskii_fluid"
+run_name = "sheath-bc_cheb_braginskii_fixed_vth"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+electron_ionization_frequency = 0.0
+constant_ionization_rate = false
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 9
+z_nelement = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 17
+vpa_nelement = 10
+vpa_L = 8.0
+vpa_bc = "periodic"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.002
\ No newline at end of file
diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index cb61b8f11..6ae64d659 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -2,26 +2,37 @@ module electron_fluid_equations
 
 export calculate_electron_density!
 export calculate_electron_upar_from_charge_conservation!
+export electron_energy_equation!
 export calculate_electron_qpar!
+export calculate_electron_parallel_friction_force!
 
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
+using ..input_structs: braginskii_fluid
 
 """
 use quasineutrality to obtain the electron density from the initial
 densities of the various ion species:
     sum_i dens_i = dens_e
 inputs:
-    dens_e - electron density at previous time level
-    dens_i - updated ion density
+    dens_e = electron density at previous time level
+    updated = flag indicating if the electron density is updated
+    dens_i = updated ion density
 output:
-    dens_e - updated electron density
+    dens_e = updated electron density
+    updated = flag indicating that the electron density has been updated
 """
-function calculate_electron_density!(dens_e, dens_i)
-    dens_e .= 0.0
-    @loop_s_r_z is ir iz begin
-        dens_e[iz,ir] += dens_i[iz,ir,is]
+function calculate_electron_density!(dens_e, updated, dens_i)
+    # only update the electron density if it has not already been updated
+    if !updated
+        dens_e .= 0.0
+        # enforce quasineutrality
+        @loop_s_r_z is ir iz begin
+            dens_e[iz,ir] += dens_i[iz,ir,is]
+        end
     end
+    # set flag indicating that the electron density has been updated
+    updated = true
     return nothing
 end
 
@@ -30,48 +41,179 @@ use charge conservation equation to solve for the electron parallel flow density
     d/dz(sum_i n_i upar_i - n_e upar_e) = 0
     ==> [sum_i n_i upar_i](z) - [sum_i n_i upar_i](zbound) = [n_e upar_e](z) - [n_e upar_e](zbound)
 inputs: 
-    upar_e - should contain updated electron parallel flow density at boundaries in zed
-    dens_e - electron particle density
-    upar_i - ion parallel flow density
-    dens_i - ion particle density
+    upar_e = should contain updated electron parallel flow at boundaries in zed
+    updated = flag indicating whether the electron parallel flow is already updated
+    dens_e = electron particle density
+    upar_i = ion parallel flow density
+    dens_i = ion particle density
 output:
-    upar_e - contains the updated electron parallel flow density
+    upar_e = contains the updated electron parallel flow
 """
-function calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i, electron_model)
-    # get the number of zed grid points, nz
-    nz = size(upar_e,1)
-    # initialise the electron parallel flow density to zero
-    upar_e .= 0.0
-    # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
-    # is equal and opposite to the ion parallel current
-    if electron_model == boltzmann_electron_response_with_simple_sheath
-        # loop over ion species, adding each species contribution to the
-        # ion parallel particle flux at the boundaries in zed
-        @loop_s_r is ir begin
-            # electron_upar at this intermediate stage is actually
-            # the electron particle flux
-            upar_e[1,ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
-            upar_e[end,ir] += dens_i[end,ir,is] * upar_i[end,ir,is]
-        end
-        @loop_r ir begin
-            # fix the boundary flux
-            boundary_flux = upar_e[1,ir]
-            for iz in 2:nz-1
-                # initialise the electron particle flux to its value at the boundary in zed
-                upar_e[iz,ir] = boundary_flux
-                # add the contributions to the electron particle flux from the various ion species
-                # particle fluxes
-                @loop_s is begin
-                    upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
+function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
+    # only calculate the electron parallel flow if it is not already updated
+    if !updated
+        # get the number of zed grid points, nz
+        nz = size(upar_e,1)
+        # initialise the electron parallel flow density to zero
+        upar_e .= 0.0
+        # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
+        # is equal and opposite to the ion parallel current
+        if electron_model in (boltzmann_electron_response_with_simple_sheath, braginskii_fluid)
+            # loop over ion species, adding each species contribution to the
+            # ion parallel particle flux at the boundaries in zed
+            @loop_s_r is ir begin
+                # electron_upar at this intermediate stage is actually
+                # the electron particle flux
+                upar_e[1,ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
+                upar_e[end,ir] += dens_i[end,ir,is] * upar_i[end,ir,is]
+            end
+            @loop_r ir begin
+                # fix the boundary flux
+                boundary_flux = upar_e[1,ir]
+                for iz in 2:nz-1
+                    # initialise the electron particle flux to its value at the boundary in zed
+                    upar_e[iz,ir] = boundary_flux
+                    # add the contributions to the electron particle flux from the various ion species
+                    # particle fluxes
+                    @loop_s is begin
+                        upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
+                    end
+                    # convert from parallel particle flux to parallel particle density
+                    upar_e[iz,ir] /= dens_e[iz,ir]
                 end
-                # convert from parallel particle flux to parallel particle density
-                upar_e[iz,ir] /= dens_e[iz,ir]
+                # convert from parallel particle flux to parallel particle density for boundary points
+                upar_e[1,ir] /= dens_e[1,ir]
+                upar_e[end,ir] /= dens_e[end,ir]
             end
-            # convert from parallel particle flux to parallel particle density for boundary points
-            upar_e[1,ir] /= dens_e[1,ir]
-            upar_e[end,ir] /= dens_e[end,ir]
         end
     end
+    updated = true
+    return nothing
+end
+
+"""
+use the electron energy equation to evolve the electron temperature via
+an explicit time advance.
+NB: so far, this is only set up for 1D problem, where we can assume
+an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
+"""
+function electron_energy_equation!(ppar, fvec, moments, collisions, dt, spectral, composition,
+                                   num_diss_params, dens)
+    begin_r_z_region()
+    # define some abbreviated variables for convenient use in rest of function
+    me_over_mi = composition.me_over_mi
+    nu_ei = collisions.nu_ei
+    # calculate contribution to rhs of energy equation (formulated in terms of pressure)
+    # arising from derivatives of ppar, qpar and upar
+    @loop_r_z ir iz begin
+        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.electron.dppar_dz_upwind[iz,ir]
+                    + (2/3)*moments.electron.dqpar_dz[iz,ir]
+                    + (5/3)*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
+    end
+    # compute the contribution to the rhs of the energy equation
+    # arising from artificial diffusion
+    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    if diffusion_coefficient > 0.0
+        @loop_r_z ir iz begin
+            ppar[iz,ir] += dt*diffusion_coefficient*moments.electron.d2ppar_dz2[iz,ir]
+        end
+    end
+    # compute the contribution to the rhs of the energy equation
+    # arising from electron-ion collisions
+    if nu_ei > 0.0
+        @loop_r_z ir iz begin
+            ppar[iz,ir] += dt*(2 * me_over_mi * nu_ei * (fvec.ppar[iz,ir]*fvec.electron_density[iz,ir]/fvec.density[iz,ir]
+                                                     - fvec.electron_ppar[iz,ir])
+                        + (2/3) * moments.electron.parallel_friction[iz,ir] * (fvec.upar[iz,ir]-fvec.electron_upar[iz,ir]))
+        end
+    end
+    # add in contributions due to charge exchange/ionization collisions
+    if composition.n_neutral_species > 0
+        if abs(collisions.charge_exchange_electron) > 0.0
+            @loop_s_r_z is ir iz begin
+                ppar[iz,ir] +=
+                    dt * me_over_mi * collisions.charge_exchange_electron * (
+                    2*(fvec.electron_density[iz,ir]*fvec.pz_neutral[iz,ir,is] -
+                    fvec.density_neutral[iz,ir,is]*fvec.electron_ppar[iz,ir]) +
+                    (2/3)*fvec.electron_density[iz,ir]*fvec.density_neutral[iz,ir,is] *
+                    (fvec.uz_neutral[iz,ir,is] - fvec.electron_upar[iz,ir])^2)
+            end
+        end
+        if abs(collisions.ionization_electron) > 0.0
+            @loop_s_r_z is ir iz begin
+                ppar[iz,ir] -=
+                    dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
+                    fvec.electron_ppar[iz,ir] +
+                    (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
+            end
+        end
+    end
+    # calculate the external electron heat source, if any
+    calculate_electron_heat_source!(moments.electron.heat_source, fvec.electron_ppar, moments.electron.dupar_dz,
+                                    fvec.density_neutral, collisions.ionization,
+                                    fvec.electron_upar, moments.electron.dppar_dz_upwind, fvec.electron_density,
+                                    dens, dt)
+    # add the contribution from the electron heat source                                
+    @loop_r_z ir iz begin
+        ppar[iz,ir] += dt * moments.electron.heat_source[iz,ir]
+    end
+    return nothing
+end
+
+"""
+solve the electron force balance (parallel momentum) equation for the
+parallel electric field, Epar:
+    Epar = -dphi/dz = (2/n_e) * (-dppar_e/dz + friction_force + n_e * m_e  * n_n * R_en * (u_n - u_e))
+    NB: in 1D only Epar is needed for update of ion pdf, so boundary phi is irrelevant
+inputs:
+    dens_e = electron density
+    dppar_dz = zed derivative of electron parallel pressure
+    nu_ei = electron-ion collision frequency
+    friction = electron-ion parallel friction force
+    n_neutral_species = number of evolved neutral species
+    charge_exchange = electron-neutral charge exchange frequency
+    me_over_mi = electron-ion mass ratio
+    dens_n = neutral density
+    upar_n = neutral parallel flow
+    upar_e = electron parallel flow
+output:
+    Epar = parallel electric field
+"""
+function calculate_Epar_from_electron_force_balance!(Epar, dens_e, dppar_dz, nu_ei, friction, n_neutral_species,
+                                                     charge_exchange, me_over_mi, dens_n, upar_n, upar_e)
+    # get the contribution to Epar from the parallel pressure
+    @loop_r_z ir iz begin
+        Epar[iz,ir] = -(2/dens_e[iz,ir]) * dppar_dz[iz,ir]
+    end
+    # if electron-ion collisions are taken into account, include electron-ion parallel friction
+    if nu_ei > 0
+        @loop_r_z ir iz begin
+            Epar[iz,ir] += (2/dens_e[iz,ir]) * friction[iz,ir]
+        end
+    end
+    # if there are neutral species evolved and accounting for charge exchange collisions with neutrals
+    if n_neutral_species > 0 && charge_exchange > 0
+        @loop_r_z ir iz begin
+            Epar[iz,ir] += 2 * me_over_mi * dens_n[iz,ir] * charge_exchange * (upar_n[iz,ir] - upar_e[iz,ir]) 
+        end
+    end
+    return nothing
+end
+
+"""
+"""
+function calculate_electron_parallel_friction_force!(friction, dens_e, upar_e, upar_i, dTe_dz,
+                                                     me_over_mi, nu_ei, electron_model)
+    if electron_model == braginskii_fluid
+        @loop_r_z ir iz begin
+            friction[iz,ir] = -0.71 * dens_e[iz,ir] * dTe_dz[iz,ir]
+        end
+        @loop_r_z ir iz begin
+            friction[iz,ir] += 0.51 * dens_e[iz,ir] * me_over_mi * nu_ei * (upar_i[iz,ir] - upar_e[iz,ir])
+        end
+    else
+        @. friction = 0.0
+    end
     return nothing
 end
 
@@ -79,23 +221,57 @@ end
 calculate the parallel component of the electron heat flux.
 there are currently two supported options for the parallel heat flux:
     Braginskii collisional closure - qpar_e = -3.16*ppar_e/(m_e*nu_ei)*dT/dz - 0.71*ppar_e*(upar_i-upar_e)
-    collisionless closure - d(qpar_e)/dz = 0 ==> qpar_e = 0
+    collisionless closure - d(qpar_e)/dz = 0 ==> qpar_e = constant
 inputs:
-    qpar_e - parallel electron heat flux at the previous time level
+    qpar_e = parallel electron heat flux at the previous time level
+    qpar_updated = flag indicating whether qpar is updated already
+    ppar_e = electron parallel pressure
+    upar_e = electron parallel flow
+    dTe_dz = zed derivative of electron temperature
+    upar_i = ion parallel flow
+    nu_ei = electron-ion collision frequency
+    me_over_mi = electron-to-ion mass ratio
+    electron_model = choice of model for electron physics 
 output:
-    qpar_e - updated parallel electron heat flux
+    qpar_e = updated parallel electron heat flux
+    qpar_updated = flag indicating that the parallel electron heat flux is updated
 """
-function calculate_electron_qpar!(qpar_e, electron_model)
-    if electron_model == "braginskii_fluid"
-        # NEED TO FIX!
-        qpar_e .= 0.0 
-    else
-        # if not using Braginskii fluid model, then assume for now
-        # that we are in the collisionless limit, for which d/dz(qpar_e) = 0.
-        # take qpar_e = 0
-        qpar_e .= 0.0
+function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz, upar_i, nu_ei, me_over_mi, electron_model)
+    # only calculate qpar_e if needs updating
+    if qpar_updated == false
+        if electron_model == braginskii_fluid
+            # use the classical Braginskii expression for the electron heat flux
+            @loop_r_z ir iz begin
+                qpar_e[iz,ir] = - 0.71 * ppar_e[iz,ir] * (upar_i[iz,ir] - upar_e[iz,ir])
+            end
+            if nu_ei > 0.0
+                @loop_r_z ir iz begin
+                    qpar_e[iz,ir] -= 3.16 * ppar_e[iz,ir] / (me_over_mi * nu_ei) * dTe_dz[iz,ir] 
+                end
+            end
+        else
+            # if not using Braginskii fluid model, then assume for now
+            # that we are in the collisionless limit, for which d/dz(qpar_e) = 0.
+            # take qpar_e = 0
+            @. qpar_e = 0.0
+        end
     end
+    # qpar has been updated
+    qpar_updated = true
     return nothing
 end
 
+function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n, ionization,
+                                         upar_e, dppar_dz_upwind, dens_e, dens_new, dt)
+    @loop_r_z ir iz begin
+        heat_source[iz,ir] = (2/3) * ppar_e[iz,ir] * dupar_dz[iz,ir]
+    end
+    n_neutral_species = size(dens_n, 3)
+    if n_neutral_species > 0 && ionization > 0.0
+       @loop_s_r_z is ir iz begin
+           heat_source[iz,ir] += ppar_e[iz,ir] * dens_n[iz,ir,is] * ionization
+       end
+    end
+end
+
 end
\ No newline at end of file
diff --git a/src/em_fields.jl b/src/em_fields.jl
index 181c2af30..bb18d49f1 100644
--- a/src/em_fields.jl
+++ b/src/em_fields.jl
@@ -14,6 +14,7 @@ using ..moment_kinetics_structs: em_fields_struct
 using ..velocity_moments: update_density!
 #using ..calculus: derivative!
 using ..derivatives: derivative_r!, derivative_z!
+using ..electron_fluid_equations: calculate_Epar_from_electron_force_balance!
 
 """
 """
@@ -28,8 +29,11 @@ end
 """
 update_phi updates the electrostatic potential, phi
 """
-function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, scratch_dummy)
+function update_phi!(fields, fvec, z, r, composition, collisions, moments, 
+    z_spectral, r_spectral, scratch_dummy)
+    # define a new variable for the number of ion species for brevity of following code
     n_ion_species = composition.n_ion_species
+    # check bounds of fields and fvec arrays
     @boundscheck size(fields.phi,1) == z.n || throw(BoundsError(fields.phi))
     @boundscheck size(fields.phi,2) == r.n || throw(BoundsError(fields.phi))
     @boundscheck size(fields.phi0,1) == z.n || throw(BoundsError(fields.phi0))
@@ -47,6 +51,8 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
     # when there is only one species.
     
     begin_serial_region()#(no_synchronize=true)
+
+    dens_e = fvec.electron_density
     # in serial as both s, r and z required locally
     if (composition.n_ion_species > 1 || true ||
         composition.electron_physics == boltzmann_electron_response_with_simple_sheath)
@@ -93,8 +99,14 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
         end
         if composition.electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
             @loop_r_z ir iz begin
-                fields.phi[iz,ir] = composition.T_e * log(fvec.electron_density[iz,ir] / N_e[ir])
+                fields.phi[iz,ir] = composition.T_e * log(dens_e[iz,ir] / N_e[ir])
             end
+        elseif composition.electron_physics == braginskii_fluid
+            calculate_Epar_from_electron_force_balance!(fields.Ez, dens_e, moments.electron.dppar_dz,
+                collisions.nu_ei, moments.electron.parallel_friction,
+                composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
+                fvec.density_neutral, fvec.uz_neutral, fvec.electron_upar)
+            calculate_phi_from_Epar!(fields.phi, fields.Ez, z.cell_width)
         end
     end
     ## can calculate phi at z = L and hence phi_wall(z=L) using jpar_i at z =L if needed
@@ -119,12 +131,26 @@ function update_phi!(fields, fvec, z, r, composition, z_spectral, r_spectral, sc
             # Er_constant defaults to 0.0 in moment_kinetics_input.jl
         end
     end
-    #Ez = - d phi / dz
-    @views derivative_z!(fields.Ez,-fields.phi,
-                scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1],
-                scratch_dummy.buffer_rs_3[:,1], scratch_dummy.buffer_rs_4[:,1],
-                z_spectral,z)
+    # if advancing electron fluid equations, solve for Ez directly from force balance
+    if composition.electron_physics != braginskii_fluid
+        # Ez = - d phi / dz
+        @views derivative_z!(fields.Ez,-fields.phi,
+                    scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1],
+                    scratch_dummy.buffer_rs_3[:,1], scratch_dummy.buffer_rs_4[:,1],
+                    z_spectral,z)
+    end
+    return nothing
+end
 
+function calculate_phi_from_Epar!(phi, Epar, dz)
+    # simple calculation of phi from Epar for now, with zero phi assumed at boundary
+    phi[1,:] .= 0.0
+    @loop_r_z ir iz begin
+        if iz > 1
+            phi[iz,ir] = phi[iz-1,ir] - dz[iz-1]*Epar[iz,ir]
+        end
+    end
+    return nothing
 end
 
 end
diff --git a/src/file_io.jl b/src/file_io.jl
index f599a35dc..56b2edaa1 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -1028,7 +1028,8 @@ function write_fields_ascii(flds, z, r, t, ascii_io)
         @inbounds begin
             for ir ∈ 1:r.n
                 for iz ∈ 1:z.n
-                    println(ascii_io,"t: ", t, "   r: ", r.grid[ir],"   z: ", z.grid[iz], "  phi: ", flds.phi[iz,ir])
+                    println(ascii_io,"t: ", t, "   r: ", r.grid[ir],"   z: ", z.grid[iz], "  phi: ", flds.phi[iz,ir],
+                            " Ez: ", flds.Ez[iz,ir])
                 end
             end
         end
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 507a7dc6c..be6f31ab0 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -36,7 +36,9 @@ using ..velocity_moments: update_ppar!, update_upar!, update_density!
 using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!
+using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
+using ..derivatives: derivative_z!
 
 using ..manufactured_solns: manufactured_solutions
 
@@ -102,8 +104,7 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, numerical_dissipation)
     electron = create_moments_electron(z.n, r.n,
-        evolve_moments.density, evolve_moments.parallel_flow,
-        evolve_moments.parallel_pressure, numerical_dissipation)
+        composition.electron_physics, numerical_dissipation)
     neutral = create_moments_neutral(z.n, r.n, composition.n_neutral_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, numerical_dissipation)
@@ -154,8 +155,8 @@ creates the normalised pdf and the velocity-space moments and populates them
 with a self-consistent initial condition
 """
 function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition, r, z,
-                               vperp, vpa, vzeta, vr, vz, vpa_spectral, vz_spectral,
-                               species, use_manufactured_solns)
+                               vperp, vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral,
+                               species, collisions, use_manufactured_solns)
     if use_manufactured_solns
         init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z,
                                              r, n_ion_species, n_neutral_species,
@@ -212,14 +213,12 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
                         vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
 
-        println("post_init_pdf_ni: ", moments.ion.dens[1,1,1])
         begin_s_r_z_region()
         # calculate the initial parallel heat flux from the initial un-normalised pdf
         update_qpar!(moments.ion.qpar, moments.ion.qpar_updated,
                      moments.ion.dens, moments.ion.upar, moments.ion.vth,
                      pdf.ion.norm, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
-        # calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
 
         if n_neutral_species > 0
             update_neutral_qz!(moments.neutral.qz, moments.neutral.qz_updated,
@@ -243,24 +242,38 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
 
     @serial_region begin 
         # initialise the electron density profile
-        init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
+        init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
         # initialise the electron parallel flow profile
-        init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
-            moments.ion.dens, composition.electron_physics)
+        init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
+            moments.ion.upar, moments.ion.dens, composition.electron_physics)
         # initialise the electron thermal speed profile
         init_electron_vth!(moments.electron.vth, composition.T_e)
-        # calculate the electron parallel pressure from the density and thermal speed
-        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
-        # calculate the electron parallel heat flux
-        calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
-        # the electron moments have now been updated
-        moments.electron.dens_updated = true
-        moments.electron.upar_updated = true
+        # calculate the electron temperature from the thermal speed
+#        @. moments.electron.temp = 0.5 * moments.electron.vth^2
+        @. moments.electron.temp = moments.electron.vth^2
+        # the electron temperature has now been updated
+        moments.electron.temp_updated = true
+        # calculate the electron parallel pressure from the density and temperature
+#        @. moments.electron.ppar = moments.electron.dens * moments.electron.temp
+        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
+        # the electron parallel pressure now been updated
         moments.electron.ppar_updated = true
+        # calculate the zed derivative of the initial electron temperature
+        # NB: ddens_dz, dupar_dz, dppar_dz and dqpar_dz are used as dummy arrays in this function
+        # NB: call, as they have yet to be calculated
+        @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
+            moments.electron.ddens_dz[:,1], moments.electron.dupar_dz[:,1], moments.electron.dppar_dz[:,1],
+            moments.electron.dqpar_dz[:,1], z_spectral, z)
+        # calculate the electron parallel heat flux
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, moments.electron.ppar,
+            moments.electron.upar, moments.electron.dT_dz, moments.ion.upar, 
+            collisions.nu_ei, composition.me_over_mi, composition.electron_physics)
+        # calculate the electron-ion parallel friction force
+        calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+            moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+            composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
     end
     
-    println("upar_i: ", moments.ion.upar[1,1,1], " upar_e: ", moments.electron.upar[1,1])
-
     return nothing
 end
 
@@ -450,49 +463,25 @@ end
 """
 initialise the electron density
 """
-function init_electron_density!(electron_density, ion_density, n_ion_species)
+function init_electron_density!(electron_density, updated, ion_density)
     # use quasineutrality to obtain the electron density from the initial
     # densities of the various ion species
-    calculate_electron_density!(electron_density, ion_density)
+    calculate_electron_density!(electron_density, updated, ion_density)
     return nothing
 end
 
 """
 initialise the electron parallel flow density
 """
-function init_electron_upar!(upar_e, dens_e, upar_i, dens_i, electron_model)
-    calculate_electron_upar_from_charge_conservation!(upar_e, dens_e, upar_i, dens_i, electron_model)
-    # # initialise the electron parallel flow density
-    # electron_upar .= 0.0
-    # println("electron_model: ", electron_model)
-    # # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
-    # # is equal and opposite to the ion parallel current
-    # if electron_model == boltzmann_electron_response_with_simple_sheath
-    #     # loop over ion species, adding each species contribution to the
-    #     # ion parallel particle flux at the boundaries in zed
-    #     for is ∈ 1:n_ion_species
-    #         for ir ∈ 1:nr
-    #             # electron_upar at this intermediate stage is actually
-    #             # the electron particle flux
-    #             electron_upar[1,ir] += ion_density[1,ir,is] * ion_upar[1,ir,is]
-    #             electron_upar[end,ir] += ion_density[end,ir,is] * ion_upar[end,ir,is]
-    #         end
-    #     end
-    #     # convert from electron particle flux to electron parallel flow density
-    #     for ir in 1:nr
-    #         electron_upar[1,ir] /= electron_density[1,ir]
-    #         electron_upar[end,ir] /= electron_density[end,ir]
-    #     end
-    #     # use charge conservation to solve for the electron upar for the rest of the zed domain
-    #     calculate_electron_upar_from_charge_conservation!(electron_upar, electron_density, ion_upar, ion_density)
-    # end
-    # return nothing
+function init_electron_upar!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
+    calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
+    return nothing
 end
 
 """
 initialise the electron thermal speed profile.
 for now the only initialisation option for the temperature is constant in z.
-returns vth0 = sqrt(Ts/Te) = 1.0
+returns vth0 = sqrt(2*Ts/Te)
 """
 function init_electron_vth!(vth, T_e)
     @loop_r_z ir iz begin
diff --git a/src/input_structs.jl b/src/input_structs.jl
index c5fbdc9d1..1858171cb 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -62,6 +62,7 @@ mutable struct advance_info
     continuity::Bool
     force_balance::Bool
     energy::Bool
+    electron_energy::Bool
     neutral_source_terms::Bool
     neutral_continuity::Bool
     neutral_force_balance::Bool
@@ -101,8 +102,15 @@ end
 
 """
 """
-@enum electron_physics_type boltzmann_electron_response boltzmann_electron_response_with_simple_sheath
-export electron_physics_type, boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath
+@enum electron_physics_type begin
+    boltzmann_electron_response 
+    boltzmann_electron_response_with_simple_sheath 
+    braginskii_fluid
+end
+export electron_physics_type
+export boltzmann_electron_response
+export boltzmann_electron_response_with_simple_sheath
+export braginskii_fluid
 
 """
 """
@@ -291,12 +299,20 @@ end
 """
 """
 mutable struct collisions_input
-    # charge exchange collision frequency
+    # ion-neutral charge exchange collision frequency
     charge_exchange::mk_float
+    # electron-neutral charge exchange collision frequency
+    charge_exchange_electron::mk_float
     # ionization collision frequency
     ionization::mk_float
+    # ionization collision frequency for electrons (probably should be same as for ions)
+    ionization_electron::mk_float
+    # ionization energy cost
+    ionization_energy::mk_float
     # if constant_ionization_rate = true, use an ionization term that is constant in z
     constant_ionization_rate::Bool
+    # electron-ion collision frequency
+    nu_ei::mk_float
 end
 
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index db0f3d017..92c0f893f 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -368,8 +368,8 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
         init_pdf_and_moments!(pdf, moments, boundary_distributions, composition, r, z,
-                              vperp, vpa, vzeta, vr, vz, vpa_spectral, vz_spectral,
-                              species, t_input.use_manufactured_solns_for_init)
+                              vperp, vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral,
+                              species, collisions, t_input.use_manufactured_solns_for_init)
         # initialize time variable
         code_time = 0.
     else
@@ -381,7 +381,6 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
                                             composition, r, z, vpa, vperp, vzeta, vr, vz)
         _block_synchronize()
     end
-    println("dens0_ion: ", moments.ion.dens[1,1,1], " dens0_electron: ", moments.electron.dens[1,1,1])
 
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
@@ -390,7 +389,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
     scratch, advance, scratch_dummy, manufactured_source_list =
         setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral,
-            r_spectral, composition, drive_input, moments, fields, t_input, collisions, species,
+            r_spectral, composition, moments, fields, t_input, collisions,
             geometry, boundary_distributions, num_diss_params, restarting)
     # setup i/o
     ascii_io, io_moments, io_dfns = setup_file_io(io_input, boundary_distributions, vz,
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index f4ef39521..11385b323 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -165,8 +165,12 @@ function mk_input(scan_input=Dict())
     #################### end specification of species inputs #####################
 
     collisions.charge_exchange = get(scan_input, "charge_exchange_frequency", 2.0*sqrt(species.ion[1].initial_temperature))
+    collisions.charge_exchange_electron = get(scan_input, "electron_charge_exchange_frequency", 0.0)
     collisions.ionization = get(scan_input, "ionization_frequency", collisions.charge_exchange)
+    collisions.ionization_electron = get(scan_input, "electron_ionization_frequency", collisions.ionization)
+    collisions.ionization_energy = get(scan_input, "ionization_energy", 0.0)
     collisions.constant_ionization_rate = get(scan_input, "constant_ionization_rate", false)
+    collisions.nu_ei = get(scan_input, "nu_ei", 0.0)
 
     # parameters related to the time stepping
     nstep = get(scan_input, "nstep", 5)
@@ -764,7 +768,8 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     #############################################################################
     # define default values and create corresponding mutable structs holding
     # information about the composition of the species and their initial conditions
-    if electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
+    if electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath, 
+                           braginskii_fluid)
         n_species = n_ion_species + n_neutral_species
     else
         n_species = n_ion_speces + n_neutral_species + 1
@@ -864,12 +869,22 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     drive_amplitude = 1.0
     drive_frequency = 1.0
     drive = drive_input_mutable(drive_phi, drive_amplitude, drive_frequency)
-    # charge exchange collision frequency
+    # ion-neutral charge exchange collision frequency
     charge_exchange = 0.0
+    # electron-neutral charge exchange collision frequency
+    charge_exchange_electron = 0.0
     # ionization collision frequency
     ionization = 0.0
+    # ionization collision frequency for electrons
+    ionization_electron = ionization
+    # ionization energy cost
+    ionization_energy = 0.0
     constant_ionization_rate = false
-    collisions = collisions_input(charge_exchange, ionization, constant_ionization_rate)
+    # electron-ion collision frequency
+    nu_ei = 0.0
+    # set up the collisions struct, containing info on collision frequencies
+    collisions = collisions_input(charge_exchange, charge_exchange_electron, ionization, 
+                                  ionization_electron, ionization_energy, constant_ionization_rate, nu_ei)
 
     Bzed = 1.0 # magnetic field component along z
     Bmag = 1.0 # magnetic field strength
diff --git a/src/moment_kinetics_structs.jl b/src/moment_kinetics_structs.jl
index 17ea13e4d..aed983108 100644
--- a/src/moment_kinetics_structs.jl
+++ b/src/moment_kinetics_structs.jl
@@ -21,6 +21,7 @@ struct scratch_pdf{n_distribution_ion, n_moment, n_moment_electron, n_distributi
     electron_density::MPISharedArray{mk_float, n_moment_electron}
     electron_upar::MPISharedArray{mk_float, n_moment_electron}
     electron_ppar::MPISharedArray{mk_float, n_moment_electron}
+    electron_temp::MPISharedArray{mk_float, n_moment_electron}
     # neutral particles 
     pdf_neutral::MPISharedArray{mk_float, n_distribution_neutral}
     density_neutral::MPISharedArray{mk_float, n_moment_neutral}
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 54d2a54e9..af579a631 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -19,6 +19,7 @@ using ..velocity_moments: update_neutral_density!, update_neutral_qz!
 using ..velocity_moments: update_neutral_uzeta!, update_neutral_uz!, update_neutral_ur!
 using ..velocity_moments: update_neutral_pzeta!, update_neutral_pz!, update_neutral_pr!
 using ..velocity_moments: calculate_moment_derivatives!, calculate_moment_derivatives_neutral!
+using ..velocity_moments: calculate_electron_moment_derivatives!
 using ..velocity_grid_transforms: vzvrvzeta_to_vpavperp!, vpavperp_to_vzvrvzeta!
 using ..initial_conditions: enforce_z_boundary_condition!, enforce_boundary_conditions!
 using ..initial_conditions: enforce_vpa_boundary_condition!, enforce_r_boundary_condition!
@@ -54,6 +55,10 @@ using ..advection: advection_info
 using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!
+using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
+using ..electron_fluid_equations: electron_energy_equation!
+using ..input_structs: braginskii_fluid
+using ..derivatives: derivative_z!
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
 using Dates
@@ -160,11 +165,10 @@ EM fields, and advection terms
 """
 function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
                              vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
-                             z_spectral, r_spectral, composition, drive_input, moments,
-                             fields, t_input, collisions, species, geometry,
+                             z_spectral, r_spectral, composition, moments,
+                             fields, t_input, collisions, geometry,
                              boundary_distributions, num_diss_params, restarting)
     # define some local variables for convenience/tidiness
-    n_species = composition.n_species
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
     # create array containing coefficients needed for the Runge Kutta time advance
@@ -176,7 +180,6 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     advance = setup_advance_flags(moments, composition, t_input, collisions,
                                   num_diss_params, rk_coefs, r, vperp, vpa, vzeta, vr, vz)
 
-
     begin_serial_region()
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
@@ -186,9 +189,14 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
                                    composition.n_ion_species,n_neutral_species_alloc)
-    # initialize the electrostatic potential
+   
     begin_serial_region()
-    update_phi!(fields, scratch[1], z, r, composition, z_spectral, r_spectral, scratch_dummy)
+    # update the derivatives of the electron moments as these may be needed when
+    # computing the electrostatic potential (and components of the electric field)
+    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 
+                                           z_spectral, num_diss_params, composition.electron_physics)
+    # initialize the electrostatic potential
+    update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
     @serial_region begin
         # save the initial phi(z) for possible use later (e.g., if forcing phi)
         fields.phi0 .= fields.phi
@@ -348,19 +356,43 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             scratch[1].ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
         end
 
-        # update the electron, parallel flow density and parallel pressure
+        # update the electron density, parallel flow and parallel pressure (and temperature)
         # in case the corresponding ion quantities have been changed by applying
         # constraints to the ion pdf
-        calculate_electron_density!(moments.electron.dens, moments.ion.dens)
-        calculate_electron_upar_from_charge_conservation!(moments.electron.upar, moments.electron.dens,
-                                                          moments.ion.upar, moments.ion.dens, composition.electron_physics)
-        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
+        calculate_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
+        calculate_electron_upar_from_charge_conservation!(moments.electron.upar, moments.electron.upar_updated,
+                                                          moments.electron.dens, moments.ion.upar, moments.ion.dens,
+                                                          composition.electron_physics)
+        # compute the updated electron temperature
+        # NB: not currently necessary, as initial vth is not directly dependent on ion quantities
+        @. moments.electron.temp = moments.electron.vth^2
+        # as the electron temperature has now been updated, set the appropriate flag
+        moments.electron.temp_updated = true
+        # compute the updated electron parallel pressure
+        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
+        # as the electron ppar has now been updated, set the appropriate flag
+        moments.electron.ppar_updated = true
+        # calculate the zed derivative of the initial electron temperature, potentially
+        # needed in the following calculation of the electron parallel friction force and
+        # parallel heat flux
+        @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # calculate the electron parallel heat flux
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, moments.electron.ppar,
+            moments.electron.upar, moments.electron.dT_dz, moments.ion.upar, 
+            collisions.nu_ei, composition.me_over_mi, composition.electron_physics)
+        # calculate the electron-ion parallel friction force
+        calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+            moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+            composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
         # update the electron moment entries in the scratch array
         begin_r_z_region()
         @loop_r_z ir iz begin
             scratch[1].electron_density[iz,ir] = moments.electron.dens[iz,ir]
             scratch[1].electron_upar[iz,ir] = moments.electron.upar[iz,ir]
             scratch[1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+            scratch[1].electron_temp[iz,ir] = moments.electron.temp[iz,ir]
         end
 
         begin_sn_r_z_region(no_synchronize=true)
@@ -372,11 +404,15 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
         end
     end
 
-    update_phi!(fields, scratch[1], z, r, composition, z_spectral, r_spectral,
-                scratch_dummy)
     calculate_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, num_diss_params)
+    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 
+                                           z_spectral, num_diss_params, composition.electron_physics)
     calculate_moment_derivatives_neutral!(moments, scratch[1], scratch_dummy, z,
                                           z_spectral, num_diss_params)
+    # update the electrostatic potential and components of the electric field, as pdfs and moments
+    # may have changed due to enforcing boundary/moment constraints                                      
+    update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral,
+                scratch_dummy)
 
     # Ensure all processes are synchronized at the end of the setup
     _block_synchronize()
@@ -407,6 +443,7 @@ function setup_advance_flags(moments, composition, t_input, collisions, num_diss
     advance_continuity = false
     advance_force_balance = false
     advance_energy = false
+    advance_electron_energy = false
     advance_neutral_z_advection = false
     advance_neutral_r_advection = false
     advance_neutral_vz_advection = false
@@ -505,7 +542,11 @@ function setup_advance_flags(moments, composition, t_input, collisions, num_diss
                 advance_neutral_energy = true
             end
         end
-
+        # if treating the electrons as a fluid with Braginskii closure, 
+        # then advance the electron energy equation
+        if composition.electron_physics == braginskii_fluid
+            advance_electron_energy = true
+        end
         # flag to determine if a d^2/dr^2 operator is present
         r_diffusion = (advance_numerical_dissipation && num_diss_params.r_dissipation_coefficient > 0.0)
         # flag to determine if a d^2/dvpa^2 operator is present
@@ -521,7 +562,7 @@ function setup_advance_flags(moments, composition, t_input, collisions, num_diss
                         advance_ionization, advance_ionization_1V,
                         advance_ionization_source, advance_numerical_dissipation,
                         advance_sources, advance_continuity, advance_force_balance,
-                        advance_energy, advance_neutral_sources,
+                        advance_energy, advance_electron_energy, advance_neutral_sources,
                         advance_neutral_continuity, advance_neutral_force_balance,
                         advance_neutral_energy, rk_coefs, manufactured_solns_test,
                         r_diffusion, vpa_diffusion, vz_diffusion)
@@ -655,6 +696,7 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
         electron_density_array = allocate_shared_float(moment_electron_dims...)
         electron_upar_array = allocate_shared_float(moment_electron_dims...)
         electron_ppar_array = allocate_shared_float(moment_electron_dims...)
+        electron_temp_array = allocate_shared_float(moment_electron_dims...)
 
         pdf_neutral_array = allocate_shared_float(pdf_neutral_dims...)
         density_neutral_array = allocate_shared_float(moment_neutral_dims...)
@@ -665,19 +707,21 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
         scratch[istage] = scratch_pdf(pdf_array, density_array, upar_array,
                                       ppar_array, temp_z_s_array,
                                       electron_density_array, electron_upar_array,
-                                      electron_ppar_array,
+                                      electron_ppar_array, electron_temp_array,
                                       pdf_neutral_array, density_neutral_array,
                                       uz_neutral_array, pz_neutral_array)
         @serial_region begin
+            # initialise the scratch arrays for the ion pdf and velocity moments
             scratch[istage].pdf .= pdf_ion_in
             scratch[istage].density .= moments.ion.dens
             scratch[istage].upar .= moments.ion.upar
             scratch[istage].ppar .= moments.ion.ppar
-
+            # initialise the scratch arrays for the electron velocity moments
             scratch[istage].electron_density .= moments.electron.dens
             scratch[istage].electron_upar .= moments.electron.upar
             scratch[istage].electron_ppar .= moments.electron.ppar
-
+            scratch[istage].electron_temp .= moments.electron.temp
+            # initialise the scratch arrays for the neutral velocity moments and pdf
             scratch[istage].pdf_neutral .= pdf_neutral_in
             scratch[istage].density_neutral .= moments.neutral.dens
             scratch[istage].uz_neutral .= moments.neutral.uz
@@ -1115,7 +1159,7 @@ or update them by taking the appropriate velocity moment of the evolved pdf
 """
 function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, vr, vzeta,
                     vpa, vperp, z, r, advect_objects, rk_coefs, istage, composition,
-                    geometry, num_diss_params, z_spectral, r_spectral, advance,
+                    collisions, geometry, num_diss_params, z_spectral, r_spectral, advance,
                     scratch_dummy)
     begin_s_r_z_region()
 
@@ -1188,12 +1232,35 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                   num_diss_params)
 
-    # update the electron moments
-    calculate_electron_density!(new_scratch.electron_density, new_scratch.density)
-    calculate_electron_upar_from_charge_conservation!(new_scratch.electron_upar, new_scratch.electron_density,
-        new_scratch.density, new_scratch.upar, composition.electron_physics)
-    @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
-    calculate_electron_qpar!(moments.electron.qpar, composition.electron_physics)
+    # update the lowest three electron moments (density, upar and ppar)
+    calculate_electron_density!(new_scratch.electron_density, moments.electron.dens_updated, new_scratch.density)
+    calculate_electron_upar_from_charge_conservation!(new_scratch.electron_upar, moments.electron.upar_updated,
+        new_scratch.electron_density, new_scratch.upar, new_scratch.density, composition.electron_physics)
+    # if electron model is braginskii_fluid, then ppar is evolved via the energy equation
+    # and is already updated;
+    # otherwise update assuming electron temperature is fixed in time
+    if composition.electron_physics == braginskii_fluid
+        @loop_r_z ir iz begin
+            new_scratch.electron_ppar[iz,ir] = (rk_coefs[1]*moments.electron.ppar[iz,ir] 
+                + rk_coefs[2]*old_scratch.electron_ppar[iz,ir] + rk_coefs[3]*new_scratch.electron_ppar[iz,ir])
+        end
+    else
+        @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
+    end
+    @. moments.electron.vth = sqrt(2 * new_scratch.electron_ppar / new_scratch.electron_density)
+    # regardless of electron model, electron ppar is now updated
+    moments.electron.ppar_updated = true
+    # calculate the corresponding zed derivatives of the moments
+    calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
+                                           num_diss_params, composition.electron_physics)
+    # update the electron parallel heat flux
+    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.electron_ppar,
+        new_scratch.electron_upar, moments.electron.dT_dz, moments.ion.upar, collisions.nu_ei,
+        composition.me_over_mi, composition.electron_physics)
+    # update the electron parallel friction force
+    calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
+        new_scratch.electron_upar, new_scratch.upar, moments.electron.dT_dz, composition.me_over_mi,
+        collisions.nu_ei, composition.electron_physics)
 
     ##
     # update the neutral particle distribution and moments
@@ -1258,7 +1325,9 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     end
 
     # update the electrostatic potential phi
-    update_phi!(fields, scratch[istage+1], z, r, composition, z_spectral, r_spectral, scratch_dummy)
+    update_phi!(fields, scratch[istage+1], z, r, composition, collisions, moments, 
+    z_spectral, r_spectral, scratch_dummy)
+    
     if !(( moments.evolve_upar || moments.evolve_ppar) &&
               istage == length(scratch)-1)
         # _block_synchronize() here because phi needs to be read on different ranks than
@@ -1390,6 +1459,7 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         first_scratch.electron_density[iz,ir] = moments.electron.dens[iz,ir]
         first_scratch.electron_upar[iz,ir] = moments.electron.upar[iz,ir]
         first_scratch.electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+        first_scratch.electron_temp[iz,ir] = moments.electron.temp[iz,ir]
     end
 
     if composition.n_neutral_species > 0
@@ -1423,8 +1493,8 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
             num_diss_params, advance, istage)
         @views rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, vr,
                           vzeta, vpa, vperp, z, r, advect_objects,
-                          advance.rk_coefs[:,istage], istage, composition, geometry,
-                          num_diss_params, spectral_objects.z_spectral,
+                          advance.rk_coefs[:,istage], istage, composition, collisions, 
+                          geometry, num_diss_params, spectral_objects.z_spectral,
                           spectral_objects.r_spectral, advance, scratch_dummy)
     end
 
@@ -1448,6 +1518,7 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
         moments.electron.dens[iz,ir] = final_scratch.electron_density[iz,ir]
         moments.electron.upar[iz,ir] = final_scratch.electron_upar[iz,ir]
         moments.electron.ppar[iz,ir] = final_scratch.electron_ppar[iz,ir]
+        moments.electron.temp[iz,ir] = final_scratch.electron_temp[iz,ir]
     end
     if composition.n_neutral_species > 0
         # No need to synchronize here as we only change neutral quantities and previous
@@ -1672,6 +1743,10 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
         neutral_energy_equation!(fvec_out.pz_neutral, fvec_in, moments, collisions, dt,
                                  z_spectral, composition, num_diss_params)
     end
+    if advance.electron_energy
+        electron_energy_equation!(fvec_out.electron_ppar, fvec_in, moments, collisions, dt,
+                                  z_spectral, composition, num_diss_params, fvec_out.density)
+    end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not
     reset_moments_status!(moments)
@@ -1699,6 +1774,7 @@ function update_solution_vector!(evolved, moments, istage, composition, vpa, vpe
         new_evolved.electron_density[iz,ir] = old_evolved.electron_density[iz,ir]
         new_evolved.electron_upar[iz,ir] = old_evolved.electron_upar[iz,ir]
         new_evolved.electron_ppar[iz,ir] = old_evolved.electron_ppar[iz,ir]
+        new_evolved.electron_temp[iz,ir] = old_evolved.electron_temp[iz,ir]
     end
     if composition.n_neutral_species > 0
         begin_sn_r_z_region()
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index 381bb08a6..ee78dc6d3 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -28,6 +28,7 @@ using ..calculus: integral
 using ..communication
 using ..derivatives: derivative_z!
 using ..looping
+using ..input_structs: braginskii_fluid
 
 #global tmpsum1 = 0.0
 #global tmpsum2 = 0.0
@@ -110,27 +111,29 @@ mutable struct moments_electron_substruct
     ppar::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not ppar needs updating before use
     ppar_updated::Bool
+    # this is the temperature
+    temp::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not temp needs updating before use
+    temp_updated::Bool
     # this is the parallel heat flux
     qpar::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not qpar needs updating before use
     qpar_updated::Bool
     # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
     vth::MPISharedArray{mk_float,2}
+    # this is the parallel friction force between ions and electrons
+    parallel_friction::MPISharedArray{mk_float,2}
+    # this is the electron heat source
+    heat_source::MPISharedArray{mk_float,2}
     # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
     # a factor of vth for each power of wpa in velocity space integrals.
     # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
     # and it is one otherwise
     v_norm_fac::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the upwinded z-derivative of the particle density
-    ddens_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the second-z-derivative of the particle density
-    d2dens_dz2::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the z-derivative of the parallel flow
     dupar_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the upwinded z-derivative of the parallel flow
-    dupar_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the second-z-derivative of the parallel flow
-    d2upar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the z-derivative of the parallel pressure
     dppar_dz::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the upwinded z-derivative of the parallel pressure
@@ -139,8 +142,10 @@ mutable struct moments_electron_substruct
     d2ppar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the z-derivative of the parallel heat flux
     dqpar_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
-    dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel temperature Tpar = ppar/dens
+    dT_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the temperature Tpar = ppar/dens
+    dT_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
 end
 
 """
@@ -305,8 +310,7 @@ end
 """
 create a moment struct containing information about the electron moments
 """
-function create_moments_electron(nz, nr, evolve_density, evolve_upar,
-                                evolve_ppar, numerical_dissipation)
+function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     # allocate array used for the particle density
     density = allocate_shared_float(nz, nr)
     # initialise Bool variable that indicates if the density is updated for each species
@@ -319,73 +323,50 @@ function create_moments_electron(nz, nr, evolve_density, evolve_upar,
     parallel_pressure = allocate_shared_float(nz, nr)
     # allocate Bool variable that indicates if the parallel pressure is updated for each species
     parallel_pressure_updated = false
+    # allocate array used for the temperature
+    temperature = allocate_shared_float(nz, nr)
+    # allocate Bool variable that indicates if the temperature is updated for each species
+    temperature_updated = false
     # allocate array used for the parallel flow
     parallel_heat_flux = allocate_shared_float(nz, nr)
     # allocate Bool variables that indicates if the parallel flow is updated for each species
     parallel_heat_flux_updated = false
+    # allocate array used for the election-ion parallel friction force
+    parallel_friction_force = allocate_shared_float(nz, nr)
+    # allocate array used for electron heat source
+    heat_source = allocate_shared_float(nz, nr)
     # allocate array used for the thermal speed
     thermal_speed = allocate_shared_float(nz, nr)
-    if evolve_ppar
-        v_norm_fac = thermal_speed
-    else
-        v_norm_fac = allocate_shared_float(nz, nr)
-        @serial_region begin
-            v_norm_fac .= 1.0
-        end
-    end
-
-    if evolve_density
-        ddens_dz_upwind = allocate_shared_float(nz, nr)
-    else
-        ddens_dz_upwind = nothing
-    end
-    if evolve_density &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
-
-        d2dens_dz2 = allocate_shared_float(nz, nr)
-    else
-        d2dens_dz2 = nothing
-    end
-    if evolve_density || evolve_upar || evolve_ppar
-        dupar_dz = allocate_shared_float(nz, nr)
-    else
-        dupar_dz = nothing
-    end
-    if evolve_upar
-        dupar_dz_upwind = allocate_shared_float(nz, nr)
-    else
-        dupar_dz_upwind = nothing
-    end
-    if evolve_upar &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
-
-        d2upar_dz2 = allocate_shared_float(nz, nr)
-    else
-        d2upar_dz2 = nothing
-    end
-    if evolve_upar
-        dppar_dz = allocate_shared_float(nz, nr)
+    # if evolving the electron pdf, it will be a function of the vth-normalised peculiar velocity
+    v_norm_fac = thermal_speed
+    # dn/dz is needed to obtain dT/dz (appearing in, e.g., Braginskii qpar) from dppar/dz
+    ddens_dz = allocate_shared_float(nz, nr)
+    # need dupar/dz to obtain, e.g., the updated electron temperature
+    dupar_dz = allocate_shared_float(nz, nr)
+    dppar_dz = allocate_shared_float(nz, nr)
+    if electron_model == braginskii_fluid
+        dppar_dz_upwind = allocate_shared_float(nz, nr)
+        dT_dz_upwind = allocate_shared_float(nz, nr)
     else
-        dppar_dz = nothing
+        dppar_dz_upwind = nothing
+        dT_dz_upwind = nothing
     end
-    if evolve_ppar
-        dppar_dz_upwind = allocate_shared_float(nz, nr)
+    if numerical_dissipation.moment_dissipation_coefficient > 0.0
         d2ppar_dz2 = allocate_shared_float(nz, nr)
-        dqpar_dz = allocate_shared_float(nz, nr)
-        dvth_dz = allocate_shared_float(nz, nr)
     else
-        dppar_dz_upwind = nothing
         d2ppar_dz2 = nothing
-        dqpar_dz = nothing
-        dvth_dz = nothing
     end
+    dqpar_dz = allocate_shared_float(nz, nr)
+    dT_dz = allocate_shared_float(nz, nr)
     
     # return struct containing arrays needed to update moments
     return moments_electron_substruct(density, density_updated, parallel_flow,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,
-        parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, v_norm_fac,
-        ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2, dppar_dz,
-        dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz)
+        temperature, temperature_updated, 
+        parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
+        parallel_friction_force, heat_source, v_norm_fac,
+        ddens_dz, dupar_dz, dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, 
+        dT_dz, dT_dz_upwind)
 end
 
 # neutral particles have natural mean velocities 
@@ -857,6 +838,61 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
     end
 end
 
+"""
+Pre-calculate spatial derivatives of the electron moments that will be needed for the time advance
+"""
+function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spectral,
+                                                numerical_dissipation, electron_model)
+    begin_r_region()
+
+    dens = scratch.electron_density
+    upar = scratch.electron_upar
+    ppar = scratch.electron_ppar
+    qpar = moments.electron.qpar
+    dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
+    buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
+    buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
+    buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
+    buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
+    buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
+    buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
+       
+    @views derivative_z!(moments.electron.dupar_dz, upar, buffer_r_1,
+                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+
+    # Upwinded using upar as advection velocity, to be used in energy equation
+    @loop_r_z ir iz begin
+        dummy_zr[iz,ir] = -upar[iz,ir]
+    end
+    if electron_model == braginskii_fluid
+        @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
+                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                             buffer_r_5, buffer_r_6, z_spectral, z)
+    end
+
+    # centred second derivative for dissipation
+    if numerical_dissipation.moment_dissipation_coefficient > 0.0
+        @views derivative_z!(dummy_zr, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
+                             buffer_r_4, z_spectral, z)
+        @views derivative_z!(moments.electron.d2ppar_dz2, dummy_zr, buffer_r_1,
+                             buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    end
+
+    @views derivative_z!(moments.electron.ddens_dz, dens, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    @views derivative_z!(moments.electron.dppar_dz, ppar, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    @views derivative_z!(moments.electron.dqpar_dz, qpar, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    # calculate the zed derivative of the electron temperature
+    @loop_r_z ir iz begin
+        # store the temperature in dummy_zr
+        dummy_zr[iz,ir] = 2*ppar[iz,ir]/dens[iz,ir]
+    end
+    @views derivative_z!(moments.electron.dT_dz, dummy_zr, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+end
+
 """
 update velocity moments of the evolved neutral pdf
 """
@@ -1610,6 +1646,10 @@ function reset_moments_status!(moments)
     moments.neutral.pzeta_updated .= false
     moments.neutral.pr_updated .= false
     moments.neutral.qz_updated .= false
+    moments.electron.dens_updated = false
+    moments.electron.upar_updated = false
+    moments.electron.ppar_updated = false
+    moments.electron.qpar_updated = false
 end
 
 end

From 72e103e3f19fd633ed7e1d9db1bba0617cef1109 Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Tue, 27 Jun 2023 09:16:32 +0100
Subject: [PATCH 036/428] added pz to the neutral moments ascii file, commented
 out the electron heat_source used for testing, and changed the boundary
 condition on electron ppar to be Te=Ti rather than Te=Tw.

---
 src/electron_fluid_equations.jl | 88 ++++++++++++++++++++-------------
 src/file_io.jl                  | 27 +++++-----
 src/initial_conditions.jl       | 23 +++------
 src/moment_kinetics.jl          |  2 +-
 src/time_advance.jl             | 12 +++--
 5 files changed, 83 insertions(+), 69 deletions(-)

diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index 6ae64d659..78ae5b9b8 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -97,8 +97,8 @@ an explicit time advance.
 NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
-function electron_energy_equation!(ppar, fvec, moments, collisions, dt, spectral, composition,
-                                   num_diss_params, dens)
+function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt, composition,
+                                   num_diss_params)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
     me_over_mi = composition.me_over_mi
@@ -121,10 +121,10 @@ function electron_energy_equation!(ppar, fvec, moments, collisions, dt, spectral
     # compute the contribution to the rhs of the energy equation
     # arising from electron-ion collisions
     if nu_ei > 0.0
-        @loop_r_z ir iz begin
-            ppar[iz,ir] += dt*(2 * me_over_mi * nu_ei * (fvec.ppar[iz,ir]*fvec.electron_density[iz,ir]/fvec.density[iz,ir]
-                                                     - fvec.electron_ppar[iz,ir])
-                        + (2/3) * moments.electron.parallel_friction[iz,ir] * (fvec.upar[iz,ir]-fvec.electron_upar[iz,ir]))
+        @loop_s_r_z is ir iz begin
+            ppar[iz,ir] += dt * (2 * me_over_mi * nu_ei * (fvec.ppar[iz,ir,is] - fvec.electron_ppar[iz,ir]))
+            ppar[iz,ir] += dt * ((2/3) * moments.electron.parallel_friction[iz,ir]
+                                * (fvec.upar[iz,ir,is]-fvec.electron_upar[iz,ir]))
         end
     end
     # add in contributions due to charge exchange/ionization collisions
@@ -141,22 +141,24 @@ function electron_energy_equation!(ppar, fvec, moments, collisions, dt, spectral
         end
         if abs(collisions.ionization_electron) > 0.0
             @loop_s_r_z is ir iz begin
-                ppar[iz,ir] -=
+                ppar[iz,ir] +=
                     dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
-                    fvec.electron_ppar[iz,ir] +
+                    fvec.electron_ppar[iz,ir] -
                     (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
             end
         end
     end
     # calculate the external electron heat source, if any
     calculate_electron_heat_source!(moments.electron.heat_source, fvec.electron_ppar, moments.electron.dupar_dz,
-                                    fvec.density_neutral, collisions.ionization,
-                                    fvec.electron_upar, moments.electron.dppar_dz_upwind, fvec.electron_density,
-                                    dens, dt)
+                                    fvec.density_neutral, collisions.ionization, collisions.ionization_energy,
+                                    fvec.electron_density, fvec.ppar, collisions.nu_ei, composition.me_over_mi,
+                                    composition.T_wall)
     # add the contribution from the electron heat source                                
     @loop_r_z ir iz begin
         ppar[iz,ir] += dt * moments.electron.heat_source[iz,ir]
     end
+    # enforce the parallel boundary condtion on the electron parallel pressure
+    enforce_parallel_BC_on_electron_pressure!(ppar, dens_i, composition.T_wall, fvec.ppar)
     return nothing
 end
 
@@ -193,8 +195,8 @@ function calculate_Epar_from_electron_force_balance!(Epar, dens_e, dppar_dz, nu_
     end
     # if there are neutral species evolved and accounting for charge exchange collisions with neutrals
     if n_neutral_species > 0 && charge_exchange > 0
-        @loop_r_z ir iz begin
-            Epar[iz,ir] += 2 * me_over_mi * dens_n[iz,ir] * charge_exchange * (upar_n[iz,ir] - upar_e[iz,ir]) 
+        @loop_s_r_z is ir iz begin
+            Epar[iz,ir] += 2 * me_over_mi * dens_n[iz,ir] * charge_exchange * (upar_n[iz,ir,is] - upar_e[iz,ir]) 
         end
     end
     return nothing
@@ -206,10 +208,10 @@ function calculate_electron_parallel_friction_force!(friction, dens_e, upar_e, u
                                                      me_over_mi, nu_ei, electron_model)
     if electron_model == braginskii_fluid
         @loop_r_z ir iz begin
-            friction[iz,ir] = -0.71 * dens_e[iz,ir] * dTe_dz[iz,ir]
+            friction[iz,ir] = -(1/2) * 0.71 * dens_e[iz,ir] * dTe_dz[iz,ir]
         end
-        @loop_r_z ir iz begin
-            friction[iz,ir] += 0.51 * dens_e[iz,ir] * me_over_mi * nu_ei * (upar_i[iz,ir] - upar_e[iz,ir])
+        @loop_s_r_z is ir iz begin
+            friction[iz,ir] += 0.51 * dens_e[iz,ir] * me_over_mi * nu_ei * (upar_i[iz,ir,is] - upar_e[iz,ir])
         end
     else
         @. friction = 0.0
@@ -239,21 +241,17 @@ output:
 function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz, upar_i, nu_ei, me_over_mi, electron_model)
     # only calculate qpar_e if needs updating
     if qpar_updated == false
+        @. qpar_e = 0.0
         if electron_model == braginskii_fluid
             # use the classical Braginskii expression for the electron heat flux
-            @loop_r_z ir iz begin
-                qpar_e[iz,ir] = - 0.71 * ppar_e[iz,ir] * (upar_i[iz,ir] - upar_e[iz,ir])
+            @loop_s_r_z is ir iz begin
+                qpar_e[iz,ir] -= 0.71 * ppar_e[iz,ir] * (upar_i[iz,ir,is] - upar_e[iz,ir])
             end
             if nu_ei > 0.0
                 @loop_r_z ir iz begin
-                    qpar_e[iz,ir] -= 3.16 * ppar_e[iz,ir] / (me_over_mi * nu_ei) * dTe_dz[iz,ir] 
+                    qpar_e[iz,ir] -= (1/2) * 3.16 * ppar_e[iz,ir] / (me_over_mi * nu_ei) * dTe_dz[iz,ir] 
                 end
             end
-        else
-            # if not using Braginskii fluid model, then assume for now
-            # that we are in the collisionless limit, for which d/dz(qpar_e) = 0.
-            # take qpar_e = 0
-            @. qpar_e = 0.0
         end
     end
     # qpar has been updated
@@ -261,16 +259,38 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz,
     return nothing
 end
 
-function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n, ionization,
-                                         upar_e, dppar_dz_upwind, dens_e, dens_new, dt)
-    @loop_r_z ir iz begin
-        heat_source[iz,ir] = (2/3) * ppar_e[iz,ir] * dupar_dz[iz,ir]
-    end
-    n_neutral_species = size(dens_n, 3)
-    if n_neutral_species > 0 && ionization > 0.0
-       @loop_s_r_z is ir iz begin
-           heat_source[iz,ir] += ppar_e[iz,ir] * dens_n[iz,ir,is] * ionization
-       end
+function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n, ionization, ionization_energy,
+                                         dens_e, ppar_i, nu_ei, me_over_mi, T_wall)
+    # heat_source currently only used for testing                                         
+    # @loop_r_z ir iz begin
+    #     heat_source[iz,ir] = (2/3) * ppar_e[iz,ir] * dupar_dz[iz,ir]
+    # end
+    # if nu_ei > 0.0
+    #     @loop_s_r_z is ir iz begin
+    #         heat_source[iz,ir] -= 2 * me_over_mi * nu_ei * (ppar_i[iz,ir,is] - ppar_e[iz,ir])
+    #     end
+    # end
+    # n_neutral_species = size(dens_n, 3)
+    # if n_neutral_species > 0 && ionization > 0.0
+    #    @loop_s_r_z is ir iz begin
+    #        heat_source[iz,ir] += (2/3) * dens_n[iz,ir,is] * dens_e[iz,ir] * ionization * ionization_energy
+    #    end
+    # end
+    # @loop_r ir begin
+    #     heat_source[1,ir] += 20 * 0.5 * (T_wall * dens_e[1,ir] - ppar_e[1,ir])
+    #     heat_source[end,ir] += 20 * (0.5 * T_wall * dens_e[end,ir] - ppar_e[end,ir])
+    # end
+    @. heat_source = 0.0
+    return nothing
+end
+
+function enforce_parallel_BC_on_electron_pressure!(ppar, dens, T_wall, ppar_i)
+    # assume T_e = T_i at boundaries in z
+    @loop_r ir begin
+        #ppar[1,ir] = 0.5 * dens[1,ir] * T_wall
+        #ppar[end,ir] = 0.5 * dens[end,ir] * T_wall
+        ppar[1,ir] = ppar_i[1,ir,1]
+        ppar[end,ir] = ppar_i[end,ir,1]
     end
 end
 
diff --git a/src/file_io.jl b/src/file_io.jl
index 56b2edaa1..1f01b41c1 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -127,14 +127,14 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
         out_prefix = string(io_input.output_dir, "/", io_input.run_name)
 
         if io_input.ascii_output
-            #ff_io = open_ascii_output_file(out_prefix, "f_vs_t")
+            ff_io = open_ascii_output_file(out_prefix, "f_vs_t")
             mom_ion_io = open_ascii_output_file(out_prefix, "moments_ion_vs_t")
             mom_eon_io = open_ascii_output_file(out_prefix, "moments_electron_vs_t")
             mom_ntrl_io = open_ascii_output_file(out_prefix, "moments_neutral_vs_t")
             fields_io = open_ascii_output_file(out_prefix, "fields_vs_t")
-            ascii = ascii_ios(mom_ion_io, mom_eon_io, mom_ntrl_io, fields_io)
+            ascii = ascii_ios(ff_io, mom_ion_io, mom_eon_io, mom_ntrl_io, fields_io)
         else
-            ascii = ascii_ios(nothing, nothing, nothing, nothing)
+            ascii = ascii_ios(nothing, nothing, nothing, nothing, nothing)
         end
 
         io_moments = setup_moments_io(out_prefix, io_input.binary_format, r, z,
@@ -905,7 +905,8 @@ include("file_io_hdf5.jl")
 
 """
 """
-function write_data_to_ascii(moments, fields, vpa, vperp, z, r, t, n_ion_species,
+#function write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, t, n_ion_species,
+function write_data_to_ascii(moments, fields, z, r, t, n_ion_species,
                              n_neutral_species, ascii_io::Union{ascii_ios,Nothing})
     if ascii_io === nothing || ascii_io.moments_ion === nothing
         # ascii I/O is disabled
@@ -915,7 +916,7 @@ function write_data_to_ascii(moments, fields, vpa, vperp, z, r, t, n_ion_species
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        #write_f_ascii(ff, z, vpa, t, ascii_io.ff)
+        #write_f_ascii(pdf, z, vpa, t, ascii_io.ff)
         write_moments_ion_ascii(moments.ion, z, r, t, n_ion_species, ascii_io.moments_ion)
         write_moments_electron_ascii(moments.electron, z, r, t, ascii_io.moments_electron)
         if n_neutral_species > 0
@@ -934,18 +935,19 @@ function write_f_ascii(f, z, vpa, t, ascii_io)
         # Only read/write from first process in each 'block'
 
         @inbounds begin
-            n_species = size(f,3)
-            for is ∈ 1:n_species
+            #n_species = size(f,3)
+            #for is ∈ 1:n_species
                 for j ∈ 1:vpa.n
                     for i ∈ 1:z.n
-                        println(ascii_io,"t: ", t, "   spec: ", is, ",   z: ", z.grid[i],
-                            ",  vpa: ", vpa.grid[j], ",   f: ", f[i,j,is])
+                        println(ascii_io,"t: ", t, "   z: ", z.grid[i],
+                            "  vpa: ", vpa.grid[j], "   fion: ", f.ion.norm[i,j,1], 
+                            "   fneutral: ", f.neutral.norm[i,j,1])
                     end
                     println(ascii_io)
                 end
                 println(ascii_io)
-            end
-            println(ascii_io)
+            #end
+            #println(ascii_io)
         end
     end
     return nothing
@@ -1008,7 +1010,8 @@ function write_moments_neutral_ascii(mom, z, r, t, n_species, ascii_io)
                     for iz ∈ 1:z.n
                         println(ascii_io,"t: ", t, "   species: ", is, "   r: ", r.grid[ir], "   z: ", z.grid[iz],
                             "  dens: ", mom.dens[iz,ir,is], "   uz: ", mom.uz[iz,ir,is],
-                            "   ur: ", mom.ur[iz,ir,is], "   uzeta: ", mom.uzeta[iz,ir,is])
+                            "   ur: ", mom.ur[iz,ir,is], "   uzeta: ", mom.uzeta[iz,ir,is],
+                            "   pz: ", mom.pz[iz,ir,is])
                     end
                 end
             end
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index be6f31ab0..aba6a537c 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -183,18 +183,6 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                 @. moments.neutral.pz = 0.5 * moments.neutral.dens * moments.neutral.vth^2
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
             end
-
-            # # initialise the electron density profile
-            # init_electron_density!(moments.electron.dens, moments.ion.dens, n_ion_species)
-
-            # println("ni: ", moments.ion.dens[1,1,1], " ne: ", moments.electron.dens[1,1])
-            # # initialise the electron parallel flow profile
-            # init_electron_upar!(moments.electron.upar, moments.electron.dens, moments.ion.upar, 
-            #     moments.ion.upar, n_ion_species, r.n, composition.electron_physics)
-            # # initialise the electron thermal speed profile
-            # init_electron_vth!(moments.electron.vth)
-            # # calculate the electron parallel pressure from the density and thermal speed
-            # @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.vth^2
         end
         moments.ion.dens_updated .= true
         moments.ion.upar_updated .= true
@@ -247,14 +235,12 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
             moments.ion.upar, moments.ion.dens, composition.electron_physics)
         # initialise the electron thermal speed profile
-        init_electron_vth!(moments.electron.vth, composition.T_e)
+        init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e)
         # calculate the electron temperature from the thermal speed
-#        @. moments.electron.temp = 0.5 * moments.electron.vth^2
         @. moments.electron.temp = moments.electron.vth^2
         # the electron temperature has now been updated
         moments.electron.temp_updated = true
         # calculate the electron parallel pressure from the density and temperature
-#        @. moments.electron.ppar = moments.electron.dens * moments.electron.temp
         @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
         # the electron parallel pressure now been updated
         moments.electron.ppar_updated = true
@@ -483,9 +469,12 @@ initialise the electron thermal speed profile.
 for now the only initialisation option for the temperature is constant in z.
 returns vth0 = sqrt(2*Ts/Te)
 """
-function init_electron_vth!(vth, T_e)
+function init_electron_vth!(vth_e, vth_i, T_e)
+    # @loop_r_z ir iz begin
+    #     vth_e[iz,ir] = sqrt(T_e)
+    # end
     @loop_r_z ir iz begin
-        vth[iz,ir] = sqrt(T_e)
+        vth_e[iz,ir] = vth_i[iz,ir,1]
     end
 end
 
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 92c0f893f..b29ee2aeb 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -396,7 +396,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
         vr, vzeta, vpa, vperp, z, r, composition, collisions, moments.evolve_density,
         moments.evolve_upar, moments.evolve_ppar, input_dict)
     # write initial data to ascii files
-    write_data_to_ascii(moments, fields, vpa, vperp, z, r, code_time, composition.n_ion_species, composition.n_neutral_species, ascii_io)
+    write_data_to_ascii(moments, fields, z, r, code_time, composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
     write_moments_data_to_binary(moments, fields, code_time, composition.n_ion_species,
diff --git a/src/time_advance.jl b/src/time_advance.jl
index af579a631..d98ebb0f1 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -823,7 +823,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                              Dates.format(now(), dateformat"H:MM:SS"))
                 end
             end
-            write_data_to_ascii(moments, fields, vpa, vperp, z, r, t,
+            write_data_to_ascii(moments, fields, z, r, t,
                                 composition.n_ion_species, composition.n_neutral_species,
                                 ascii_io)
             write_moments_data_to_binary(moments, fields, t, composition.n_ion_species,
@@ -1247,7 +1247,9 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     else
         @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
     end
-    @. moments.electron.vth = sqrt(2 * new_scratch.electron_ppar / new_scratch.electron_density)
+    @. moments.electron.temp = 2 * new_scratch.electron_ppar / new_scratch.electron_density
+    moments.electron.temp_updated = true
+    @. moments.electron.vth = sqrt(moments.electron.temp)
     # regardless of electron model, electron ppar is now updated
     moments.electron.ppar_updated = true
     # calculate the corresponding zed derivatives of the moments
@@ -1255,7 +1257,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                                            num_diss_params, composition.electron_physics)
     # update the electron parallel heat flux
     calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.electron_ppar,
-        new_scratch.electron_upar, moments.electron.dT_dz, moments.ion.upar, collisions.nu_ei,
+        new_scratch.electron_upar, moments.electron.dT_dz, new_scratch.upar, collisions.nu_ei,
         composition.me_over_mi, composition.electron_physics)
     # update the electron parallel friction force
     calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
@@ -1744,8 +1746,8 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                  z_spectral, composition, num_diss_params)
     end
     if advance.electron_energy
-        electron_energy_equation!(fvec_out.electron_ppar, fvec_in, moments, collisions, dt,
-                                  z_spectral, composition, num_diss_params, fvec_out.density)
+        electron_energy_equation!(fvec_out.electron_ppar, fvec_out.density, fvec_in, moments, collisions, dt,
+                                  composition, num_diss_params)
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not

From 9b68be0d9b12b0c84fe45f9b785302f59731f39a Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Wed, 28 Jun 2023 12:10:06 +0100
Subject: [PATCH 037/428] added input toml files for Braginskii electron fluid
 tests (June 2023 ExCALIBUR report)

---
 ...l+sheath-bc_braginskii_boltzmann_test.toml | 79 ++++++++++++++++++
 runs/wall+sheath-bc_braginskii_colls.toml     | 80 +++++++++++++++++++
 2 files changed, 159 insertions(+)
 create mode 100644 runs/wall+sheath-bc_braginskii_boltzmann_test.toml
 create mode 100644 runs/wall+sheath-bc_braginskii_colls.toml

diff --git a/runs/wall+sheath-bc_braginskii_boltzmann_test.toml b/runs/wall+sheath-bc_braginskii_boltzmann_test.toml
new file mode 100644
index 000000000..13b94d2ff
--- /dev/null
+++ b/runs/wall+sheath-bc_braginskii_boltzmann_test.toml
@@ -0,0 +1,79 @@
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "braginskii_fluid"
+run_name = "sheath-bc_cheb_braginskii_boltzmann_test"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+electron_ionization_frequency = 2.0
+ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 40000
+dt = 0.0005
+nwrite = 200
+nwrite_dfns = 1000
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 9
+z_nelement = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 17
+vpa_nelement = 10
+vpa_L = 8.0
+vpa_bc = "periodic"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01
diff --git a/runs/wall+sheath-bc_braginskii_colls.toml b/runs/wall+sheath-bc_braginskii_colls.toml
new file mode 100644
index 000000000..eb0476bc5
--- /dev/null
+++ b/runs/wall+sheath-bc_braginskii_colls.toml
@@ -0,0 +1,80 @@
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+electron_physics = "braginskii_fluid"
+run_name = "sheath-bc_cheb_braginskii_zdiss0p01"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 2.0
+nu_ei = 100000.0
+ionization_frequency = 2.0
+electron_ionization_frequency = 2.0
+ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 40000
+dt = 0.0001
+nwrite = 50
+nwrite_dfns = 1000
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 9
+z_nelement = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 17
+vpa_nelement = 10
+vpa_L = 8.0
+vpa_bc = "periodic"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01

From d8d4636d4359d41c791b2ddfc6e73c159f972a0a Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Wed, 28 Jun 2023 13:57:48 +0100
Subject: [PATCH 038/428] Added toml file for run with hard-wired electron
 Boltzmann response used in June 2023 Excalibur report.

---
 runs/wall+sheath-bc_boltzmann.toml | 78 ++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)
 create mode 100644 runs/wall+sheath-bc_boltzmann.toml

diff --git a/runs/wall+sheath-bc_boltzmann.toml b/runs/wall+sheath-bc_boltzmann.toml
new file mode 100644
index 000000000..1844deebf
--- /dev/null
+++ b/runs/wall+sheath-bc_boltzmann.toml
@@ -0,0 +1,78 @@
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "boltzmann_electron_response_with_simple_sheath"
+run_name = "sheath-bc_cheb_boltzmann"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+electron_ionization_frequency = 0.0
+ionization_energy = 0.0
+constant_ionization_rate = false
+nstep = 40000
+dt = 0.0005
+nwrite = 200
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 9
+z_nelement = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 17
+vpa_nelement = 10
+vpa_L = 8.0
+vpa_bc = "periodic"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01
\ No newline at end of file

From 32013d5aabc40077c3a434ee970df92cd9308621 Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Tue, 30 Jan 2024 16:03:45 +0000
Subject: [PATCH 039/428] many changes to get the kinetic treatment of
 electrons up and running, but this has not been merged with changes from
 master for several months. use at your own peril!!

---
 runs/wall+sheath-bc_braginskii.toml |    5 +-
 src/calculus.jl                     |   30 +-
 src/coordinates.jl                  |   42 +-
 src/derivatives.jl                  |   53 ++
 src/electron_fluid_equations.jl     |   56 +-
 src/electron_kinetic_equation.jl    | 1315 +++++++++++++++++++++++++++
 src/electron_vpa_advection.jl       |   59 ++
 src/electron_z_advection.jl         |   54 ++
 src/em_fields.jl                    |    8 +-
 src/file_io.jl                      |    2 +-
 src/initial_conditions.jl           |  288 +++++-
 src/input_structs.jl                |    2 +
 src/moment_kinetics.jl              |   45 +-
 src/moment_kinetics_input.jl        |    5 +-
 src/moment_kinetics_structs.jl      |    5 +-
 src/post_processing.jl              |   35 +-
 src/post_processing_input.jl        |    6 +-
 src/time_advance.jl                 |  176 ++--
 src/velocity_moments.jl             |   11 +-
 19 files changed, 2056 insertions(+), 141 deletions(-)
 create mode 100644 src/electron_kinetic_equation.jl
 create mode 100644 src/electron_vpa_advection.jl
 create mode 100644 src/electron_z_advection.jl

diff --git a/runs/wall+sheath-bc_braginskii.toml b/runs/wall+sheath-bc_braginskii.toml
index 13a21e1b0..60c8cf4d0 100644
--- a/runs/wall+sheath-bc_braginskii.toml
+++ b/runs/wall+sheath-bc_braginskii.toml
@@ -2,7 +2,7 @@ n_ion_species = 1
 n_neutral_species = 1
 #boltzmann_electron_response = false
 electron_physics = "braginskii_fluid"
-run_name = "sheath-bc_cheb_braginskii_fixed_vth"
+run_name = "sheath-bc_cheb_braginskii_Rion_test"
 evolve_moments_density = false
 evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
@@ -45,7 +45,8 @@ charge_exchange_frequency = 2.0
 electron_charge_exchange_frequency = 0.0
 nu_ei = 0.0
 ionization_frequency = 2.0
-electron_ionization_frequency = 0.0
+electron_ionization_frequency = 2.0
+ionization_energy = 1.0
 constant_ionization_rate = false
 nstep = 40000
 #nstep = 1
diff --git a/src/calculus.jl b/src/calculus.jl
index b88fcb9f6..113c31bf3 100644
--- a/src/calculus.jl
+++ b/src/calculus.jl
@@ -128,12 +128,12 @@ function second_derivative!(d2f, f, Q, coord, spectral)
         # points are not set by a boundary condition.
         # Full grid may be across processes and bc only applied to extreme ends of the
         # domain.
-        if coord.irank == 0
-            d2f[1] = 0.0
-        end
-        if coord.irank == coord.nrank - 1
-            d2f[end] = 0.0
-        end
+        #if coord.irank == 0
+        #    d2f[1] = 0.0
+        #end
+        #if coord.irank == coord.nrank - 1
+        #    d2f[end] = 0.0
+        #end
     elseif coord.bc == "periodic"
         # Need to get first derivatives from opposite ends of grid
         if coord.nelement_local != coord.nelement_global
@@ -507,6 +507,24 @@ function reconcile_element_boundaries_MPI!(df1d::AbstractArray{mk_float,Ndims},
     _block_synchronize()
 end
 
+# """
+# compute the 1D differentiation matrix for the given coordinate;
+# currently assumes no upwinding
+# """
+# function compute_1d_differentiation_matrix(coord, spectral::Union{Bool,<:chebyshev_info})
+#     # initialise all elements of the coord.scratch array to zero
+#     coord.scratch = 0.0
+#     for i ∈ coord.n
+#         # provide a unit impulse at the i-th grid point
+#         coord.scratch[i] = 1.0
+#         # compute the derivative of the unit impulse;
+#         # the solution is the ith column vector of the differentiation matrix
+#         @views derivative!(coord.differentiation_matrix[:, i], coord.scratch, coord, spectral)
+#         # zero out the impulse for the next iteration
+#         coord.scratch[i] = 0.0
+#     end
+# end
+
 """
 Computes the integral of the integrand, using the input wgts
 """
diff --git a/src/coordinates.jl b/src/coordinates.jl
index 704f67023..d30fc5fe3 100644
--- a/src/coordinates.jl
+++ b/src/coordinates.jl
@@ -5,6 +5,7 @@ module coordinates
 export define_coordinate, write_coordinate
 export equally_spaced_grid
 
+using LinearAlgebra
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_int
 using ..calculus: derivative!
@@ -17,7 +18,7 @@ using MPI
 """
 structure containing basic information related to coordinates
 """
-struct coordinate
+struct coordinate{T}
     # name is the name of the variable associated with this coordiante
     name::String
     # n_global is the total number of grid points associated with this coordinate
@@ -84,6 +85,10 @@ struct coordinate
     local_io_range::UnitRange{Int64}
     # global range to write into in output file
     global_io_range::UnitRange{Int64}
+    # array containing the LU-decomopsed, 1D differentiation matrix wrt this coordinate
+    #differentiation_matrix::Array{mk_float,2}
+    #differentiation_matrix::LU{mk_float,Array{mk_float,2},Vector{mk_int}}
+    differentiation_matrix::T
 end
 
 """
@@ -121,6 +126,22 @@ function define_coordinate(input, parallel_io::Bool=false)
     scratch_2d = allocate_float(input.ngrid, input.nelement_local)
     # struct containing the advection speed options/inputs for this coordinate
     advection = input.advection
+    # D_matrix will be the 1D differentiation matrix wrt this coordinate
+    # initialise to the identity matrix
+    D_matrix = rand(n_global, n_global)
+    #D_matrix .= 0.0
+    #for i ∈ 1:n_global
+    #    D_matrix[i,i] = 1.0
+    #end
+    # create an LU object for the differentiation matrix
+    # to be modified later once the actual D_matrix is calculated
+    LU_matrix = factorize(D_matrix)
+    #LU_matrix = LU{mk_float, Matrix{mk_float}, Vector{mk_int}}
+
+    #println("coord: ", input.name)
+    #println("D_matrix: ", D_matrix)
+    #println("LU_matrix: ", LU_matrix)
+
     # buffers for cyclic communication of boundary points
     # each chain of elements has only two external (off-rank)
     # endpoints, so only two pieces of information must be shared
@@ -149,7 +170,7 @@ function define_coordinate(input, parallel_io::Bool=false)
         cell_width, igrid, ielement, imin, imax, input.discretization, input.fd_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch),
         scratch_2d, copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
-        local_io_range, global_io_range)
+        local_io_range, global_io_range, LU_matrix)
 
     if input.discretization == "chebyshev_pseudospectral" && coord.n > 1
         # create arrays needed for explicit Chebyshev pseudospectral treatment in this
@@ -158,6 +179,23 @@ function define_coordinate(input, parallel_io::Bool=false)
         spectral = setup_chebyshev_pseudospectral(coord)
         # obtain the local derivatives of the uniform grid with respect to the used grid
         derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
+        # operate on the unit vectors to obtain the column vectors of the differentiation matrix
+        unit_vector = allocate_float(coord.n)
+        unit_vector .= 0.0
+        for i ∈ 1:coord.n
+            unit_vector[i] = 1.0
+            @views derivative!(D_matrix[:,i], unit_vector, coord, spectral)
+            unit_vector[i] = 0.0
+        end
+
+        println("coord: ", coord.name)
+        # replace the differentiaation matrix with its LU decomposition
+        if (coord.name == "z")
+            LU_matrix = factorize(D_matrix)
+            @. coord.differentiation_matrix.L = LU_matrix.L
+            @. coord.differentiation_matrix.U = LU_matrix.U
+            @. coord.differentiation_matrix.p = LU_matrix.p
+        end
     else
         # create dummy Bool variable to return in place of the above struct
         spectral = false
diff --git a/src/derivatives.jl b/src/derivatives.jl
index 8ef4bc21a..d44a001ba 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -154,6 +154,30 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
 	
 end
 
+#4D version for f[vpa,vperp,z,r] -> dfn electron particles
+function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_float,4},
+	dfdz_lower_endpoints::AbstractArray{mk_float,3},
+	dfdz_upper_endpoints::AbstractArray{mk_float,3},
+	z_send_buffer::AbstractArray{mk_float,3},
+	z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z)
+
+	# differentiate f w.r.t z
+	@loop_r_vperp_vpa ir ivperp ivpa begin
+		@views derivative!(dfdz[ivpa,ivperp,:,ir], f[ivpa,ivperp,:,ir], z, z_spectral)
+		# get external endpoints to reconcile via MPI
+		dfdz_lower_endpoints[ivpa,ivperp,ir] = z.scratch_2d[1,1]
+		dfdz_upper_endpoints[ivpa,ivperp,ir] = z.scratch_2d[end,end]
+	end
+	# now reconcile element boundaries across
+	# processes with large message involving all y
+	if z.nelement_local < z.nelement_global
+		reconcile_element_boundaries_MPI!(dfdz,
+		dfdz_lower_endpoints, dfdz_upper_endpoints,
+		z_send_buffer, z_receive_buffer, z)
+	end
+
+end
+
 #6D version for f[vz,vr,vzeta,z,r] -> dfn neutral particles
 function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6},
         dfdz_lower_endpoints::AbstractArray{mk_float,5},
@@ -338,6 +362,35 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
 	
 end
 
+#4D version for f[vpa,vperp,z,r] -> dfn electron particles
+function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_float,4},
+	advect, adv_fac_lower_buffer::AbstractArray{mk_float,3},
+	adv_fac_upper_buffer::AbstractArray{mk_float,3},
+	dfdz_lower_endpoints::AbstractArray{mk_float,3},
+	dfdz_upper_endpoints::AbstractArray{mk_float,3},
+	z_send_buffer::AbstractArray{mk_float,3},
+	z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z)
+
+# differentiate the pdf f w.r.t z
+@loop_r_vperp_vpa ir ivperp ivpa begin
+	@views derivative!(dfdz[ivpa,ivperp,:,ir], f[ivpa,ivperp,:,ir], z, advect[1].adv_fac[:,ivpa,ivperp,ir], z_spectral)
+	# get external endpoints to reconcile via MPI
+	dfdz_lower_endpoints[ivpa,ivperp,ir] = z.scratch_2d[1,1]
+	dfdz_upper_endpoints[ivpa,ivperp,ir] = z.scratch_2d[end,end]
+	adv_fac_lower_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[1,ivpa,ivperp,ir]
+	adv_fac_upper_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[end,ivpa,ivperp,ir]
+end
+# now reconcile element boundaries across
+# processes with large message
+if z.nelement_local < z.nelement_global
+	reconcile_element_boundaries_MPI!(dfdz,
+	 adv_fac_lower_buffer, adv_fac_upper_buffer,
+	 dfdz_lower_endpoints,dfdz_upper_endpoints,
+	 z_send_buffer, z_receive_buffer, z)
+end
+
+end
+
 #6D version for f[vz,vr,vzeta,z,r,sn] -> dfn neutral particles
 function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_float,6},
         advect, adv_fac_lower_buffer::AbstractArray{mk_float,5},
diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index 78ae5b9b8..4fe56acb7 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -5,10 +5,12 @@ export calculate_electron_upar_from_charge_conservation!
 export electron_energy_equation!
 export calculate_electron_qpar!
 export calculate_electron_parallel_friction_force!
+export calculate_electron_qpar_from_pdf!
 
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
-using ..input_structs: braginskii_fluid
+using ..input_structs: braginskii_fluid, kinetic_electrons
+using ..velocity_moments: integrate_over_vspace
 
 """
 use quasineutrality to obtain the electron density from the initial
@@ -58,7 +60,7 @@ function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens
         upar_e .= 0.0
         # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
         # is equal and opposite to the ion parallel current
-        if electron_model in (boltzmann_electron_response_with_simple_sheath, braginskii_fluid)
+        if electron_model ∈ (boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons)
             # loop over ion species, adding each species contribution to the
             # ion parallel particle flux at the boundaries in zed
             @loop_s_r is ir begin
@@ -98,7 +100,7 @@ NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
 function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt, composition,
-                                   num_diss_params)
+                                   num_diss_params, zgrid)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
     me_over_mi = composition.me_over_mi
@@ -107,9 +109,14 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     # arising from derivatives of ppar, qpar and upar
     @loop_r_z ir iz begin
         ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.electron.dppar_dz_upwind[iz,ir]
-                    + (2/3)*moments.electron.dqpar_dz[iz,ir]
-                    + (5/3)*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
+                    + moments.electron.dqpar_dz[iz,ir]
+                    + 3*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
     end
+    # @loop_r_z ir iz begin
+    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.electron.dppar_dz_upwind[iz,ir]
+    #                 + (2/3)*moments.electron.dqpar_dz[iz,ir]
+    #                 + (5/3)*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
+    # end
     # compute the contribution to the rhs of the energy equation
     # arising from artificial diffusion
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
@@ -140,11 +147,17 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
             end
         end
         if abs(collisions.ionization_electron) > 0.0
+            # @loop_s_r_z is ir iz begin
+            #     ppar[iz,ir] +=
+            #         dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
+            #         fvec.electron_ppar[iz,ir] -
+            #         (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
+            # end
             @loop_s_r_z is ir iz begin
                 ppar[iz,ir] +=
                     dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
                     fvec.electron_ppar[iz,ir] -
-                    (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
+                    fvec.electron_density[iz,ir] * collisions.ionization_energy)
             end
         end
     end
@@ -152,12 +165,13 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     calculate_electron_heat_source!(moments.electron.heat_source, fvec.electron_ppar, moments.electron.dupar_dz,
                                     fvec.density_neutral, collisions.ionization, collisions.ionization_energy,
                                     fvec.electron_density, fvec.ppar, collisions.nu_ei, composition.me_over_mi,
-                                    composition.T_wall)
+                                    composition.T_wall, zgrid)
     # add the contribution from the electron heat source                                
     @loop_r_z ir iz begin
         ppar[iz,ir] += dt * moments.electron.heat_source[iz,ir]
     end
     # enforce the parallel boundary condtion on the electron parallel pressure
+    #println("!!!NO PARALLEL BC IS BEING ENFORCED ON ELECTRON PRESSURE!!!")
     enforce_parallel_BC_on_electron_pressure!(ppar, dens_i, composition.T_wall, fvec.ppar)
     return nothing
 end
@@ -227,18 +241,22 @@ there are currently two supported options for the parallel heat flux:
 inputs:
     qpar_e = parallel electron heat flux at the previous time level
     qpar_updated = flag indicating whether qpar is updated already
+    pdf = electron pdf
     ppar_e = electron parallel pressure
     upar_e = electron parallel flow
+    vth_e = electron thermal speed
     dTe_dz = zed derivative of electron temperature
     upar_i = ion parallel flow
     nu_ei = electron-ion collision frequency
     me_over_mi = electron-to-ion mass ratio
     electron_model = choice of model for electron physics 
+    vpa = struct containing information about the vpa coordinate
 output:
     qpar_e = updated parallel electron heat flux
     qpar_updated = flag indicating that the parallel electron heat flux is updated
 """
-function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz, upar_i, nu_ei, me_over_mi, electron_model)
+function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth_e, dTe_dz, upar_i, 
+                                  nu_ei, me_over_mi, electron_model, vpa)
     # only calculate qpar_e if needs updating
     if qpar_updated == false
         @. qpar_e = 0.0
@@ -252,6 +270,9 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz,
                     qpar_e[iz,ir] -= (1/2) * 3.16 * ppar_e[iz,ir] / (me_over_mi * nu_ei) * dTe_dz[iz,ir] 
                 end
             end
+        elseif electron_model == kinetic_electrons
+            # use the modified electron pdf to calculate the electron heat flux
+            calculate_electron_qpar_from_pdf!(qpar_e, ppar_e, vth_e, pdf, vpa)
         end
     end
     # qpar has been updated
@@ -259,8 +280,20 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, ppar_e, upar_e, dTe_dz,
     return nothing
 end
 
+"""
+calculate the parallel component of the electron heat flux,
+defined as qpar = 2 * ppar * vth * int dwpa (pdf * wpa^3)
+"""
+function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
+    # specialise to 1D for now
+    ivperp = 1
+    @loop_r_z ir iz begin
+        @views qpar[iz, ir] = 2*ppar[iz,ir]*vth[iz,ir]*integrate_over_vspace(pdf[:, ivperp, iz, ir], vpa.grid.^3, vpa.wgts)
+    end
+end
+
 function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n, ionization, ionization_energy,
-                                         dens_e, ppar_i, nu_ei, me_over_mi, T_wall)
+                                         dens_e, ppar_i, nu_ei, me_over_mi, T_wall, zgrid)
     # heat_source currently only used for testing                                         
     # @loop_r_z ir iz begin
     #     heat_source[iz,ir] = (2/3) * ppar_e[iz,ir] * dupar_dz[iz,ir]
@@ -280,7 +313,10 @@ function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n,
     #     heat_source[1,ir] += 20 * 0.5 * (T_wall * dens_e[1,ir] - ppar_e[1,ir])
     #     heat_source[end,ir] += 20 * (0.5 * T_wall * dens_e[end,ir] - ppar_e[end,ir])
     # end
-    @. heat_source = 0.0
+    # Gaussian heat deposition profile, with a decay of 5 e-foldings at the ends of the domain in z
+    @loop_r_z ir iz begin
+        heat_source[iz,ir] = 50*exp(-5*(zgrid[iz]/zgrid[end])^2)
+    end
     return nothing
 end
 
diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
new file mode 100644
index 000000000..bbba05a4d
--- /dev/null
+++ b/src/electron_kinetic_equation.jl
@@ -0,0 +1,1315 @@
+module electron_kinetic_equation
+
+using LinearAlgebra
+
+export get_electron_critical_velocities
+
+using ..looping
+using ..derivatives: derivative_z!
+using ..calculus: derivative!, integral
+using ..interpolation: interpolate_to_grid_1d!
+using ..type_definitions: mk_float
+using ..array_allocation: allocate_float
+using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
+using ..electron_fluid_equations: electron_energy_equation!
+using ..electron_z_advection: electron_z_advection!
+using ..electron_vpa_advection: electron_vpa_advection!
+using ..velocity_moments: integrate_over_vspace
+
+"""
+update_electron_pdf is a function that uses the electron kinetic equation 
+to solve for the updated electron pdf
+
+The electron kinetic equation is:
+    zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa = pdf * pre_factor
+
+    INPUTS:
+    pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
+    dens = electron density
+    vthe = electron thermal speed
+    ppar = electron parallel pressure
+    ddens_dz = z-derivative of the electron density
+    dppar_dz = z-derivative of the electron parallel pressure
+    dqpar_dz = z-derivative of the electron parallel heat flux
+    dvth_dz = z-derivative of the electron thermal speed
+    z = struct containing z-coordinate information
+    vpa = struct containing vpa-coordinate information
+    z_spectral = struct containing spectral information for the z-coordinate
+    vpa_spectral = struct containing spectral information for the vpa-coordinate
+    scratch_dummy = dummy arrays to be used for temporary storage
+    dt = time step size
+    max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
+OUTPUT:
+    pdf = updated (modified) electron pdf
+"""
+function update_electron_pdf!(fvec, pdf, moments, dens, vthe, 
+    ppar, qpar, qpar_updated, 
+    phi, ddens_dz, dppar_dz, 
+    dqpar_dz, dvth_dz, z, vpa, z_spectral, 
+    vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition, 
+    num_diss_params, max_electron_pdf_iterations)
+
+    # set the method to use to solve the electron kinetic equation
+    solution_method = "artificial_time_derivative"
+    #solution_method = "shooting_method"
+    #solution_method = "picard_iteration"
+    # solve the electron kinetic equation using the specified method
+    if solution_method == "artificial_time_derivative"
+        update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
+            moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
+            z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
+            num_diss_params, max_electron_pdf_iterations)
+    elseif solution_method == "shooting_method"
+        update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar, qpar_updated, 
+            phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy, composition)
+    elseif solution_method == "picard_iteration"
+        update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz,
+            z, vpa, vpa_spectral, scratch_dummy, max_electron_pdf_iterations)
+    else
+        println("!!! invalid solution method specified !!!")
+    end
+    return nothing    
+end
+
+"""
+update_electron_pdf_with_time_advance is a function that introduces an artifical time derivative to advance 
+the electron kinetic equation until a steady-state solution is reached.
+
+The electron kinetic equation is:
+    zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa = pdf * pre_factor
+
+    INPUTS:
+    pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
+    dens = electron density
+    vthe = electron thermal speed
+    ppar = electron parallel pressure
+    ddens_dz = z-derivative of the electron density
+    dppar_dz = z-derivative of the electron parallel pressure
+    dqpar_dz = z-derivative of the electron parallel heat flux
+    dvth_dz = z-derivative of the electron thermal speed
+    z = struct containing z-coordinate information
+    vpa = struct containing vpa-coordinate information
+    z_spectral = struct containing spectral information for the z-coordinate
+    vpa_spectral = struct containing spectral information for the vpa-coordinate
+    scratch_dummy = dummy arrays to be used for temporary storage
+    dt = time step size
+    max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
+OUTPUT:
+    pdf = updated (modified) electron pdf
+"""
+function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
+    moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition,
+    z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
+    num_diss_params, max_electron_pdf_iterations)
+
+    println("TMP FOR TESTING: SETTING UPAR_I = UPAR_E = 0!!!")
+    moments.electron.upar .= 0.0
+
+    # there will be a better way of doing this
+    # store the incoming ppar
+    ppar_old = allocate_float(z.n,1)
+    ppar_old .= ppar
+
+    # create a (z,r) dimension dummy array for use in taking derivatives
+    dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
+    # create several (z) dimension dummy arrays for use in taking derivatives
+    buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
+    buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
+    buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
+    buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
+    buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
+    buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
+
+    # compute the z-derivative of the input electron parallel flow, needed for the electron kinetic equation
+    @views derivative_z!(moments.electron.dupar_dz, moments.electron.upar, buffer_r_1, buffer_r_2, buffer_r_3,
+                         buffer_r_4, z_spectral, z)
+
+    # compute the z-derivative of the input electron parallel pressure, needed for the electron kinetic equation
+    @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
+                         buffer_r_4, z_spectral, z)
+
+    #println("TMP FOR TESTING -- dens and ddens_dz artificially set!!!")
+    #@. dens = 1.0
+    #@. ddens_dz = (dppar_dz / ppar)*0.9
+
+    # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
+    # and parallel pressure
+    @. dvth_dz = 0.5 * vthe * (dppar_dz / ppar - ddens_dz / dens)
+    # update the electron thermal speed itself using the updated electron parallel pressure
+    @. vthe = sqrt(abs(2.0 * ppar / (dens * composition.me_over_mi)))
+    # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
+    @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
+                         buffer_r_4, z_spectral, z)
+
+    #dt_electron = dt * sqrt(composition.me_over_mi)
+    dt_max = 1.0
+    dt_energy = 0.001
+    time = 0.0
+
+    # define residual to point to a dummy array;
+    # to be filled with the contributions to the electron kinetic equation; i.e., d(pdf)/dt = -residual
+    residual = scratch_dummy.buffer_vpavperpzr_6
+    max_term = scratch_dummy.buffer_vpavperpzr_5
+    single_term = scratch_dummy.buffer_vpavperpzr_4
+    # initialise the number of iterations in the solution of the electron kinetic equation to be 1
+    iteration = 1
+    # initialise the electron pdf convergence flag to false
+    electron_pdf_converged = false
+    # calculate the residual of the electron kinetic equation for the initial guess of the electron pdf
+    dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vthe, ppar, 
+                                        ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+                                        z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                        num_diss_params, dt_max)
+    # open files to write the electron heat flux and pdf to file                                
+    io_upar = open("upar.txt", "w")
+    io_qpar = open("qpar.txt", "w")
+    io_ppar = open("ppar.txt", "w")
+    io_pdf = open("pdf.txt", "w")
+    io_vth = open("vth.txt", "w")
+    if !electron_pdf_converged
+        # need to exit or handle this appropriately
+        @loop_vpa ivpa begin
+            @loop_z iz begin
+                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+            end
+            println(io_pdf,"")
+        end
+        @loop_z iz begin
+            println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+            println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+            println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+            println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+        end
+        println(io_upar,"")
+        println(io_qpar,"")
+        println(io_ppar,"")
+        println(io_vth,"")
+    end
+    #stop()
+
+    io_pdf_stages = open("pdf_zright.txt", "w")
+
+    # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
+    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+    #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
+    # evolve (artificially) in time until the residual is less than the tolerance
+    while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+        # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
+        @. pdf -= dt_electron * residual
+
+        # enforce the boundary condition(s) on the electron pdf
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
+        
+        if (mod(iteration,50)==1)
+            @loop_vpa ivpa begin
+                println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " iteration: ", iteration, " flag: ", 1)
+            end
+            println(io_pdf_stages,"")
+        end
+
+        # update the time following the pdf update
+        time += dt_electron
+
+        qpar_updated = false
+        # update the electron heat flux
+        calculate_electron_qpar_from_pdf!(qpar, ppar, vthe, pdf, vpa)
+        qpar_updated = true
+        
+        # calculate the updated dqpar/dz
+        for iz ∈ 2:z.n-1
+            dqpar_dz[iz,1] = 0.5*(qpar[iz+1,1]-qpar[iz,1])/z.cell_width[iz] + 0.5*(qpar[iz]-qpar[iz-1,1])/z.cell_width[iz-1]
+        end
+        dqpar_dz[1,1] = (qpar[2,1]-qpar[1,1])/z.cell_width[1]
+        dqpar_dz[z.n,1] = (qpar[end,1]-qpar[end-1,1])/z.cell_width[end-1]
+
+        # compute the z-derivative of the parallel electron heat flux
+        #@views derivative_z!(dqpar_dz, qpar, buffer_r_1,
+        #    buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+
+        # Compute the upwinded z-derivative of the electron parallel pressure for the 
+        # electron energy equation
+        dummy_zr .= -moments.electron.upar
+        @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
+                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                             buffer_r_5, buffer_r_6, z_spectral, z)
+         # centred second derivative for dissipation
+        if num_diss_params.moment_dissipation_coefficient > 0.0
+            @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
+                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+        end
+
+        if (mod(iteration,50) == 0)
+            println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,1], " average_residual: ", average_residual)
+            @loop_z iz begin
+                println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+            end
+            println(io_upar,"")
+            println(io_qpar,"")
+            println(io_ppar,"")
+            println(io_vth,"")
+        end
+
+        dt_energy = dt_electron
+
+        # get an updated iterate of the electron parallel pressure
+        #ppar .= ppar_old
+        for i in 1:1000
+            # Compute the upwinded z-derivative of the electron parallel pressure for the 
+            # electron energy equation
+            dummy_zr .= -moments.electron.upar
+            @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
+                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                             buffer_r_5, buffer_r_6, z_spectral, z)
+            # centred second derivative for dissipation
+            if num_diss_params.moment_dissipation_coefficient > 0.0
+                @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
+                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+            end
+
+            dt_energy = dt_electron
+            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z.grid)
+            fvec.electron_ppar .= ppar
+        
+        # compute the z-derivative of the updated electron parallel pressure
+        @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
+                             buffer_r_4, z_spectral, z)
+        # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
+        # and parallel pressure
+        @. dvth_dz = 0.5 * vthe * (dppar_dz / ppar - ddens_dz / dens)
+        # update the electron thermal speed itself using the updated electron parallel pressure
+        @. vthe = sqrt(abs(2.0 * ppar / (dens * composition.me_over_mi)))
+        end
+
+        #@views derivative_z!(dqpar_dz, qpar, 
+        #        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        #        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # TMP FOR TESTING
+        #dqpar_dz .= 0.0
+        # calculate the residual of the electron kinetic equation for the updated electron pdf
+        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vthe, ppar, ddens_dz, 
+                                            dppar_dz, dqpar_dz, dvth_dz, 
+                                            z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                            num_diss_params, dt_max)
+        # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+        if electron_pdf_converged
+            break
+        end
+        iteration += 1
+    end
+    if !electron_pdf_converged
+        # need to exit or handle this appropriately
+        println("!!!max number of iterations for electron pdf update exceeded!!!")
+        @loop_vpa ivpa begin
+            @loop_z iz begin
+                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+            end
+            println(io_pdf,"")
+        end
+    end
+    close(io_upar)
+    close(io_qpar)
+    close(io_ppar)
+    close(io_vth)
+    close(io_pdf)
+    close(io_pdf_stages)
+    return nothing    
+end
+
+function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
+    # first enforce the boundary condition at z_min.
+    # this involves forcing the pdf to be zero for electron travelling faster than the max speed
+    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
+    # for electrons with positive velocities less than this critical value, they must have the same
+    # pdf as electrons with negative velocities of the same magnitude.
+    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
+    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
+
+    # pdf_adjustment_option determines the velocity-dependent pre-factor for the
+    # corrections to the pdf needed to ensure moment constraints are satisfied
+    pdf_adjustment_option = "vpa4_gaussian"
+
+    # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
+    wpa_values = vpa.scratch
+    # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
+    interpolated_pdf = vpa.scratch2
+
+    # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
+    ivpa_zero = (vpa.n+1)÷2
+
+    @loop_r ir begin
+        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+        @. wpa_values = vpa.grid - upar[1,ir] / vthe[1,ir]
+        # interpolate the pdf onto this grid
+        @views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+        # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
+        @. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
+        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+        @. wpa_values = vpa.grid + upar[1,ir] / vthe[1,ir]
+        # interpolate back onto the original wpa grid
+        @views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+        # construct wpa * pdf
+        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
+        # calculate the first moment of the normalised pdf
+        #first_vspace_moment = 0.0
+        first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
+        ivpa_max = vpa.n
+        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+        # if the first moment is positive, then the cutoff speed needs to be reduced
+        while first_vspace_moment > 0.0
+            # zero out the pdf at the current cutoff velocity
+            pdf[ivpa_max,1,1,ir] = 0.0
+            # update wpa * pdf
+            vpa.scratch3[ivpa_max] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # if first_vspace_moment > 0.0
+            ivpa_max -= 1
+            # end
+        end
+        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+        phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+        # of the modified pdf are 1, 0 and 1/2 respectively
+        # will need vpa / vthe = wpa + upar/vthe
+        @. vpa.scratch2 = vpa.grid + upar[1,ir] / vthe[1,ir]
+        # first need to calculate int dwpa pdf = zeroth_moment
+        zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
+        # calculate int dwpa wpa^2 * pdf = wpa2_moment
+        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
+        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * pdf = vpa2_moment
+        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
+        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
+        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
+        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        
+        if pdf_adjustment_option == "absvpa"
+            # calculate int dwpa |vpa| * pdf = absvpa_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
+            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            # calculate the 'B' normalisation constant
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+            # calculate the 'A' normalisation constant
+            normalisation_constant_A = (1 + normalisation_constant_B * 
+                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+            # calculate the 'C' normalisation constant
+            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                + vpa.scratch2^2 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4"
+            # calculate int dwpa vpa^4 * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                + vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4_gaussian"
+            afac = 0.1
+            bfac = 0.2
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        else
+            println("pdf_adjustment_option not recognised")
+            stop()
+        end
+
+        # smooth the pdf at the boundary
+        #for ivpa ∈ 2:ivpa_max-1
+        #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
+        #end
+    end
+
+    # next enforce the boundary condition at z_max.
+    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
+    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
+    # for electrons with negative velocities less than this critical value, they must have the same
+    # pdf as electrons with positive velocities of the same magnitude.
+    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
+    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
+    
+    # io_pdf_stages = open("pdf_stages.txt", "w")
+    # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+    # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid^2
+    # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+    # @loop_vpa ivpa begin
+    #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+    #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 0)
+    # end
+    # println(io_pdf_stages,"")
+
+    @loop_r ir begin
+        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+        @. wpa_values = vpa.grid - upar[end,ir] / vthe[end,ir]
+        # interpolate the pdf onto this grid
+        @views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+        # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
+        @. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
+        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+        @. wpa_values = vpa.grid + upar[end,ir] / vthe[end,ir]
+        # interpolate back onto the original wpa grid
+        @views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
+        # end
+        # println(io_pdf_stages,"")
+
+        # construct wpa * pdf
+        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
+        # calculate the first moment of the normalised pdf
+        first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
+        ivpa_min = 1
+        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+        # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
+        while first_vspace_moment < 0.0
+            # zero out the pdf at the current cutoff velocity
+            pdf[ivpa_min,1,end,ir] = 0.0
+            # update wpa * pdf
+            vpa.scratch3[ivpa_min] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # if first_vspace_moment < 0.0
+            ivpa_min += 1
+            # end
+        end
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
+        # end
+        # println(io_pdf_stages,"")
+
+        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+        phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
+        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+        # of the modified pdf are 1, 0 and 1/2 respectively
+        # will need vpa / vthe = wpa + upar/vthe
+        @. vpa.scratch2 = vpa.grid + upar[end,ir] / vthe[end,ir]
+        # first need to calculate int dwpa pdf = zeroth_moment
+        zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
+        # calculate int dwpa wpa^2 * pdf = wpa2_moment
+        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
+        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * pdf = vpa2_moment
+        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
+        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
+        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
+        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        
+        if pdf_adjustment_option == "absvpa"
+            # calculate int dwpa |vpa| * pdf = absvpa_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
+            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            # calculate the 'B' normalisation constant
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+            # calculate the 'A' normalisation constant
+            normalisation_constant_A = (1 + normalisation_constant_B * 
+                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+            # calculate the 'C' normalisation constant
+            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                + vpa.scratch2^2 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4"
+            # calculate int dwpa vpa^4 * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                + vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4_gaussian"
+            afac = 0.1
+            bfac = 0.2
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        else
+            println("pdf_adjustment_option not recognised")
+            stop()
+        end
+
+        # smooth the pdf at the boundary
+        #for ivpa ∈ ivpa_min+1:vpa.n-1
+        #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
+        #end
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
+        # end
+        # println(io_pdf_stages,"")
+    end
+
+    # # initialise the electron pdf for positive vpa to be the mirror reflection about vpa = 0
+    # ivpa_zero = (vpa.n+1)÷2
+    # ivpa_max = vpa.n
+    # @. pdf[ivpa_zero+1:end,:,1,:] = pdf[ivpa_zero-1:-1:1,:,1,:]
+    # # calculate the zeroth v-space moment of the normalised electron pdf:
+    # # if unity to within specified tolerance, then the boundary condition is satisfied;
+    # # otherwise, modify the cutoff velocity and repeat
+    # @loop_r ir begin
+    #     unity = 2.0
+    #     while unity > 1.0
+    #         unity = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
+    #         # if unity > 1.0, then the cutoff velocity is too high so reduce it
+    #         if unity > 1.0
+    #             pdf[ivpa_max,1,1,ir] = 0.0
+    #             ivpa_max -= 1
+    #         end
+    #     end
+    #     phi[1,ir] = vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+    # end
+    # # repeat the above procedure for the boundary at z_max
+    # @. pdf[ivpa_zero-1:-1:1,:,end,:] = pdf[ivpa_zero+1:end,:,end,:]
+    # ivpa_max = 1
+    # @loop_r ir begin
+    #     unity = 2.0
+    #     while unity > 1.0
+    #         unity = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
+    #         if unity > 1.0
+    #             pdf[ivpa_max,1,end,ir] = 0.0
+    #             ivpa_max += 1
+    #         end
+    #         phi[end,ir] = vthe[end,ir]^2 * vpa.grid[ivpa_max]^2
+    #     end
+    #     #println("unity: ", unity)
+    # end
+end
+
+
+"""
+update_electron_pdf_with_shooting_method is a function that using a shooting method to solve for the electron pdf
+
+The electron kinetic equation is:
+    zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa = pdf * pre_factor
+The shooting method is 'explicit' in z, solving
+    zdot_i * (pdf_{i+1} - pdf_{i})/dz_{i} + wpadot_{i} * d(pdf_{i})/dwpa = pdf_{i} * pre_factor_{i}
+
+    INPUTS:
+    pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
+    dens = electron density
+    vthe = electron thermal speed
+    ppar = electron parallel pressure
+    ddens_dz = z-derivative of the electron density
+    dppar_dz = z-derivative of the electron parallel pressure
+    dqpar_dz = z-derivative of the electron parallel heat flux
+    dvth_dz = z-derivative of the electron thermal speed
+    z = struct containing z-coordinate information
+    vpa = struct containing vpa-coordinate information
+    vpa_spectral = struct containing spectral information for the vpa-coordinate
+    scratch_dummy = dummy arrays to be used for temporary storage
+OUTPUT:
+    pdf = updated (modified) electron pdf
+"""
+function update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar, qpar_updated, phi, 
+    ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy, composition)
+
+    # it is convenient to create a pointer to a scratch array to store the RHS of the linear 
+    # system that includes all terms evaluated at the previous z location in the shooting method
+    rhs = scratch_dummy.buffer_vpavperpzr_1
+    # initialise the RHS to zero
+    rhs .= 0.0
+    # get critical velocities beyond which electrons are lost to the wall
+    crit_speed_zmin, crit_speed_zmax = get_electron_critical_velocities(phi, vthe, composition.me_over_mi)
+    # add the contribution to rhs from the term proportional to the pdf (rather than its derivatives)
+    add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+    # add the contribution to rhs from the wpa advection term
+    add_contribution_from_wpa_advection!(rhs, pdf, vthe, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
+    # shoot in z from incoming boundary (using sign of zdot to determine direction)
+    # pdf_{i+1} = pdf_{i} - dz_{i} * rhs_{i}
+    @loop_r_vperp ir ivperp begin
+        for ivpa ∈ 1:vpa.n
+            # deal with case of positive z-advection speed
+            if vpa.grid[ivpa] > eps(mk_float)
+                for iz ∈ 1:z.n-1
+                    pdf[ivpa, ivperp, iz + 1, ir] = pdf[ivpa, ivperp, iz, ir] - (z.cell_width[iz] * rhs[ivpa, ivperp, iz, ir] 
+                                                              / (vpa.grid[ivpa] * vthe[iz, ir]))
+                end
+            # deal with case of negative z-advection speed
+            elseif vpa.grid[ivpa] < -eps(mk_float)
+                for iz ∈ z.n:-1:2
+                    pdf[ivpa, ivperp, iz - 1, ir] = pdf[ivpa, ivperp, iz, ir] + (z.cell_width[iz-1] * rhs[ivpa, ivperp, iz, ir] 
+                                                              / (vpa.grid[ivpa] * vthe[iz, ir]))
+                end
+            # deal with case of zero z-advection speed
+            else
+                # hack for now until I figure out how to deal with the singular case
+                @. pdf[ivpa, ivperp, :, ir] = 0.5 * (pdf[ivpa+1, ivperp, :, ir] + pdf[ivpa-1, ivperp, :, ir])
+            end
+        end
+        # enforce the boundary condition at the walls
+        for ivpa ∈ 1:vpa.n
+            # find the ivpa index corresponding to the same magnitude 
+            # of vpa but with opposite sign
+            ivpamod = vpa.n - ivpa + 1
+            # deal with wall at z_min
+            if vpa.grid[ivpa] > eps(mk_float)
+                iz = 1
+                # electrons travelling sufficiently fast in the negative z-direction are lost to the wall
+                # the maximum posative z-velocity that can be had is the amount by which electrons
+                # can be accelerated by the electric field between the wall and the simulation boundary
+                if vpa.grid[ivpa] > crit_speed_zmin
+                    pdf[ivpa, ivperp, iz, ir] = 0.0
+                else
+                    pdf[ivpa, ivperp, iz, ir] = pdf[ivpamod, ivperp, iz, ir]
+                end
+            # deal with wall at z_max
+            elseif vpa.grid[ivpa] < -eps(mk_float)
+                iz = z.n
+                # electrons travelling sufficiently fast in the positive z-direction are lost to the wall
+                # the maximum negative z-velocity that can be had is the amount by which electrons
+                # can be accelerated by the electric field between the wall and the simulation boundary
+                if vpa.grid[ivpa] < crit_speed_zmax
+                    pdf[ivpa, ivperp, iz, ir] = 0.0
+                else
+                    pdf[ivpa, ivperp, iz, ir] = pdf[ivpamod, ivperp, iz, ir]
+                end
+            end
+        end
+    end
+    # the electron parallel heat flux is no longer consistent with the electron pdf
+    qpar_updated = false
+    # calculate the updated electron parallel heat flux
+    calculate_electron_qpar_from_pdf!(qpar, ppar, vthe, pdf, vpa)
+    for iz ∈ 1:z.n
+        for ivpa ∈ 1:vpa.n
+            println("z: ", z.grid[iz], " vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " qpar: ", qpar[iz, 1], 
+                " dppardz: ", dppar_dz[iz,1], " vth: ", vthe[iz,1], " ppar: ", ppar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1],
+                " dvth_dz: ", dvth_dz[iz,1], " ddens_dz: ", ddens_dz[iz,1])
+        end
+        println()
+    end
+    return nothing    
+end
+
+"""
+use Picard iteration to solve the electron kinetic equation
+
+The electron kinetic equation is:
+    zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa = pdf * pre_factor
+Picard iteration uses the previous iteration of the electron pdf to calculate the next iteration:
+    zdot * d(pdf^{i+1})/dz + wpadot^{i} * d(pdf^{i})/dwpa = pdf^{i} * prefactor^{i}
+
+INPUTS:
+    pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
+    dens = electron density
+    vthe = electron thermal speed
+    ppar = electron parallel pressure
+    ddens_dz = z-derivative of the electron density
+    dppar_dz = z-derivative of the electron parallel pressure
+    dqpar_dz = z-derivative of the electron parallel heat flux
+    dvth_dz = z-derivative of the electron thermal speed
+    z = struct containing z-coordinate information
+    vpa = struct containing vpa-coordinate information
+    z_spectral = struct containing spectral information for the z-coordinate
+    vpa_spectral = struct containing spectral information for the vpa-coordinate
+    scratch_dummy = dummy arrays to be used for temporary storage
+    max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
+OUTPUT:
+    pdf = updated (modified) electron pdf
+"""
+function update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz,
+    z, vpa, vpa_spectral, scratch_dummy, max_electron_pdf_iterations)
+
+    # it is convenient to create a pointer to a scratch array to store the RHS of the linear 
+    # system that will be solved iteratively
+    rhs = scratch_dummy.buffer_vpavperpzr_1
+    # define residual to point to a scratch array;
+    # to be filled with the difference between successive iterations of the electron pdf
+    residual = scratch_dummy.buffer_vpavperpzr_2
+    # define pdf_new to point to a scratch array;
+    # to be filled with the updated electron pdf
+    pdf_new = scratch_dummy.buffer_vpavperpzr_3
+    # initialise the electron pdf convergence flag to false
+    electron_pdf_converged = false
+    # initialise the number of iterations in the solution of the electron kinetic equation to be 1
+    iteration = 1
+    while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+        # calculate the RHS of the linear system, consisting of all terms in which the
+        # pdf is evaluated at the previous iteration.
+        
+        # initialise the RHS to zero
+        rhs .= 0.0
+        # add the contribution to rhs from the term proportional to the pdf (rather than its derivatives)
+        add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+        # add the contribution to rhs from the wpa advection term
+        #add_contribution_from_wpa_advection!(rhs, pdf, vthe, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
+        # loop over wpa locations, solving the linear system at each location;
+        # care must be taken when wpa = 0, as the linear system is singular in this case
+        @loop_r_vperp_vpa ir ivperp ivpa begin
+            # modify the rhs vector to account for the pre-factor multiplying d(pdf)/dz
+            if abs(vpa.grid[ivpa]) < eps()
+                # hack for now until I figure out how to deal with the singular linear system
+                rhs[ivpa, ivperp, :, ir] .= 0.0
+            else
+                @loop_z iz begin
+                    #rhs[ivpa, ivperp, iz, ir] /= -(vpa.grid[ivpa] * vthe[iz, ir])
+                    rhs[ivpa, ivperp, iz, ir] = -rhs[ivpa, ivperp, iz, ir]
+                end
+            end
+            # solve the linear system at this (wpa, wperp, r) location
+            pdf_new[ivpa, ivperp, :, ir] = z.differentiation_matrix \ rhs[ivpa, ivperp, :, ir]
+        end
+        # calculate the difference between successive iterations of the electron pdf
+        @. residual = pdf_new - pdf
+        # check to see if the electron pdf has converged to within the specified tolerance
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual)
+        # prepare for the next iteration by updating the pdf
+        @. pdf = pdf_new
+        # if converged, exit the loop; otherwise, increment the iteration counter
+        if electron_pdf_converged
+            break
+        else
+            iteration +=1
+        end
+    end
+    # if the maximum number of iterations has been exceeded, print a warning
+    if !electron_pdf_converged
+        # need to exit or handle this appropriately
+        println("!!! max number of iterations for electron pdf update exceeded !!!")
+        @loop_z iz begin
+            println("z: ", z.grid[iz], " pdf: ", pdf[10, 1, iz, 1])
+        end
+    end
+    
+    return nothing
+end
+
+
+"""
+electron_kinetic_equation_residual! calculates the residual of the (time-independent) electron kinetic equation
+INPUTS:
+    residual = dummy array to be filled with the residual of the electron kinetic equation
+OUTPUT:
+    residual = updated residual of the electron kinetic equation
+"""
+function electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vth, ppar, 
+                                             ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+                                             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                             num_diss_params, dt_electron)
+
+    # initialise the residual to zero                                             
+    residual .= 0.0
+    # calculate the contribution to the residual from the z advection term
+    #electron_z_advection!(residual, pdf, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
+    dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
+    single_term .= residual
+    max_term .= abs.(residual)
+    #println("z_advection: ", sum(residual), " dqpar_dz: ", sum(abs.(dqpar_dz)))
+    #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
+    # add in the contribution to the residual from the wpa advection term
+    #electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
+    #                        vpa_advect, vpa, vpa_spectral, scratch_dummy)
+    dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
+    @. single_term = residual - single_term
+    max_term .= max.(max_term, abs.(single_term))
+    @. single_term = residual
+    #add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
+    # add in the contribution to the residual from the term proportional to the pdf
+    add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+    @. single_term = residual - single_term
+    max_term .= max.(max_term, abs.(single_term))
+    @. single_term = residual
+    # @loop_vpa ivpa begin
+    #     @loop_z iz begin
+    #         println("LHS: ", residual[ivpa,1,iz,1], " vpa: ", vpa.grid[ivpa], " z: ", z.grid[iz], " dvth_dz: ", dvth_dz[iz,1], " type: ", 1) 
+    #     end
+    #     println("")
+    # end
+    # println("")
+    # add in numerical dissipation terms
+    add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
+    @. single_term = residual - single_term
+    max_term .= max.(max_term, abs.(single_term))
+    # add in particle and heat source term(s)
+    #@. single_term = residual
+    #add_source_term!(residual, vpa.grid, z.grid, dvth_dz)
+    #@. single_term = residual - single_term
+    #max_term .= max.(max_term, abs.(single_term))
+    #stop()
+    # @loop_vpa ivpa begin
+    #     @loop_z iz begin
+    #         println("total_residual: ", residual[ivpa,1,iz,1], " vpa: ", vpa.grid[ivpa], " z: ", z.grid[iz], " dvth_dz: ", dvth_dz[iz,1], " type: ", 2) 
+    #     end
+    #     println("")
+    # end
+    # stop()
+    dt_max = min(dt_max_zadv, dt_max_vadv)
+    #println("dt_max: ", dt_max, " dt_max_zadv: ", dt_max_zadv, " dt_max_vadv: ", dt_max_vadv)
+    return dt_max
+end
+
+function simple_z_advection!(advection_term, pdf, vth, z, vpa, dt_max_in)
+    dt_max = dt_max_in
+    # take the z derivative of the input pdf
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        speed = vth[iz,ir] * vpa[ivpa]
+        dt_max = min(dt_max, 0.5*z.cell_width[iz]/(max(abs(speed),1e-3)))
+        if speed > 0
+            if iz == 1
+                #dpdf_dz  = pdf[ivpa,ivperp,iz,ir]/z.cell_width[1]
+                dpdf_dz = 0.0
+            else
+                dpdf_dz = (pdf[ivpa,ivperp,iz,ir] - pdf[ivpa,ivperp,iz-1,ir])/z.cell_width[iz-1]
+            end
+        elseif speed < 0
+            if iz == z.n
+                #dpdf_dz = -pdf[ivpa,ivperp,iz,ir]/z.cell_width[z.n-1]
+                dpdf_dz = 0.0
+            else
+                dpdf_dz = (pdf[ivpa,ivperp,iz+1,ir] - pdf[ivpa,ivperp,iz,ir])/z.cell_width[iz]
+            end
+        else
+            if iz == 1
+                dpdf_dz = (pdf[ivpa,ivperp,iz+1,ir] - pdf[ivpa,ivperp,iz,ir])/z.cell_width[1]
+            elseif iz == z.n
+                dpdf_dz = (pdf[ivpa,ivperp,iz,ir] - pdf[ivpa,ivperp,iz-1,ir])/z.cell_width[z.n-1]
+            else
+                dpdf_dz = (pdf[ivpa,ivperp,iz+1,ir] - pdf[ivpa,ivperp,iz,ir])/z.cell_width[1]
+            end
+        end
+        advection_term[ivpa,ivperp,iz,ir] += speed * dpdf_dz
+    end
+    return dt_max
+end
+
+#################
+# NEED TO ADD PARTICLE/HEAT SOURCE INTO THE ELECTRON KINETIC EQUATION TO FIX amplitude
+#################
+
+
+function simple_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_max_in)
+    dt_max = dt_max_in
+    # take the vpa derivative in the input pdf
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        speed = ((vth[iz,ir] * dppar_dz[iz,ir] + vpa.grid[ivpa] * dqpar_dz[iz,ir]) 
+                / (2 * ppar[iz,ir]) - vpa.grid[ivpa]^2 * dvth_dz[iz,ir])
+        dt_max = min(dt_max, 0.5*vpa.cell_width[ivpa]/max(abs(speed),1e-3))
+        if speed > 0
+            if ivpa == 1
+                dpdf_dvpa  = pdf[ivpa,ivperp,iz,ir]/vpa.cell_width[1]
+            else
+                dpdf_dvpa = (pdf[ivpa,ivperp,iz,ir] - pdf[ivpa-1,ivperp,iz,ir])/vpa.cell_width[ivpa-1]
+            end
+        elseif speed < 0
+            if ivpa == vpa.n
+                dpdf_dvpa = -pdf[ivpa,ivperp,iz,ir]/vpa.cell_width[vpa.n-1]
+            else
+                dpdf_dvpa = (pdf[ivpa+1,ivperp,iz,ir] - pdf[ivpa,ivperp,iz,ir])/vpa.cell_width[ivpa]
+            end
+        else
+            if ivpa == 1
+                dpdf_dvpa = (pdf[ivpa+1,ivperp,iz,ir] - pdf[ivpa,ivperp,iz,ir])/vpa.cell_width[1]
+            elseif ivpa == vpa.n
+                dpdf_dvpa = (pdf[ivpa,ivperp,iz,ir] - pdf[ivpa-1,ivperp,iz,ir])/vpa.cell_width[vpa.n-1]
+            else
+                dpdf_dvpa = (pdf[ivpa+1,ivperp,iz,ir] - pdf[ivpa,ivperp,iz,ir])/vpa.cell_width[1]
+            end
+        end
+        advection_term[ivpa,ivperp,iz,ir] += speed * dpdf_dvpa
+    end
+    return dt_max
+end
+
+function add_source_term!(source_term, vpa, z, dvth_dz)
+    # add in particle and heat source term(s)
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+    #    source_term[ivpa,ivperp,iz,ir] -= 40*exp(-vpa[ivpa]^2)*exp(-z[iz]^2)
+    #    source_term[ivpa,ivperp,iz,ir] -= vpa[ivpa]*exp(-vpa[ivpa]^2)*(2*vpa[ivpa]^2-3)*dvth_dz[iz,ir]
+    end
+    return nothing
+end
+
+function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
+    dummy_zr1 = @view scratch_dummy.dummy_zrs[:,:,1]
+    dummy_zr2 = @view scratch_dummy.buffer_vpavperpzr_1[1,1,:,:]
+    buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
+    buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
+    buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
+    buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
+    # add in numerical dissipation terms
+    @loop_vperp_vpa ivperp ivpa begin
+        @views derivative_z!(dummy_zr1, pdf[ivpa,ivperp,:,:], buffer_r_1, buffer_r_2, buffer_r_3,
+                             buffer_r_4, z_spectral, z)
+        @views derivative_z!(dummy_zr2, dummy_zr1, buffer_r_1, buffer_r_2, buffer_r_3,
+                             buffer_r_4, z_spectral, z)
+        @. residual[ivpa,ivperp,:,:] -= num_diss_params.z_dissipation_coefficient * dummy_zr2
+    end
+    @loop_r_z_vperp ir iz ivperp begin
+        @views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
+        @views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
+        @. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
+    end
+    #stop()
+    return nothing
+end
+
+"""
+update_electron_pdf! iterates to find a solution for the electron pdf
+from the electron kinetic equation:
+    zdot * d(pdf^{i})/dz + wpadot^{i} * d(pdf^{i})/dwpa + wpedot^{i} * d(pdf^{i})/dwpe  = pdf^{i+1} * pre_factor^{i}
+NB: zdot, wpadot, wpedot and pre-factor contain electron fluid quantities, that will have been updated
+independently of the electron pdf;
+NB: wpadot, wpedot and pre-factor all contain the electron parallel heat flux,
+which itself depends on the electron pdf.
+
+INPUTS:
+    pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
+OUTPUT:
+    pdf = updated (modified) electron pdf
+"""
+# function update_electron_pdf!(pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+#                               max_electron_pdf_iterations, z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+#     # iterate the electron kinetic equation until the electron pdf is converged
+#     # or the specified maximum number of iterations is exceeded
+#     electron_pdf_converged = false
+#     iteration = 1
+
+#     println("pdf_update_electron_pdf: ", sum(abs.(pdf)))
+#     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+#         # calculate the contribution to the kinetic equation due to z advection
+#         # and store in the dummy array scratch_dummy.buffer_vpavperpzr_1
+#         @views calculate_contribution_from_z_advection!(scratch_dummy.buffer_vpavperpzr_1, pdf, vthe, z, vpa.grid, z_spectral, scratch_dummy)
+#         # calculate the contribution to the kinetic equation due to w_parallel advection and add to the z advection term
+#         @views add_contribution_from_wpa_advection!(scratch_dummy.buffer_vpavperpzr_1, pdf, vthe, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
+#         # calculate the pre-factor multiplying the modified electron pdf
+#         # and store in the dummy array scratch_dummy.buffer_vpavperpzr_2
+#         @views calculate_pdf_dot_prefactor!(scratch_dummy.buffer_vpavperpzr_2, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid)
+#         # update the electron pdf
+#         #pdf_new = scratch_dummy.buffer_vpavperpzr_1
+#         #println("advection_terms: ", sum(abs.(advection_terms)), "pdf_dot_prefactor: ", sum(abs.(pdf_dot_prefactor)))
+#         @. scratch_dummy.buffer_vpavperpzr_1 /= scratch_dummy.buffer_vpavperpzr_2
+#         # check to see if the electron pdf has converged to within the specified tolerance
+#         check_electron_pdf_convergence!(electron_pdf_converged, scratch_dummy.buffer_vpavperpzr_1, pdf)
+#         # prepare for the next iteration
+#         @. pdf = scratch_dummy.buffer_vpavperpzr_1
+#         iteration += 1
+#     end
+#     if !electron_pdf_converged
+#         # need to exit or handle this appropriately
+#         println("!!!max number of iterations for electron pdf update exceeded!!!")
+#     end
+#     return nothing
+# end
+
+
+
+# """
+# use the biconjugate gradient stabilized method to solve the electron kinetic equation
+# """
+# function update_electron_pdf_bicgstab!(pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+#     max_electron_pdf_iterations, z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+ 
+#     for iz in 1:z.n
+#         println("z: ", z.grid[iz], " pdf: ", pdf[10, 1, iz, 1], " iteration: ", 1, " z.n: ", z.n)
+#     end
+
+#     # set various arrays to point to a corresponding dummy array scratch_dummy.buffer_vpavperpzr_X
+#     # so that it is easier to understand what is going on
+#     residual = scratch_dummy.buffer_vpavperpzr_1
+#     residual_hat = scratch_dummy.buffer_vpavperpzr_2
+#     pvec = scratch_dummy.buffer_vpavperpzr_3
+#     vvec = scratch_dummy.buffer_vpavperpzr_4
+#     tvec = scratch_dummy.buffer_vpavperpzr_5
+
+#     # calculate the residual of the electron kinetic equation for the initial guess of the electron pdf
+#     electron_kinetic_equation_residual!(residual, pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+#                                         z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+#     @. residual = -residual
+
+#     for iz in 1:z.n
+#         println("z: ", z.grid[iz], " residual: ", residual[10, 1, iz, 1], " iteration: ", -1, " z.n: ", z.n)
+#     end
+
+#     residual_hat .= residual + pdf
+
+#     # calculate the inner product of the residual and residual_hat
+#     rho = dot(residual_hat, residual)
+#     #@views rho = integral(residual_hat[10, 1, :, 1], residual[10,1,:,1], z.wgts)
+
+#     println("rho: ", rho)
+
+#     pvec .= residual
+
+#     # iterate the electron kinetic equation until the electron pdf is converged
+#     # or the specified maximum number of iterations is exceeded
+#     electron_pdf_converged = false
+#     iteration = 1
+
+#     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+#         # calculate the residual of the electron kinetic equation with the residual of the previous iteration
+#         # as the input electron pdf
+#         electron_kinetic_equation_residual!(vvec, pvec, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+#                                             z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+#         alpha = rho / dot(residual_hat, vvec)
+#         #alpha = rho / integral(residual_hat[10,1,:,1], vvec[10,1,:,1], z.wgts)
+#         println("alpha: ", alpha)
+#         @. pdf += alpha * pvec
+#         @. residual -= alpha * vvec
+#         # check to see if the electron pdf has converged to within the specified tolerance
+#         check_electron_pdf_convergence!(electron_pdf_converged, residual)
+#         if electron_pdf_converged
+#             break
+#         end
+#         electron_kinetic_equation_residual!(tvec, residual, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
+#                                             z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+#         omega = dot(tvec, residual) / dot(tvec, tvec)
+#         #omega = integral(tvec[10,1,:,1],residual[10,1,:,1],z.wgts)/integral(tvec[10,1,:,1],tvec[10,1,:,1],z.wgts)
+#         @. pdf += omega * residual
+#         @. residual -= omega * tvec
+#         check_electron_pdf_convergence!(electron_pdf_converged, residual)
+#         if electron_pdf_converged
+#             break
+#         end
+#         #rho_new = integral(residual_hat[10,1,:,1],residual[10,1,:,1],z.wgts)
+#         rho_new = dot(residual_hat, residual)
+#         beta = (rho_new / rho) * (alpha / omega)
+#         println("rho_new: ", rho_new, " beta: ", beta)
+#         #@. pvec -= omega * vvec
+#         #@. pvec *= beta
+#         #@. pvec += residual
+#         @. pvec = residual + beta * (pvec - omega * vvec)
+#         # prepare for the next iteration
+#         rho = rho_new
+#         iteration += 1
+#         println()
+#         if iteration == 2
+#             for iz in 1:z.n
+#                 println("z: ", z.grid[iz], " pvec: ", pvec[10, 1, iz, 1], " iteration: ", -2, " z.n: ", z.n)
+#             end
+#         end
+#         for iz in 1:z.n
+#             println("z: ", z.grid[iz], " pdf: ", pdf[10, 1, iz, 1], " iteration: ", iteration, " z.n: ", z.n)
+#         end
+#     end
+#     if !electron_pdf_converged
+#         # need to exit or handle this appropriately
+#         println("!!!max number of iterations for electron pdf update exceeded!!!")
+#     end
+#     return nothing
+# end
+
+"""
+calculates the contribution to the residual of the electron kinetic equation from the z advection term:
+residual = zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa - pdf * pre_factor
+INPUTS:
+    z_advection_term = dummy array to be filled with the contribution to the residual from the z advection term
+    pdf = modified electron pdf used in the kinetic equation = (true electron pdf / dens_e) * vth_e
+    vthe = electron thermal speed
+    z = z grid
+    vpa = v_parallel grid
+    z_spectral = spectral representation of the z grid
+    scratch_dummy = dummy array to be used for temporary storage
+OUTPUT:
+    z_advection_term = updated contribution to the residual from the z advection term
+"""
+function calculate_contribution_from_z_advection!(z_advection_term, pdf, vthe, z, vpa, spectral, scratch_dummy)
+    # calculate the z-derivative of the electron pdf and store in the dummy array
+    # scratch_dummy.buffer_vpavperpzr_1
+    derivative_z!(z_advection_term, pdf,
+                  scratch_dummy.buffer_vpavperpr_1,
+                  scratch_dummy.buffer_vpavperpr_2, scratch_dummy.buffer_vpavperpr_3,
+                  scratch_dummy.buffer_vpavperpr_4, spectral, z)
+    # contribution from the z-advection term is wpa * vthe * d(pdf)/dz
+    @loop_vperp_vpa ivperp ivpa begin
+        @. z_advection_term[ivpa, ivperp, :, :] = z_advection_term[ivpa, ivperp, :, :] #* vpa[ivpa] * vthe[:, :]
+    end
+    return nothing
+end
+
+function add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz,
+                                                    vpa, vpa_spectral)
+    @loop_r_z_vperp ir iz ivperp begin
+        # calculate the wpa-derivative of the pdf and store in the scratch array vpa.scratch
+        @views derivative!(vpa.scratch, pdf[:, ivperp, iz, ir], vpa, vpa_spectral)
+        # contribution from the wpa-advection term is ( ) * d(pdf)/dwpa
+        @. residual[:, ivperp, iz, ir] += vpa.scratch[:] * (0.5 * vth[iz, ir] / ppar[iz, ir] * dppar_dz[iz, ir]
+                + 0.5 * vpa.grid[:] / ppar[iz, ir] * dqpar_dz[iz, ir] - vpa.grid[:]^2 * dvth_dz[iz, ir])
+    end
+    return nothing
+end
+
+# calculate the pre-factor multiplying the modified electron pdf
+function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa)
+    @loop_vperp_vpa ivperp ivpa begin
+        @. pdf_dot_prefactor[ivpa, ivperp, :, :] = 0.5 * dqpar_dz[:, :] / ppar[:, :] - vpa[ivpa] * vth[:, :] * 
+        (ddens_dz[:, :] / dens[:, :] + dvth_dz[:, :] / vth[:, :])
+    end
+    return nothing
+end
+
+# add contribution to the residual coming from the term proporational to the pdf
+function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z)
+    @loop_vperp_vpa ivperp ivpa begin
+        @. residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :] - vpa[ivpa] * vth[:, :] * 
+                                (ddens_dz[:, :] / dens[:, :] - dvth_dz[:, :] / vth[:, :])) * pdf[ivpa, ivperp, :, :]
+        #@. residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf[ivpa, ivperp, :, :]
+    end
+    return nothing
+end
+
+# function check_electron_pdf_convergence!(electron_pdf_converged, pdf_new, pdf)
+#     # check to see if the electron pdf has converged to within the specified tolerance
+#     # NB: the convergence criterion is based on the average relative difference between the
+#     # new and old electron pdfs
+#     println("sum(abs(pdf)): ", sum(abs.(pdf)), " sum(abs(pdf_new)): ", sum(abs.(pdf_new)))
+#     println("sum(abs(pdfdiff)): ", sum(abs.(pdf_new - pdf)) / sum(abs.(pdf)), " length: ", length(pdf))
+#     average_relative_error = sum(abs.(pdf_new - pdf)) / sum(abs.(pdf))
+#     electron_pdf_converged = (average_relative_error < 1e-3)
+#     println("average_relative_error: ", average_relative_error, " electron_pdf_converged: ", electron_pdf_converged)
+#     return nothing
+# end
+
+function get_electron_critical_velocities(phi, vthe, me_over_mi)
+    # get the critical velocities beyond which electrons are lost to the wall
+    crit_speed_zmin = sqrt(phi[1, 1] / (me_over_mi * vthe[1, 1]^2))
+    crit_speed_zmax = -sqrt(max(phi[end, 1],0.0) / (me_over_mi * vthe[end, 1]^2))
+    return crit_speed_zmin, crit_speed_zmax
+end
+
+function check_electron_pdf_convergence(residual, norm)
+    #average_residual = 0.0
+    # @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+    #     if norm[ivpa, ivperp, iz, ir] > eps(mk_float)
+    #         average_residual += abs(residual[ivpa, ivperp, iz, ir]) / norm[ivpa, ivperp, iz, ir]
+    #     else
+    #         average_residual += abs(residual[ivpa, ivperp, iz, ir])
+    #     end
+    #     average_residual /= length(residual)
+    # end
+    average_residual = sum(abs.(residual)) / sum(norm)
+    electron_pdf_converged = (average_residual < 1e-2)
+    return average_residual, electron_pdf_converged
+end
+
+function check_electron_pdf_convergence(residual)
+    average_residual = sum(abs.(residual)) / length(residual)
+    electron_pdf_converged = (average_residual < 1e-3)
+    return average_residual, electron_pdf_converged
+end
+
+end
\ No newline at end of file
diff --git a/src/electron_vpa_advection.jl b/src/electron_vpa_advection.jl
new file mode 100644
index 000000000..863cbff1c
--- /dev/null
+++ b/src/electron_vpa_advection.jl
@@ -0,0 +1,59 @@
+"""
+"""
+module electron_vpa_advection
+
+export electron_vpa_advection!
+export update_electron_speed_vpa!
+
+using ..looping
+using ..calculus: derivative!, second_derivative!
+
+"""
+calculate the wpa-advection term for the electron kinetic equation 
+= (vthe / 2 ppare * dppare/dz + wpa / 2 ppare * dqpare/dz - wpa^2 * dvthe/dz) * df/dwpa
+"""
+function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
+                                 advect, vpa, spectral, scratch_dummy)
+    # create a reference to a scratch_dummy array to store the wpa-derivative of the electron pdf
+    dpdf_dvpa = scratch_dummy.buffer_vpavperpzr_1
+    d2pdf_dvpa2 = scratch_dummy.buffer_vpavperpzr_2
+    begin_r_z_vperp_region()
+    # get the updated speed along the wpa direction using the current pdf
+    @views update_electron_speed_vpa!(advect[1], ppar[:,:], vth[:,:], dppar_dz[:,:], dqpar_dz[:,:], dvth_dz[:,:], vpa.grid)
+    # update adv_fac -- note that there is no factor of dt here because
+    # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
+    @views @. advect[1].adv_fac[:,:,:,:] = -advect[1].speed[:,:,:,:]
+    #calculate the upwind derivative of the electron pdf w.r.t. wpa
+    @loop_r_z_vperp ir iz ivperp begin
+        @views derivative!(dpdf_dvpa[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, advect[1].adv_fac[:,ivperp,iz,ir], spectral)
+    end
+    @loop_r_z_vperp ir iz ivperp begin
+        @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], 1, vpa, spectral)
+    end
+    # calculate the advection term
+    @loop_vpa ivpa begin
+        @. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:] + 0.0001*d2pdf_dvpa2[ivpa,:,:,:]
+    end
+    #@loop_vpa ivpa begin
+    #    println("electron_vpa_advection: ", advection_term[ivpa,1,10,1], " vpa: ", vpa.grid[ivpa], " dpdf_dvpa: ", dpdf_dvpa[ivpa,1,10,1],
+    #        " pdf: ", pdf[ivpa,1,10,1])
+    #end
+    #exit()
+    return nothing
+end
+
+"""
+calculate the electron advection speed in the wpa-direction at each grid point
+"""
+function update_electron_speed_vpa!(advect, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa)
+    # calculate the advection speed in wpa
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        # TMP FOR TESTING
+        #advect.speed[ivpa,ivperp,iz,ir] = vth[iz,ir] * dppar_dz[iz,ir] / (2 * ppar[iz,ir])
+        advect.speed[ivpa,ivperp,iz,ir] = ((vth[iz,ir] * dppar_dz[iz,ir] + vpa[ivpa] * dqpar_dz[iz,ir]) 
+                                           / (2 * ppar[iz,ir]) - vpa[ivpa]^2 * dvth_dz[iz,ir])
+    end
+    return nothing
+end
+
+end
diff --git a/src/electron_z_advection.jl b/src/electron_z_advection.jl
new file mode 100644
index 000000000..816ac979a
--- /dev/null
+++ b/src/electron_z_advection.jl
@@ -0,0 +1,54 @@
+"""
+"""
+module electron_z_advection
+
+export electron_z_advection!
+export update_electron_speed_z!
+
+using ..advection: advance_f_df_precomputed!
+using ..chebyshev: chebyshev_info
+using ..looping
+using ..derivatives: derivative_z!
+using ..calculus: second_derivative!, derivative!
+
+"""
+calculate the z-advection term for the electron kinetic equation = wpa * vthe * df/dz
+"""
+function electron_z_advection!(advection_term, pdf, vth, advect, z, vpa, spectral, scratch_dummy)
+    # create a pointer to a scratch_dummy array to store the z-derivative of the electron pdf
+    dpdf_dz = scratch_dummy.buffer_vpavperpzr_1
+    d2pdf_dz2 = scratch_dummy.buffer_vpavperpzr_2
+    begin_r_vperp_vpa_region()
+    # get the updated speed along the z direction using the current pdf
+    @views update_electron_speed_z!(advect[1], vth[:,:], vpa)
+    # update adv_fac -- note that there is no factor of dt here because
+    # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
+    @views @. advect[1].adv_fac[:,:,:,:] = -advect[1].speed[:,:,:,:]
+    #calculate the upwind derivative
+    derivative_z!(dpdf_dz, pdf,
+                  advect, scratch_dummy.buffer_vpavperpr_1,
+                  scratch_dummy.buffer_vpavperpr_2, scratch_dummy.buffer_vpavperpr_3,
+                  scratch_dummy.buffer_vpavperpr_4, scratch_dummy.buffer_vpavperpr_5,
+                  scratch_dummy.buffer_vpavperpr_6, spectral, z)
+    @loop_r_vperp_vpa ir ivperp ivpa begin
+        @views second_derivative!(d2pdf_dz2[ivpa,ivperp,:,ir], pdf[ivpa,ivperp,:,ir], 1, z, spectral)
+    end
+    # calculate the advection term
+    @loop_z iz begin
+        @. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:] + 0.0001*d2pdf_dz2[:,:,iz,:]
+    end
+    return nothing
+end
+
+"""
+calculate the electron advection speed in the z-direction at each grid point
+"""
+function update_electron_speed_z!(advect, vth, vpa)
+    # the electron advection speed in z is v_par = w_par * v_the
+    @loop_r_vperp_vpa ir ivperp ivpa begin
+        @. @views advect.speed[:,ivpa,ivperp,ir] = vpa[ivpa] * vth[:,ir]
+    end
+    return nothing
+end
+
+end
diff --git a/src/em_fields.jl b/src/em_fields.jl
index bb18d49f1..b9aee1f77 100644
--- a/src/em_fields.jl
+++ b/src/em_fields.jl
@@ -55,7 +55,7 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
     dens_e = fvec.electron_density
     # in serial as both s, r and z required locally
     if (composition.n_ion_species > 1 || true ||
-        composition.electron_physics == boltzmann_electron_response_with_simple_sheath)
+        composition.electron_physics ∈ (boltzmann_electron_response_with_simple_sheath))
         # If there is more than 1 ion species, the ranks that handle species 1 have to
         # read density for all the other species, so need to synchronize here.
         # If composition.electron_physics ==
@@ -101,7 +101,7 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
             @loop_r_z ir iz begin
                 fields.phi[iz,ir] = composition.T_e * log(dens_e[iz,ir] / N_e[ir])
             end
-        elseif composition.electron_physics == braginskii_fluid
+        elseif composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
             calculate_Epar_from_electron_force_balance!(fields.Ez, dens_e, moments.electron.dppar_dz,
                 collisions.nu_ei, moments.electron.parallel_friction,
                 composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
@@ -132,7 +132,7 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
         end
     end
     # if advancing electron fluid equations, solve for Ez directly from force balance
-    if composition.electron_physics != braginskii_fluid
+    if composition.electron_physics ∉ (braginskii_fluid, kinetic_electrons)
         # Ez = - d phi / dz
         @views derivative_z!(fields.Ez,-fields.phi,
                     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1],
@@ -144,7 +144,7 @@ end
 
 function calculate_phi_from_Epar!(phi, Epar, dz)
     # simple calculation of phi from Epar for now, with zero phi assumed at boundary
-    phi[1,:] .= 0.0
+    phi[1,:] .= 3.0
     @loop_r_z ir iz begin
         if iz > 1
             phi[iz,ir] = phi[iz-1,ir] - dz[iz-1]*Epar[iz,ir]
diff --git a/src/file_io.jl b/src/file_io.jl
index 1f01b41c1..683cb75fe 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -24,7 +24,7 @@ structure containing the various input/output streams
 """
 struct ascii_ios{T <: Union{IOStream,Nothing}}
     # corresponds to the ascii file to which the distribution function is written
-    #ff::T
+    ff::T
     # corresponds to the ascii file to which velocity space moments of the
     # distribution function such as density and pressure are written
     moments_ion::T
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index aba6a537c..75248305a 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -23,6 +23,7 @@ using ..calculus: reconcile_element_boundaries_MPI!
 using ..coordinates: coordinate
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
+using ..em_fields: update_phi!
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
@@ -37,7 +38,8 @@ using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
-using ..input_structs: boltzmann_electron_response_with_simple_sheath
+using ..electron_kinetic_equation: update_electron_pdf!, get_electron_critical_velocities
+using ..input_structs: boltzmann_electron_response_with_simple_sheath, kinetic_electrons
 using ..derivatives: derivative_z!
 
 using ..manufactured_solns: manufactured_solutions
@@ -55,8 +57,8 @@ end
 struct pdf_struct
     #ion particles: s + r + z + vperp + vpa
     ion::pdf_substruct{5}
-    # electron particles: 1 + r + z + vperp + vpa
-    electron::pdf_substruct{5}
+    # electron particles: r + z + vperp + vpa
+    electron::pdf_substruct{4}
     #neutral particles: s + r + z + vzeta + vr + vz
     neutral::pdf_substruct{6}
 end
@@ -92,7 +94,8 @@ end
 Creates the structs for the pdf and the velocity-space moments
 """
 function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
-                                  evolve_moments, collisions, numerical_dissipation)
+                                  evolve_moments, collisions, numerical_dissipation, t_input)
+    # create the ion, electron and neutral pdf structs                             
     pdf = create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
 
     # create the 'moments' struct that contains various v-space moments and other
@@ -127,7 +130,54 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
     boundary_distributions = create_boundary_distributions(vz, vr, vzeta, vpa, vperp, z,
                                                            composition)
 
-    return pdf, moments, boundary_distributions
+    # create an array of structs containing scratch arrays for the pdf and low-order moments
+    # that may be evolved separately via fluid equations
+    scratch = allocate_scratch_arrays(moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+
+    return pdf, moments, boundary_distributions, scratch
+end
+
+function allocate_scratch_arrays(moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
+    # create n_rk_stages+1 structs, each of which will contain one pdf,
+    # one density, and one parallel flow array
+    scratch = Vector{scratch_pdf{5,3,4,2,6,3}}(undef, n_rk_stages+1)
+    pdf_dims = size(pdf_ion_in)
+    moment_ion_dims = size(moments.ion.dens)
+    pdf_electron_dims = size(pdf_electron_in)
+    moment_electron_dims = size(moments.electron.dens)
+    pdf_neutral_dims = size(pdf_neutral_in)
+    moment_neutral_dims = size(moments.neutral.dens)
+    # populate each of the structs
+    for istage ∈ 1:n_rk_stages+1
+        # Allocate arrays in temporary variables so that we can identify them
+        # by source line when using @debug_shared_array
+
+        # these are the pdf and moment arrays for the ion species
+        pdf_array = allocate_shared_float(pdf_dims...)
+        density_array = allocate_shared_float(moment_ion_dims...)
+        upar_array = allocate_shared_float(moment_ion_dims...)
+        ppar_array = allocate_shared_float(moment_ion_dims...)
+        temp_z_s_array = allocate_shared_float(moment_ion_dims...)
+        # these are the pdf and moment arrays for the electron species
+        pdf_electron_array = allocate_shared_float(pdf_electron_dims...)
+        electron_density_array = allocate_shared_float(moment_electron_dims...)
+        electron_upar_array = allocate_shared_float(moment_electron_dims...)
+        electron_ppar_array = allocate_shared_float(moment_electron_dims...)
+        electron_temp_array = allocate_shared_float(moment_electron_dims...)
+        # these are the pdf and moment arrays for the neutral species
+        pdf_neutral_array = allocate_shared_float(pdf_neutral_dims...)
+        density_neutral_array = allocate_shared_float(moment_neutral_dims...)
+        uz_neutral_array = allocate_shared_float(moment_neutral_dims...)
+        pz_neutral_array = allocate_shared_float(moment_neutral_dims...)
+        # construct the (uninitialized) scratch struct for this stage in the RK solve
+        scratch[istage] = scratch_pdf(pdf_array, density_array, upar_array,
+                                      ppar_array, temp_z_s_array, pdf_electron_array,
+                                      electron_density_array, electron_upar_array,
+                                      electron_ppar_array, electron_temp_array,
+                                      pdf_neutral_array, density_neutral_array,
+                                      uz_neutral_array, pz_neutral_array)
+    end
+    return scratch
 end
 
 """
@@ -139,10 +189,10 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     pdf_ion_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
-    pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, 1)
+    pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
     # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
     # to emulate ion and neutral behaviour
-    pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, 1)
+    pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
 
     return pdf_struct(pdf_substruct(pdf_ion_norm, pdf_ion_buffer),
                       pdf_substruct(pdf_electron_norm, pdf_electron_buffer),
@@ -151,12 +201,13 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
 end
 
 """
-creates the normalised pdf and the velocity-space moments and populates them
+creates the normalised pdfs and the velocity-space moments and populates them
 with a self-consistent initial condition
 """
-function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition, r, z,
-                               vperp, vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral,
-                               species, collisions, use_manufactured_solns)
+function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, composition, r, z,
+                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, vz_spectral,
+                               species, collisions, use_manufactured_solns,
+                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs)
     if use_manufactured_solns
         init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z,
                                              r, n_ion_species, n_neutral_species,
@@ -165,41 +216,47 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
         n_ion_species = composition.n_ion_species
         n_neutral_species = composition.n_neutral_species
         @serial_region begin
-            # initialise the density profile
+            # initialise the ion density profile
             init_density!(moments.ion.dens, z, r, species.ion, n_ion_species)
-            # initialise the parallel flow profile
+            # initialise the ion parallel flow profile
             init_upar!(moments.ion.upar, z, r, species.ion, n_ion_species)
-            # initialise the parallel thermal speed profile
+            # initialise the ion parallel thermal speed profile
             init_vth!(moments.ion.vth, z, r, species.ion, n_ion_species)
             @. moments.ion.ppar = 0.5 * moments.ion.dens * moments.ion.vth^2
-
+            # initialise the fluid moments for neutrals, if present
             if n_neutral_species > 0
-                #neutral particles
+                # initialise the neutral density profile
                 init_density!(moments.neutral.dens, z, r, species.neutral, n_neutral_species)
+                # initialise the z-component of the neutral flow
                 init_uz!(moments.neutral.uz, z, r, species.neutral, n_neutral_species)
+                # initialise the r-component of the neutral flow
                 init_ur!(moments.neutral.ur, z, r, species.neutral, n_neutral_species)
+                # initialise the zeta-component of the neutral flow
                 init_uzeta!(moments.neutral.uzeta, z, r, species.neutral, n_neutral_species)
+                # initialise the neutral thermal speed
                 init_vth!(moments.neutral.vth, z, r, species.neutral, n_neutral_species)
+                # calculate the z-component of the neutral pressure
                 @. moments.neutral.pz = 0.5 * moments.neutral.dens * moments.neutral.vth^2
+                # calculate the total neutral pressure
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
             end
         end
+        # reflect the fact that the ion moments have now been updated
         moments.ion.dens_updated .= true
         moments.ion.upar_updated .= true
         moments.ion.ppar_updated .= true
-        # moments.electron.dens_updated = true
-        # moments.electron.upar_updated = true
-        # moments.electron.ppar_updated = true
+        # account for the fact that the neutral moments have now been updated
         moments.neutral.dens_updated .= true
         moments.neutral.uz_updated .= true
         moments.neutral.pz_updated .= true
-        # create and initialise the normalised particle distribution function (pdf)
+        # create and initialise the normalised, ion particle distribution function (pdf)
         # such that ∫dwpa pdf.norm = 1, ∫dwpa wpa * pdf.norm = 0, and ∫dwpa wpa^2 * pdf.norm = 1/2
         # note that wpa = vpa - upar, unless moments.evolve_ppar = true, in which case wpa = (vpa - upar)/vth
         # the definition of pdf.norm changes accordingly from pdf_unnorm / density to pdf_unnorm * vth / density
         # when evolve_ppar = true.
         initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
-                        vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
+                        vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral, species,
+                        scratch_dummy)
 
         begin_s_r_z_region()
         # calculate the initial parallel heat flux from the initial un-normalised pdf
@@ -229,44 +286,107 @@ function init_pdf_and_moments!(pdf, moments, boundary_distributions, composition
                                  z, r, composition)
 
     @serial_region begin 
+        moments.electron.dens_updated = false
         # initialise the electron density profile
         init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
         # initialise the electron parallel flow profile
         init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
             moments.ion.upar, moments.ion.dens, composition.electron_physics)
         # initialise the electron thermal speed profile
-        init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e)
+        init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
         # calculate the electron temperature from the thermal speed
-        @. moments.electron.temp = moments.electron.vth^2
+        @. moments.electron.temp = composition.me_over_mi * moments.electron.vth^2
         # the electron temperature has now been updated
         moments.electron.temp_updated = true
         # calculate the electron parallel pressure from the density and temperature
         @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
         # the electron parallel pressure now been updated
         moments.electron.ppar_updated = true
+
+        # calculate the zed derivative of the initial electron density
+        @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         # calculate the zed derivative of the initial electron temperature
-        # NB: ddens_dz, dupar_dz, dppar_dz and dqpar_dz are used as dummy arrays in this function
-        # NB: call, as they have yet to be calculated
         @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
-            moments.electron.ddens_dz[:,1], moments.electron.dupar_dz[:,1], moments.electron.dppar_dz[:,1],
-            moments.electron.dqpar_dz[:,1], z_spectral, z)
-        # calculate the electron parallel heat flux
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, moments.electron.ppar,
-            moments.electron.upar, moments.electron.dT_dz, moments.ion.upar, 
-            collisions.nu_ei, composition.me_over_mi, composition.electron_physics)
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # calculate the zed derivative of the initial electron thermal speed
+        @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # calculate the zed derivative of the initial electron parallel pressure
+        @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # calculate the electron parallel heat flux;
+        # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+            moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
+            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+        # calculate the zed derivative of the initial electron parallel heat flux
+        @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         # calculate the electron-ion parallel friction force
         calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
             moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
             composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
+        
+        # initialize the scratch arrays containing pdfs and moments for the first RK stage
+        # the electron pdf is yet to be initialised but with the current code logic, the scratch
+        # arrays need to exist and be otherwise initialised in order to compute the initial electron pdf
+        initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+        # get the initial electrostatic potential and parallel electric field
+        update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+        # initialize the electron pdf that satisfies the electron kinetic equation
+        initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, z, vpa, vperp, z_spectral, vpa_spectral, 
+                                 advection_structs.electron_z_advect, advection_structs.electron_vpa_advect,
+                                 scratch_dummy, collisions, composition, 
+                                 num_diss_params, t_input.dt)
+        # re-initialize the scratch arrays now that the electron pdf has been initialised
+        initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
     end
+
+    begin_s_r_z_region()
     
     return nothing
 end
 
+"""
+initialize the array of structs containing scratch arrays for the normalised pdf and low-order moments
+that may be evolved separately via fluid equations
+"""
+function initialize_scratch_arrays!(scratch, moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
+    # populate each of the structs
+    for istage ∈ 1:n_rk_stages+1
+        @serial_region begin
+            # initialise the scratch arrays for the ion pdf and velocity moments
+            scratch[istage].pdf .= pdf_ion_in
+            scratch[istage].density .= moments.ion.dens
+            scratch[istage].upar .= moments.ion.upar
+            scratch[istage].ppar .= moments.ion.ppar
+            # initialise the scratch arrays for the electron pdf and velocity moments
+            scratch[istage].pdf_electron .= pdf_electron_in
+            scratch[istage].electron_density .= moments.electron.dens
+            scratch[istage].electron_upar .= moments.electron.upar
+            scratch[istage].electron_ppar .= moments.electron.ppar
+            scratch[istage].electron_temp .= moments.electron.temp
+            # initialise the scratch arrays for the neutral velocity moments and pdf
+            scratch[istage].pdf_neutral .= pdf_neutral_in
+            scratch[istage].density_neutral .= moments.neutral.dens
+            scratch[istage].uz_neutral .= moments.neutral.uz
+            scratch[istage].pz_neutral .= moments.neutral.pz
+        end
+    end
+    return nothing
+end
+
 """
 """
 function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
-                         vpa, vzeta, vr, vz, vpa_spectral, vz_spectral, species)
+                         vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral, species,
+                         scratch_dummy)
     wall_flux_0 = allocate_float(r.n, composition.n_ion_species)
     wall_flux_L = allocate_float(r.n, composition.n_ion_species)
 
@@ -328,9 +448,68 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
         end
     end
 
+    # @serial_region begin
+    #     @loop_r ir begin
+    #         # this is the initial guess for the electron pdf
+    #         # it will be iteratively updated to satisfy the time-independent
+    #         # electron kinetic equation
+    #         @views init_electron_pdf_over_density!(pdf.electron.norm[:,:,:,ir], moments.electron.dens[:,ir],
+    #             moments.electron.upar[:,ir], moments.electron.vth[:,ir], z, vpa, vperp)
+    #     end
+    #     # now that we have our initial guess for the electron pdf, we iterate
+    #     # using the time-independent electron kinetic equation to find a self-consistent
+    #     # solution for the electron pdf
+    #     max_electron_pdf_iterations = 100
+    #     @views update_electron_pdf!(pdf.electron.norm, moments.electron.dens, moments.electron.vth, moments.electron.ppar, 
+    #                                 moments.electron.ddens_dz, moments.electron.dppar_dz, moments.electron.dqpar_dz, moments.electron.dvth_dz,
+    #                                 max_electron_pdf_iterations, z, vpa, z_spectral, vpa_spectral, scratch_dummy)
+    # end
+
+
+
     return nothing
 end
 
+function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spectral, vpa_spectral, z_advect, vpa_advect,
+                                  scratch_dummy, collisions, composition, num_diss_params, dt)
+    @serial_region begin
+        @loop_r ir begin
+            # this is the initial guess for the electron pdf
+            # it will be iteratively updated to satisfy the time-independent
+            # electron kinetic equation
+            @views init_electron_pdf_over_density!(pdf.electron.norm[:,:,:,ir], moments.electron.dens[:,ir],
+                moments.electron.upar[:,ir], moments.electron.vth[:,ir], phi[:,ir], z, vpa, vperp, composition.me_over_mi)
+        end
+        # update the electron pdf in the fvec struct
+        fvec.pdf_electron .= pdf.electron.norm
+        # now that the initial electron pdf is given, the electron parallel heat flux should be updated
+        # if using kinetic electrons
+        if composition.electron_physics == kinetic_electrons
+            moments.electron.qpar_updated = false
+            calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+                moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
+                moments.electron.dT_dz, moments.ion.upar, 
+                collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+            # update dqpar/dz for electrons
+            # calculate the zed derivative of the initial electron parallel heat flux
+            @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
+                scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+                scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+            # now that we have our initial guess for the electron pdf, we iterate
+            # using the time-independent electron kinetic equation to find a self-consistent
+            # solution for the electron pdf
+            #max_electron_pdf_iterations = 500000
+            max_electron_pdf_iterations = 10000
+            @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
+                                        moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
+                                        phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
+                                        moments.electron.dqpar_dz, moments.electron.dvth_dz, z, vpa, z_spectral, 
+                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
+                                        num_diss_params, max_electron_pdf_iterations)
+        end
+    end
+end
+
 """
 for now the only initialisation option for the temperature is constant in z
 returns vth0 = sqrt(2Ts/ms) / sqrt(2Te/ms) = sqrt(Ts/Te)
@@ -357,10 +536,12 @@ end
 """
 function init_density!(dens, z, r, spec, n_species)
     for is ∈ 1:n_species
+        println("init_option: ", spec[is].z_IC.initialization_option)
         for ir ∈ 1:r.n
             if spec[is].z_IC.initialization_option == "gaussian"
                 # initial condition is an unshifted Gaussian
                 @. dens[:,ir,is] = spec[is].initial_density + exp(-(z.grid/spec[is].z_IC.width)^2)
+                println("ion_dens: ", dens[1,1,is], " init: ", spec[is].initial_density, " sum_factor: ", exp(-(z.grid[1]/spec[is].z_IC.width)^2))
             elseif spec[is].z_IC.initialization_option == "sinusoid"
                 # initial condition is sinusoid in z
                 @. dens[:,ir,is] =
@@ -469,12 +650,13 @@ initialise the electron thermal speed profile.
 for now the only initialisation option for the temperature is constant in z.
 returns vth0 = sqrt(2*Ts/Te)
 """
-function init_electron_vth!(vth_e, vth_i, T_e)
+function init_electron_vth!(vth_e, vth_i, T_e, me_over_mi, z)
     # @loop_r_z ir iz begin
     #     vth_e[iz,ir] = sqrt(T_e)
     # end
     @loop_r_z ir iz begin
-        vth_e[iz,ir] = vth_i[iz,ir,1]
+        vth_e[iz,ir] = vth_i[iz,ir,1] / sqrt(me_over_mi)
+        #vth_e[iz,ir] = exp(-5*(z[iz]/z[end])^2)/sqrt(me_over_mi)
     end
 end
 
@@ -861,6 +1043,44 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
     return nothing
 end
 
+"""
+init_electron_pdf_over_density initialises the normalised electron pdf = pdf_e * vth_e / dens_e;
+care is taken to ensure that the parallel boundary condition is satisfied;
+NB: as the electron pdf is obtained via a time-independent equation,
+this 'initital' value for the electron will just be the first guess in an iterative solution
+"""
+function init_electron_pdf_over_density!(pdf, density, upar, vth, phi, z, vpa, vperp, me_over_mi)
+    if z.bc == "wall"
+        # get critical velocities beyond which electrons are lost to the wall
+        vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi)
+        println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
+        # loop over all z values on this rank, initialising a shifted Maxwellian velocity distribution
+        sharp_fac = 10
+        blend_fac = 100
+        @loop_z_vperp iz ivperp begin
+            #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
+            #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
+            @. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2) * (
+                                  (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
+                                  tanh(sharp_fac*(vpa.grid[:]-vpa_crit_zmin))) *
+                                  (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
+                                  tanh(-sharp_fac*(vpa.grid[:]-vpa_crit_zmax)))) #/
+                                  #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
+                                  #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
+#                                   exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
+                                   #                                      exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
+
+            # ensure that the normalised electron pdf integrates to unity
+            norm_factor = integrate_over_vspace(pdf[:,ivperp,iz], vpa.wgts)
+            @. pdf[:,ivperp,iz] /= norm_factor
+            #println("TMP FOR TESTING -- init electron pdf")
+            #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
+        end
+    else
+        println("!!! currently, only the wall BC is supported for kinetic electrons !!!")
+    end
+end
+
 function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z, r, n_ion_species, n_neutral_species, geometry,composition)
     manufactured_solns_list = manufactured_solutions(r.L,z.L,r.bc,z.bc,geometry,composition,r.n)
     dfni_func = manufactured_solns_list.dfni_func
diff --git a/src/input_structs.jl b/src/input_structs.jl
index 1858171cb..0dd6982b6 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -106,11 +106,13 @@ end
     boltzmann_electron_response 
     boltzmann_electron_response_with_simple_sheath 
     braginskii_fluid
+    kinetic_electrons
 end
 export electron_physics_type
 export boltzmann_electron_response
 export boltzmann_electron_response_with_simple_sheath
 export braginskii_fluid
+export kinetic_electrons
 
 """
 """
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index b29ee2aeb..387dc71b0 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -33,16 +33,19 @@ include("electron_fluid_equations.jl")
 include("em_fields.jl")
 include("bgk.jl")
 include("manufactured_solns.jl") # MRH Here?
-include("initial_conditions.jl")
-include("moment_constraints.jl")
 include("advection.jl")
 include("vpa_advection.jl")
 include("z_advection.jl")
 include("r_advection.jl")
 include("vperp_advection.jl")
+include("electron_z_advection.jl")
+include("electron_vpa_advection.jl")
 include("neutral_r_advection.jl")
 include("neutral_z_advection.jl")
 include("neutral_vz_advection.jl")
+include("electron_kinetic_equation.jl")
+include("initial_conditions.jl")
+include("moment_constraints.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("continuity.jl")
@@ -81,6 +84,8 @@ using .moment_kinetics_input: mk_input, read_input_file, run_type, performance_t
 using .time_advance: setup_time_advance!, time_advance!
 using .type_definitions: mk_int
 using .em_fields: setup_em_fields
+using .time_advance: setup_dummy_and_buffer_arrays
+using .time_advance: allocate_advection_structs
 
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
 
@@ -353,23 +358,34 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
             vperp=vperp.n, vpa=vpa.n, vzeta=vzeta.n, vr=vr.n, vz=vz.n)
     end
 
-    # Allocate arrays and create the pdf and moments structs
-    pdf, moments, boundary_distributions =
-        allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
-                                 evolve_moments, collisions, num_diss_params)
-
     # create the "fields" structure that contains arrays
     # for the electrostatic potential phi (and eventually the electromagnetic fields)
     fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude, 
                              drive_input.frequency, drive_input.force_Er_zero_at_wall)
 
+    # Allocate arrays and create the pdf and moments structs
+    pdf, moments, boundary_distributions, scratch =
+        allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
+                                 evolve_moments, collisions, num_diss_params, t_input)
+
+    # create structs containing the information needed to treat advection in z, r, vpa, vperp, and vz
+    # for ions, electrons and neutrals
+    # NB: the returned advection_structs are yet to be initialized
+    advection_structs = allocate_advection_structs(composition, z, r, vpa, vperp, vz, vr, vzeta)
+
+    # setup dummy arrays & buffer arrays for z r MPI                             
+    n_neutral_species_alloc = max(1, composition.n_neutral_species)
+    scratch_dummy = setup_dummy_and_buffer_arrays(r.n, z.n, vpa.n, vperp.n, vz.n, vr.n, vzeta.n, 
+        composition.n_ion_species, n_neutral_species_alloc)
+
     if restart_prefix_iblock === nothing
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
-        init_pdf_and_moments!(pdf, moments, boundary_distributions, composition, r, z,
-                              vperp, vpa, vzeta, vr, vz, z_spectral, vpa_spectral, vz_spectral,
-                              species, collisions, t_input.use_manufactured_solns_for_init)
+        init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, composition, r, z,
+                              vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, vz_spectral,
+                              species, collisions, t_input.use_manufactured_solns_for_init,
+                              scratch_dummy, scratch, t_input, num_diss_params, advection_structs)
         # initialize time variable
         code_time = 0.
     else
@@ -385,12 +401,11 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
-    moments, spectral_objects, advect_objects,
-    scratch, advance, scratch_dummy, manufactured_source_list =
-        setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
+    moments, spectral_objects, advance, manufactured_source_list =
+        setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral,
             r_spectral, composition, moments, fields, t_input, collisions,
-            geometry, boundary_distributions, num_diss_params, restarting)
+            geometry, boundary_distributions, num_diss_params, restarting, scratch_dummy)
     # setup i/o
     ascii_io, io_moments, io_dfns = setup_file_io(io_input, boundary_distributions, vz,
         vr, vzeta, vpa, vperp, z, r, composition, collisions, moments.evolve_density,
@@ -408,7 +423,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
     begin_s_r_z_vperp_region()
 
     return pdf, scratch, code_time, t_input, vz, vr, vzeta, vpa, vperp, gyrophase, z, r,
-           moments, fields, spectral_objects, advect_objects,
+           moments, fields, spectral_objects, advection_structs,
            composition, collisions, geometry, boundary_distributions,
            num_diss_params, advance, scratch_dummy, manufactured_source_list,
            ascii_io, io_moments, io_dfns
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 11385b323..dd0da24aa 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -768,11 +768,10 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     #############################################################################
     # define default values and create corresponding mutable structs holding
     # information about the composition of the species and their initial conditions
-    if electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath, 
-                           braginskii_fluid)
+    if electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
         n_species = n_ion_species + n_neutral_species
     else
-        n_species = n_ion_speces + n_neutral_species + 1
+        n_species = n_ion_species + n_neutral_species + 1
     end
     use_test_neutral_wall_pdf = false
     # electron temperature over reference temperature
diff --git a/src/moment_kinetics_structs.jl b/src/moment_kinetics_structs.jl
index aed983108..7aa8b3c7c 100644
--- a/src/moment_kinetics_structs.jl
+++ b/src/moment_kinetics_structs.jl
@@ -10,7 +10,9 @@ using ..type_definitions: mk_float
 
 """
 """
-struct scratch_pdf{n_distribution_ion, n_moment, n_moment_electron, n_distribution_neutral, n_moment_neutral}
+struct scratch_pdf{n_distribution_ion, n_moment, 
+                   n_distribution_electron, n_moment_electron, 
+                   n_distribution_neutral, n_moment_neutral}
     # ions
     pdf::MPISharedArray{mk_float, n_distribution_ion}
     density::MPISharedArray{mk_float, n_moment}
@@ -18,6 +20,7 @@ struct scratch_pdf{n_distribution_ion, n_moment, n_moment_electron, n_distributi
     ppar::MPISharedArray{mk_float, n_moment}
     temp_z_s::MPISharedArray{mk_float, n_moment}
     # electrons
+    pdf_electron::MPISharedArray{mk_float, n_distribution_electron}
     electron_density::MPISharedArray{mk_float, n_moment_electron}
     electron_upar::MPISharedArray{mk_float, n_moment_electron}
     electron_ppar::MPISharedArray{mk_float, n_moment_electron}
diff --git a/src/post_processing.jl b/src/post_processing.jl
index a507638a6..8239c2b2b 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -980,9 +980,13 @@ function plot_1D_1V_diagnostics(run_names, nwrite_movie, itime_min, itime_max,
 
     n_runs = length(run_names)
 
+    for iz in 1:size(z.grid)
+        println("z: ", z[iz], " pz: ", neutral_pz[iz,1,1,end])
+    end
+
     # plot_unnormalised() requires PyPlot, so ensure it is used for all plots for
     # consistency
-    pyplot()
+    #pyplot()
 
     # analyze the fields data
     phi_fldline_avg, delta_phi = get_tuple_of_return_values(analyze_fields_data, phi,
@@ -1230,6 +1234,33 @@ function plot_electron_moments(density, parallel_flow, parallel_pressure,
     if pp.plot_qpar0_vs_t
         plot_single_moment_z0(parallel_heat_flux, time, iz0, prefix, legend, run_names, moment_string)
     end
+    if pp.animate_phi_vs_z
+        pmin = minimum(minimum(p) for p ∈ parallel_pressure)
+        pmax = maximum(maximum(p) for p ∈ parallel_pressure)
+        # make a gif animation of vthe(z) at different times
+        anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+            plot(legend=legend)
+            for (t, this_z, p, run_label) ∈ zip(time, z, parallel_pressure, run_names)
+                @views plot!(this_z.grid, p[:,i], xlabel="z", ylabel="ppar",
+                             ylims=(pmin, pmax), label=run_label)
+            end
+        end
+        outfile = string(prefix, "_electron_ppar_vs_z.gif")
+        gif(anim, outfile, fps=5)
+
+        pmin = minimum(minimum(p) for p ∈ thermal_speed)
+        pmax = maximum(maximum(p) for p ∈ thermal_speed)
+        # make a gif animation of vthe(z) at different times
+        anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
+            plot(legend=legend)
+            for (t, this_z, p, run_label) ∈ zip(time, z, thermal_speed, run_names)
+                @views plot!(this_z.grid, p[:,i], xlabel="z", ylabel="vth",
+                             ylims=(pmin, pmax), label=run_label)
+            end
+        end
+        outfile = string(prefix, "_electron_vth_vs_z.gif")
+        gif(anim, outfile, fps=5)
+    end
     # if pp.plot_qpar_vs_z_t
     #     plot_single_moment_vs_z_t(parallel_heat_flux, time, z, run_names, prefix, moment_string)
     # end
@@ -1608,6 +1639,8 @@ function plot_dfns(density, parallel_flow, parallel_pressure, thermal_speed, ff,
             gif(anim, outfile, fps=5)
             # make a gif animation of f(vpa,z,t)
             anim = @animate for i ∈ itime_min_pdfs:nwrite_movie_pdfs:itime_max_pdfs
+                println("i: ", i, " itime_min_pdfs: ", itime_min_pdfs, " nwrite_movie_pdfs: ", nwrite_movie_pdfs,
+                        " itime_max_pdfs: ", itime_max_pdfs)
                 #heatmap(z, vpa, log.(abs.(ff[:,:,i])), xlabel="z", ylabel="vpa", clims = (fmin,fmax), c = :deep)
                 subplots = (@views heatmap(this_z.grid, this_vpa.grid, f[:,:,is,i], xlabel="z",
                                            ylabel="vpa", c = :deep,
diff --git a/src/post_processing_input.jl b/src/post_processing_input.jl
index f63a1654f..e561b5a51 100644
--- a/src/post_processing_input.jl
+++ b/src/post_processing_input.jl
@@ -36,13 +36,13 @@ const plot_ppar_vs_z_t = false
 # if plot_qpar_vs_z_t = true, create heatmap of species parallel heat flux vs z and time
 const plot_qpar_vs_z_t = false
 # if animate_dens_vs_z = true, create animation of species density(z) at different time slices
-const animate_dens_vs_z =  false #ttrue
+const animate_dens_vs_z =  true #ttrue
 # if animate_upar_vs_z = true, create animation of species parallel flow(z) at different time slices
-const animate_upar_vs_z = false
+const animate_upar_vs_z = true
 # 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
+const animate_vth_vs_z = true
 # 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 plot_f_unnormalized_vs_vpa_z = true, create heatmap of f_unnorm(v_parallel_unnorm,z) at
diff --git a/src/time_advance.jl b/src/time_advance.jl
index d98ebb0f1..9de0da292 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -4,6 +4,7 @@ module time_advance
 
 export setup_time_advance!
 export time_advance!
+export allocate_advection_structs
 
 using MPI
 using ..type_definitions: mk_float
@@ -132,6 +133,19 @@ struct scratch_dummy_arrays
     buffer_vzvrvzetazrsn_1::MPISharedArray{mk_float,6}
     buffer_vzvrvzetazrsn_2::MPISharedArray{mk_float,6}
 
+    buffer_vpavperpzr_1::MPISharedArray{mk_float,4}
+    buffer_vpavperpzr_2::MPISharedArray{mk_float,4}
+    buffer_vpavperpzr_3::MPISharedArray{mk_float,4}
+    buffer_vpavperpzr_4::MPISharedArray{mk_float,4}
+    buffer_vpavperpzr_5::MPISharedArray{mk_float,4}
+    buffer_vpavperpzr_6::MPISharedArray{mk_float,4}
+
+    buffer_vpavperpr_1::MPISharedArray{mk_float,3}
+    buffer_vpavperpr_2::MPISharedArray{mk_float,3}
+    buffer_vpavperpr_3::MPISharedArray{mk_float,3}
+    buffer_vpavperpr_4::MPISharedArray{mk_float,3}
+    buffer_vpavperpr_5::MPISharedArray{mk_float,3}
+    buffer_vpavperpr_6::MPISharedArray{mk_float,3}
 end 
 
 struct advect_object_struct
@@ -139,6 +153,8 @@ struct advect_object_struct
     vperp_advect::Vector{advection_info{4,5}}
     z_advect::Vector{advection_info{4,5}}
     r_advect::Vector{advection_info{4,5}}
+    electron_z_advect::Vector{advection_info{4,5}}
+    electron_vpa_advect::Vector{advection_info{4,5}}
     neutral_z_advect::Vector{advection_info{5,6}}
     neutral_r_advect::Vector{advection_info{5,6}}
     neutral_vz_advect::Vector{advection_info{5,6}}
@@ -156,6 +172,68 @@ struct spectral_object_struct{Tvz,Tvr,Tvzeta,Tvpa,Tvperp,Tz,Tr}
     r_spectral::Tr
 end
 
+function allocate_advection_structs(composition, z, r, vpa, vperp, vz, vr, vzeta)
+    # define some local variables for convenience/tidiness
+    n_ion_species = composition.n_ion_species
+    n_neutral_species = composition.n_neutral_species
+    n_neutral_species_alloc = max(1,composition.n_neutral_species)
+    ##                              ##
+    # ion particle advection structs #
+    ##                              ##
+    # create structure z_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the part of the ion kinetic equation dealing
+    # with advection in z
+    begin_serial_region()
+    z_advect = setup_advection(n_ion_species, z, vpa, vperp, r)
+    # create structure r_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the split part of the ion kinetic equation dealing
+    # with advection in r
+    begin_serial_region()
+    r_advect = setup_advection(n_ion_species, r, vpa, vperp, z)
+    # create structure vpa_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the split part of the ion kinetic equation dealing
+    # with advection in vpa
+    begin_serial_region()
+    vpa_advect = setup_advection(n_ion_species, vpa, vperp, z, r)
+    # create structure vperp_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the split part of the ion kinetic equation dealing
+    # with advection in vperp
+    begin_serial_region()
+    vperp_advect = setup_advection(n_ion_species, vperp, vpa, z, r)
+    ##                                   ##
+    # electron particle advection structs #
+    ##                                   ##
+    # create structure electron_z_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the part of the electron kinetic equation dealing
+    # with advection in z
+    begin_serial_region()
+    electron_z_advect = setup_advection(1, z, vpa, vperp, r)
+    # create structure vpa_advect whose members are the arrays needed to compute
+    # the advection term(s) appearing in the part of the electron kinetic equation dealing
+    # with advection in vpa
+    begin_serial_region()
+    electron_vpa_advect = setup_advection(1, vpa, vperp, z, r)
+    ##                                  ##
+    # neutral particle advection structs #
+    ##                                  ##
+    # create structure neutral_z_advect for neutral particle advection
+    begin_serial_region()
+    neutral_z_advect = setup_advection(n_neutral_species_alloc, z, vz, vr, vzeta, r)
+    # create structure neutral_r_advect for neutral particle advection
+    begin_serial_region()
+    neutral_r_advect = setup_advection(n_neutral_species_alloc, r, vz, vr, vzeta, z)
+    # create structure neutral_vz_advect for neutral particle advection
+    begin_serial_region()
+    neutral_vz_advect = setup_advection(n_neutral_species_alloc, vz, vr, vzeta, z, r)
+    ##                                                                 ##
+    # construct named list of advection structs to compactify arguments #
+    ##                                                                 ##
+    advection_structs = advect_object_struct(vpa_advect, vperp_advect, z_advect, r_advect, 
+                                             electron_z_advect, electron_vpa_advect,
+                                             neutral_z_advect, neutral_r_advect, neutral_vz_advect)
+    return advection_structs
+end
+
 """
 create arrays and do other work needed to setup
 the main time advance loop.
@@ -163,11 +241,12 @@ this includes creating and populating structs
 for Chebyshev transforms, velocity space moments,
 EM fields, and advection terms
 """
-function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
+function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
                              vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
                              z_spectral, r_spectral, composition, moments,
                              fields, t_input, collisions, geometry,
-                             boundary_distributions, num_diss_params, restarting)
+                             boundary_distributions, num_diss_params, restarting,
+                             scratch_dummy)
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
@@ -182,13 +261,14 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
 
     begin_serial_region()
 
-    # create an array of structs containing scratch arrays for the pdf and low-order moments
-    # that may be evolved separately via fluid equations
-    scratch = setup_scratch_arrays(moments, pdf.ion.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    # # initialize the array of structs containing scratch arrays for the pdf and low-order moments
+    # # that may be evolved separately via fluid equations
+    # #scratch = setup_scratch_arrays(moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    # initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
-    scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
-                                   composition.n_ion_species,n_neutral_species_alloc)
+    # scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
+    #                               composition.n_ion_species,n_neutral_species_alloc)
    
     begin_serial_region()
     # update the derivatives of the electron moments as these may be needed when
@@ -217,6 +297,7 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     # with advection in z
     begin_serial_region()
     z_advect = setup_advection(n_ion_species, z, vpa, vperp, r)
+    electron_z_advect = setup_advection(1, z, vpa, vperp, r)
     # initialise the z advection speed
     begin_s_r_vperp_vpa_region()
     @loop_s is begin
@@ -294,7 +375,9 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
 
     #advect_objects = (vpa_advect = vpa_advect, vperp_advect = vperp_advect, z_advect = z_advect,
     # r_advect = r_advect, neutral_z_advect = neutral_z_advect, neutral_r_advect = neutral_r_advect)
-    advect_objects = advect_object_struct(vpa_advect, vperp_advect, z_advect, r_advect, neutral_z_advect, neutral_r_advect, neutral_vz_advect)
+    advect_objects = advect_object_struct(vpa_advect, vperp_advect, z_advect, r_advect, 
+                                          neutral_z_advect, neutral_r_advect, neutral_vz_advect,
+                                          electron_z_advect)
     #spectral_objects = (vz_spectral = vz_spectral, vr_spectral = vr_spectral, vzeta_spectral = vzeta_spectral,
     # vpa_spectral = vpa_spectral, vperp_spectral = vperp_spectral, z_spectral = z_spectral, r_spectral = r_spectral)
     spectral_objects = spectral_object_struct(vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral, r_spectral)
@@ -379,9 +462,9 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
             scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         # calculate the electron parallel heat flux
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, moments.electron.ppar,
-            moments.electron.upar, moments.electron.dT_dz, moments.ion.upar, 
-            collisions.nu_ei, composition.me_over_mi, composition.electron_physics)
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm,
+            moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
+            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
         # calculate the electron-ion parallel friction force
         calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
             moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
@@ -417,8 +500,7 @@ function setup_time_advance!(pdf, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
     # Ensure all processes are synchronized at the end of the setup
     _block_synchronize()
 
-    return moments, spectral_objects, advect_objects, scratch, advance, 
-        scratch_dummy, manufactured_source_list
+    return moments, spectral_objects, advect_objects, advance, manufactured_source_list
 end
 
 """
@@ -615,6 +697,20 @@ function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies
     buffer_vpavperpzrs_1 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
     buffer_vpavperpzrs_2 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
 
+    buffer_vpavperpzr_1 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    buffer_vpavperpzr_2 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    buffer_vpavperpzr_3 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    buffer_vpavperpzr_4 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    buffer_vpavperpzr_5 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    buffer_vpavperpzr_6 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    
+    buffer_vpavperpr_1 = allocate_shared_float(nvpa,nvperp,nr)
+    buffer_vpavperpr_2 = allocate_shared_float(nvpa,nvperp,nr)
+    buffer_vpavperpr_3 = allocate_shared_float(nvpa,nvperp,nr)
+    buffer_vpavperpr_4 = allocate_shared_float(nvpa,nvperp,nr)
+    buffer_vpavperpr_5 = allocate_shared_float(nvpa,nvperp,nr)
+    buffer_vpavperpr_6 = allocate_shared_float(nvpa,nvperp,nr)
+
     buffer_vzvrvzetazsn_1 = allocate_shared_float(nvz,nvr,nvzeta,nz,nspecies_neutral)
     buffer_vzvrvzetazsn_2 = allocate_shared_float(nvz,nvr,nvzeta,nz,nspecies_neutral)
     buffer_vzvrvzetazsn_3 = allocate_shared_float(nvz,nvr,nvzeta,nz,nspecies_neutral)
@@ -642,7 +738,9 @@ function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies
         buffer_vpavperpzrs_1,buffer_vpavperpzrs_2,
         buffer_vzvrvzetazsn_1,buffer_vzvrvzetazsn_2,buffer_vzvrvzetazsn_3,buffer_vzvrvzetazsn_4,buffer_vzvrvzetazsn_5,buffer_vzvrvzetazsn_6,
         buffer_vzvrvzetarsn_1,buffer_vzvrvzetarsn_2,buffer_vzvrvzetarsn_3,buffer_vzvrvzetarsn_4,buffer_vzvrvzetarsn_5,buffer_vzvrvzetarsn_6,
-        buffer_vzvrvzetazrsn_1, buffer_vzvrvzetazrsn_2)
+        buffer_vzvrvzetazrsn_1, buffer_vzvrvzetazrsn_2, 
+        buffer_vpavperpzr_1, buffer_vpavperpzr_2,buffer_vpavperpzr_3,buffer_vpavperpzr_4,buffer_vpavperpzr_5,buffer_vpavperpzr_6,
+        buffer_vpavperpr_1, buffer_vpavperpr_2, buffer_vpavperpr_3, buffer_vpavperpr_4, buffer_vpavperpr_5, buffer_vpavperpr_6)
 
 end
 
@@ -671,52 +769,20 @@ function normalize_pdf!(pdf, moments, scratch)
 end
 
 """
-create an array of structs containing scratch arrays for the normalised pdf and low-order moments
+initialize the array of structs containing scratch arrays for the normalised pdf and low-order moments
 that may be evolved separately via fluid equations
 """
-function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
-    # create n_rk_stages+1 structs, each of which will contain one pdf,
-    # one density, and one parallel flow array
-    scratch = Vector{scratch_pdf{5,3,2,6,3}}(undef, n_rk_stages+1)
-    pdf_dims = size(pdf_ion_in)
-    moment_ion_dims = size(moments.ion.dens)
-    moment_electron_dims = size(moments.electron.dens)
-    pdf_neutral_dims = size(pdf_neutral_in)
-    moment_neutral_dims = size(moments.neutral.dens)
+function initialize_scratch_arrays!(scratch, moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
     # populate each of the structs
     for istage ∈ 1:n_rk_stages+1
-        # Allocate arrays in temporary variables so that we can identify them
-        # by source line when using @debug_shared_array
-        pdf_array = allocate_shared_float(pdf_dims...)
-        density_array = allocate_shared_float(moment_ion_dims...)
-        upar_array = allocate_shared_float(moment_ion_dims...)
-        ppar_array = allocate_shared_float(moment_ion_dims...)
-        temp_z_s_array = allocate_shared_float(moment_ion_dims...)
-
-        electron_density_array = allocate_shared_float(moment_electron_dims...)
-        electron_upar_array = allocate_shared_float(moment_electron_dims...)
-        electron_ppar_array = allocate_shared_float(moment_electron_dims...)
-        electron_temp_array = allocate_shared_float(moment_electron_dims...)
-
-        pdf_neutral_array = allocate_shared_float(pdf_neutral_dims...)
-        density_neutral_array = allocate_shared_float(moment_neutral_dims...)
-        uz_neutral_array = allocate_shared_float(moment_neutral_dims...)
-        pz_neutral_array = allocate_shared_float(moment_neutral_dims...)
-
-
-        scratch[istage] = scratch_pdf(pdf_array, density_array, upar_array,
-                                      ppar_array, temp_z_s_array,
-                                      electron_density_array, electron_upar_array,
-                                      electron_ppar_array, electron_temp_array,
-                                      pdf_neutral_array, density_neutral_array,
-                                      uz_neutral_array, pz_neutral_array)
         @serial_region begin
             # initialise the scratch arrays for the ion pdf and velocity moments
             scratch[istage].pdf .= pdf_ion_in
             scratch[istage].density .= moments.ion.dens
             scratch[istage].upar .= moments.ion.upar
             scratch[istage].ppar .= moments.ion.ppar
-            # initialise the scratch arrays for the electron velocity moments
+            # initialise the scratch arrays for the electron pdf and velocity moments
+            scratch[istage].pdf_electron .= pdf_electron_in
             scratch[istage].electron_density .= moments.electron.dens
             scratch[istage].electron_upar .= moments.electron.upar
             scratch[istage].electron_ppar .= moments.electron.ppar
@@ -728,7 +794,7 @@ function setup_scratch_arrays(moments, pdf_ion_in, pdf_neutral_in, n_rk_stages)
             scratch[istage].pz_neutral .= moments.neutral.pz
         end
     end
-    return scratch
+    return nothing
 end
 
 """
@@ -1256,9 +1322,9 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                            num_diss_params, composition.electron_physics)
     # update the electron parallel heat flux
-    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.electron_ppar,
-        new_scratch.electron_upar, moments.electron.dT_dz, new_scratch.upar, collisions.nu_ei,
-        composition.me_over_mi, composition.electron_physics)
+    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.pdf_electron,
+        new_scratch.electron_ppar, new_scratch.electron_upar, moments.electron.vth, moments.electron.dT_dz, 
+        new_scratch.upar, collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
     # update the electron parallel friction force
     calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
         new_scratch.electron_upar, new_scratch.upar, moments.electron.dT_dz, composition.me_over_mi,
@@ -1747,7 +1813,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
     if advance.electron_energy
         electron_energy_equation!(fvec_out.electron_ppar, fvec_out.density, fvec_in, moments, collisions, dt,
-                                  composition, num_diss_params)
+                                  composition, num_diss_params, z.grid)
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index ee78dc6d3..5011a7d82 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -28,7 +28,7 @@ using ..calculus: integral
 using ..communication
 using ..derivatives: derivative_z!
 using ..looping
-using ..input_structs: braginskii_fluid
+using ..input_structs: braginskii_fluid, kinetic_electrons
 
 #global tmpsum1 = 0.0
 #global tmpsum2 = 0.0
@@ -146,6 +146,8 @@ mutable struct moments_electron_substruct
     dT_dz::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the upwinded z-derivative of the temperature Tpar = ppar/dens
     dT_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the electron thermal speed vth = sqrt(2*Tpar/m)
+    dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
 end
 
 """
@@ -344,7 +346,7 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     # need dupar/dz to obtain, e.g., the updated electron temperature
     dupar_dz = allocate_shared_float(nz, nr)
     dppar_dz = allocate_shared_float(nz, nr)
-    if electron_model == braginskii_fluid
+    if electron_model ∈ (braginskii_fluid, kinetic_electrons)
         dppar_dz_upwind = allocate_shared_float(nz, nr)
         dT_dz_upwind = allocate_shared_float(nz, nr)
     else
@@ -358,6 +360,7 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     end
     dqpar_dz = allocate_shared_float(nz, nr)
     dT_dz = allocate_shared_float(nz, nr)
+    dvth_dz = allocate_shared_float(nz, nr)
     
     # return struct containing arrays needed to update moments
     return moments_electron_substruct(density, density_updated, parallel_flow,
@@ -366,7 +369,7 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
         parallel_friction_force, heat_source, v_norm_fac,
         ddens_dz, dupar_dz, dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, 
-        dT_dz, dT_dz_upwind)
+        dT_dz, dT_dz_upwind, dvth_dz)
 end
 
 # neutral particles have natural mean velocities 
@@ -864,7 +867,7 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
     @loop_r_z ir iz begin
         dummy_zr[iz,ir] = -upar[iz,ir]
     end
-    if electron_model == braginskii_fluid
+    if electron_model ∈ (braginskii_fluid, kinetic_electrons)
         @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)

From 1bd3d21240eca2391ca1a548a527f7b4cd9dc961 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:31:40 +0000
Subject: [PATCH 040/428] Fix electron ppar evolution

1. Remove boundary condition on ppar.
2. When subcycling ppar, update qpar as vthe changes (i.e. calculate
   qpar as vthe^3*wpa3_moment).

Also increase timestep for ppar subcycling.
---
 src/electron_fluid_equations.jl  | 4 ++--
 src/electron_kinetic_equation.jl | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index 4fe56acb7..17d6fe8f0 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -172,7 +172,7 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     end
     # enforce the parallel boundary condtion on the electron parallel pressure
     #println("!!!NO PARALLEL BC IS BEING ENFORCED ON ELECTRON PRESSURE!!!")
-    enforce_parallel_BC_on_electron_pressure!(ppar, dens_i, composition.T_wall, fvec.ppar)
+    #enforce_parallel_BC_on_electron_pressure!(ppar, dens_i, composition.T_wall, fvec.ppar)
     return nothing
 end
 
@@ -330,4 +330,4 @@ function enforce_parallel_BC_on_electron_pressure!(ppar, dens, T_wall, ppar_i)
     end
 end
 
-end
\ No newline at end of file
+end
diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index bbba05a4d..531f4bd56 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -252,11 +252,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             println(io_vth,"")
         end
 
-        dt_energy = dt_electron
+        dt_energy = dt_electron * 100.0
 
         # get an updated iterate of the electron parallel pressure
         #ppar .= ppar_old
+        wpa3_moment = @. qpar / vthe^3
         for i in 1:1000
+            @. qpar = vthe^3 * wpa3_moment
             # Compute the upwinded z-derivative of the electron parallel pressure for the 
             # electron energy equation
             dummy_zr .= -moments.electron.upar

From 792624046dcf4a9abc783f18f029f3f051fd57bc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:34:49 +0000
Subject: [PATCH 041/428] Fix second derivative in electron pdf dissipation

Disable z-dissipation as we do not have a proper second derivative
function available.

Use second_derivative!() for vpa dissipation instead of calling
derivative!() twice, to avoid problems at element boundaries. Would
probably be even better to use the weak-form 2nd derivative, but need to
merge master-branch updates first.
---
 src/electron_kinetic_equation.jl | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 531f4bd56..24418c53e 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -6,7 +6,7 @@ export get_electron_critical_velocities
 
 using ..looping
 using ..derivatives: derivative_z!
-using ..calculus: derivative!, integral
+using ..calculus: derivative!, second_derivative!, integral
 using ..interpolation: interpolate_to_grid_1d!
 using ..type_definitions: mk_float
 using ..array_allocation: allocate_float
@@ -1054,17 +1054,19 @@ function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa,
     buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
     buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
     # add in numerical dissipation terms
-    @loop_vperp_vpa ivperp ivpa begin
-        @views derivative_z!(dummy_zr1, pdf[ivpa,ivperp,:,:], buffer_r_1, buffer_r_2, buffer_r_3,
-                             buffer_r_4, z_spectral, z)
-        @views derivative_z!(dummy_zr2, dummy_zr1, buffer_r_1, buffer_r_2, buffer_r_3,
-                             buffer_r_4, z_spectral, z)
-        @. residual[ivpa,ivperp,:,:] -= num_diss_params.z_dissipation_coefficient * dummy_zr2
-    end
+    #@loop_vperp_vpa ivperp ivpa begin
+    #    @views derivative_z!(dummy_zr1, pdf[ivpa,ivperp,:,:], buffer_r_1, buffer_r_2, buffer_r_3,
+    #                         buffer_r_4, z_spectral, z)
+    #    @views derivative_z!(dummy_zr2, dummy_zr1, buffer_r_1, buffer_r_2, buffer_r_3,
+    #                         buffer_r_4, z_spectral, z)
+    #    @. residual[ivpa,ivperp,:,:] -= num_diss_params.z_dissipation_coefficient * dummy_zr2
+    #end
     @loop_r_z_vperp ir iz ivperp begin
-        @views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
-        @views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
-        @. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
+        #@views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
+        #@views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
+        #@. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
+        @views second_derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
+        @. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch
     end
     #stop()
     return nothing

From 5f3417a9940ca4162bb9ddd23db7784e4e269f21 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:39:46 +0000
Subject: [PATCH 042/428] Enforce moment constraints on electron pdf away from
 wall boundaries

The boundary condition functions have a different algorithm for
enforcing constraints, so skip forcing them there, but enforce the
constraints everywhere else.
---
 src/electron_kinetic_equation.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 24418c53e..5c2e06bec 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -14,6 +14,7 @@ using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!
 using ..electron_vpa_advection: electron_vpa_advection!
+using ..moment_constraints: hard_force_moment_constraints!
 using ..velocity_moments: integrate_over_vspace
 
 """
@@ -199,6 +200,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         # enforce the boundary condition(s) on the electron pdf
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
+        for ir ∈ 1:size(ppar, 2), iz ∈ 2:size(ppar,1)-1
+            #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
+            @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
+        end
         
         if (mod(iteration,50)==1)
             @loop_vpa ivpa begin

From 7fe0670c00c47ea34342421bbe712a3713ce431d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:40:41 +0000
Subject: [PATCH 043/428] Don't set moments.electron.upar to 0.0

We need a non-zero upar to enforce the boundary condition.
---
 src/electron_kinetic_equation.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 5c2e06bec..4018edb98 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -103,8 +103,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
     num_diss_params, max_electron_pdf_iterations)
 
-    println("TMP FOR TESTING: SETTING UPAR_I = UPAR_E = 0!!!")
-    moments.electron.upar .= 0.0
+    #println("TMP FOR TESTING: SETTING UPAR_I = UPAR_E = 0!!!")
+    #moments.electron.upar .= 0.0
 
     # there will be a better way of doing this
     # store the incoming ppar

From 379a7a6af5ec6e13f72de20b30542a13a2c121e0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:43:13 +0000
Subject: [PATCH 044/428] Use finite-element versions of z-advection,
 vpa-advection

...also don't try to estimate timestep from CFL constraints. This is not
supported by the finite-element discretization, and also does not seem
to work well, as the timestep is constrained more by what happens at the
sheath edge boundaries than by the CFL conditions for advection in the
bulk of the domain.
---
 src/electron_kinetic_equation.jl | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 4018edb98..b32bfbf7c 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -143,7 +143,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                          buffer_r_4, z_spectral, z)
 
     #dt_electron = dt * sqrt(composition.me_over_mi)
-    dt_max = 1.0
+    dt_max = 1.0e-8 #1.0
     dt_energy = 0.001
     time = 0.0
 
@@ -922,16 +922,16 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     # initialise the residual to zero                                             
     residual .= 0.0
     # calculate the contribution to the residual from the z advection term
-    #electron_z_advection!(residual, pdf, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
-    dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
+    electron_z_advection!(residual, pdf, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
+    #dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
     single_term .= residual
     max_term .= abs.(residual)
     #println("z_advection: ", sum(residual), " dqpar_dz: ", sum(abs.(dqpar_dz)))
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
     # add in the contribution to the residual from the wpa advection term
-    #electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
-    #                        vpa_advect, vpa, vpa_spectral, scratch_dummy)
-    dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
+    electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
+                            vpa_advect, vpa, vpa_spectral, scratch_dummy)
+    #dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
     @. single_term = residual - single_term
     max_term .= max.(max_term, abs.(single_term))
     @. single_term = residual
@@ -965,7 +965,8 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #     println("")
     # end
     # stop()
-    dt_max = min(dt_max_zadv, dt_max_vadv)
+    #dt_max = min(dt_max_zadv, dt_max_vadv)
+    dt_max = dt_electron
     #println("dt_max: ", dt_max, " dt_max_zadv: ", dt_max_zadv, " dt_max_vadv: ", dt_max_vadv)
     return dt_max
 end

From 7876c1995937f0a8a760cbb7d77d2d604079e6f6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:48:20 +0000
Subject: [PATCH 045/428] Check convergence vs abs.(pdf) rather than max_term

The max_term is the maximum absolute value of various terms that
contribute to the residual, so is almost always smaller than the
residual itself. When max_term was used, the interation could never
converge.
---
 src/electron_kinetic_equation.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index b32bfbf7c..bbab7791a 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -191,7 +191,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     io_pdf_stages = open("pdf_zright.txt", "w")
 
     # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
-    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+    #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
     #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
     # evolve (artificially) in time until the residual is less than the tolerance
     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
@@ -301,7 +302,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                             num_diss_params, dt_max)
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
-        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+        #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
         if electron_pdf_converged
             break
         end

From 6a75234cd8ef76e912f89a8c502afefb5389ffcb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:51:35 +0000
Subject: [PATCH 046/428] Modify residual to relax CFL constraint at large wpa

Multiplying `residual` by `1 / sqrt(1 + vpa.grid^2)` relaxes the CFL
constraint from the wpa*dg/dwpa term at large wpa. It messes up time
evolution, but should not change the steady state.
---
 src/electron_kinetic_equation.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index bbab7791a..4813a8a08 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -161,6 +161,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                         ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
                                         z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                         num_diss_params, dt_max)
+    # Divide by wpa to relax CFL condition at large wpa - only looking for steady
+    # state here, so does not matter that this makes time evolution incorrect.
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        residual[ivpa,ivperp,iz,ir] /= sqrt(1.0 + vpa.grid[ivpa]^2)
+    end
     # open files to write the electron heat flux and pdf to file                                
     io_upar = open("upar.txt", "w")
     io_qpar = open("qpar.txt", "w")
@@ -301,6 +306,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                             dppar_dz, dqpar_dz, dvth_dz, 
                                             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                             num_diss_params, dt_max)
+        # Divide by wpa to relax CFL condition at large wpa - only looking for steady
+        # state here, so does not matter that this makes time evolution incorrect.
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            residual[ivpa,ivperp,iz,ir] /= sqrt(1.0 + vpa.grid[ivpa]^2)
+        end
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))

From 9392935cee5a1b0220970039d087cbd5a19d2b50 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:53:22 +0000
Subject: [PATCH 047/428] Print phi as a diagnostic from both ends of the
 domain

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

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 4813a8a08..14bbf546d 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -249,8 +249,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
 
-        if (mod(iteration,50) == 0)
-            println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,1], " average_residual: ", average_residual)
+        if (mod(iteration,output_interval) == 0)
+            println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,[1,end]], " average_residual: ", average_residual)
             @loop_z iz begin
                 println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
                 println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)

From 55c72427155a44d0e1af6e05b33ec37e7967a018 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:54:07 +0000
Subject: [PATCH 048/428] Zero out any negative values of the pdf

Helps increase stability when applying moment constraints. Useful when
'ringing' occurs near discontinuities.
---
 src/electron_kinetic_equation.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 14bbf546d..74509b379 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -206,6 +206,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         # enforce the boundary condition(s) on the electron pdf
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
+        pdf = max.(pdf, 0.0)
         for ir ∈ 1:size(ppar, 2), iz ∈ 2:size(ppar,1)-1
             #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
             @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)

From ccba9bc328c7a3e4fc9a2e1e423bfdf427412ba7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 19:58:03 +0000
Subject: [PATCH 049/428] Make calculation of cutoff and interpolation of
 electron pdf more robust

By calculating the cutoff velocity for the electron distribution
function just by integrating over the high-velocity incoming tail of the
electron distribution function, we remove sensitivity to any
shift/interpolation of the distribution function and to behaviour around
vpa=0.

At large vpa, as upar_e is negligible compared to vth_e, there is
essentially no difference between vth/n*f(vpa) and g(vth*wpa), so this
is equivalent to calculating the cutoff velocity by converting from g to
f (neglecting upar contributions as small in sqrt(mass ratio)) and
integrating over the tail of f.

Ignoring upar_e entirely (might??) make the distribution behave oddly
around vpa=0. However, interpolating from wpa to vpa, reflecting, then
interpolating back is likely to exacerbate interpolation artefacts as
there are two interpolations to do. Replace this by interpolating
directly onto a grid of wpa(-vpa)=-vpa/vth-upar/vth then replacing the
elements of the pdf array which have vpa>0 (or vpa<0 on the opposite
boundary) with the interpolated (and cut off beyond v_cutoff) values
from `reversed_pdf`.
---
 src/electron_kinetic_equation.jl | 118 +++++++++++++++++++++----------
 1 file changed, 80 insertions(+), 38 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 74509b379..4005b43a2 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -353,46 +353,68 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
     pdf_adjustment_option = "vpa4_gaussian"
 
     # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
-    wpa_values = vpa.scratch
+    #wpa_values = vpa.scratch
     # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
-    interpolated_pdf = vpa.scratch2
+    #interpolated_pdf = vpa.scratch2
+    reversed_pdf = vpa.scratch
 
     # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
     ivpa_zero = (vpa.n+1)÷2
 
     @loop_r ir begin
         # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        @. wpa_values = vpa.grid - upar[1,ir] / vthe[1,ir]
+        #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
+        #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[1,ir] / vthe[1,ir]
         # interpolate the pdf onto this grid
-        @views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+        #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+        @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+        reverse!(reversed_pdf)
         # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
-        @. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
+        #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
         # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        @. wpa_values = vpa.grid + upar[1,ir] / vthe[1,ir]
+        #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
         # interpolate back onto the original wpa grid
-        @views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+        #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
         # construct wpa * pdf
-        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
+        #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
         # calculate the first moment of the normalised pdf
         #first_vspace_moment = 0.0
-        first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
         # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
-        ivpa_max = vpa.n
+        ivpa = 0
+        ivpa_max = vpa.n + 1
         # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
         # if the first moment is positive, then the cutoff speed needs to be reduced
-        while first_vspace_moment > 0.0
-            # zero out the pdf at the current cutoff velocity
-            pdf[ivpa_max,1,1,ir] = 0.0
-            # update wpa * pdf
-            vpa.scratch3[ivpa_max] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # if first_vspace_moment > 0.0
+        upar_over_vth = upar[1,ir] / vthe[1,ir]
+        vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
+        upar_integral = 0.0
+        #while first_vspace_moment > upar_over_vth # > 0.0
+        #    # zero out the pdf at the current cutoff velocity
+        #    pdf[ivpa_max,1,1,ir] = 0.0
+        #    # update wpa * pdf
+        #    vpa.scratch3[ivpa_max] = 0.0
+        #    # calculate the updated first moment of the normalised pdf
+        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #    if first_vspace_moment > upar_over_vth #0.0
+        #        ivpa_max -= 1
+        #    end
+        #end
+        upar0 = upar[1,ir]
+        while upar_integral > upar0
+            ivpa += 1
             ivpa_max -= 1
-            # end
+            # zero out the reversed pdf at the current cutoff velocity
+            reversed_pdf[ivpa_max] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
         end
+        vmax = vpa_unnorm[ivpa_max]
         # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+        #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+        phi[1,ir] = me_over_mi * vmax^2
+        iv0 = findfirst(x -> x>0.0, vpa_unnorm)
+        pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments
         # of the modified pdf are 1, 0 and 1/2 respectively
         # will need vpa / vthe = wpa + upar/vthe
@@ -517,15 +539,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
 
     @loop_r ir begin
         # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        @. wpa_values = vpa.grid - upar[end,ir] / vthe[end,ir]
+        #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[end,ir] / vthe[end,ir]
         # interpolate the pdf onto this grid
-        @views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+        #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+        @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+        reverse!(reversed_pdf)
         # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
-        @. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
+        #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
         # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        @. wpa_values = vpa.grid + upar[end,ir] / vthe[end,ir]
+        #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
         # interpolate back onto the original wpa grid
-        @views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+        #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
 
         # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
         # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
@@ -539,24 +564,38 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         # println(io_pdf_stages,"")
 
         # construct wpa * pdf
-        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
+        #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
         # calculate the first moment of the normalised pdf
-        first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
         # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
-        ivpa_min = 1
+        ivpa = vpa.n+1
+        ivpa_min = 0
         # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
         # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
-        while first_vspace_moment < 0.0
-            # zero out the pdf at the current cutoff velocity
-            pdf[ivpa_min,1,end,ir] = 0.0
-            # update wpa * pdf
-            vpa.scratch3[ivpa_min] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # if first_vspace_moment < 0.0
+        upar_over_vth = upar[end,ir] / vthe[end,ir]
+        vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
+        upar_integral = 0.0
+        #while first_vspace_moment < upar_over_vth # < 0.0
+        #    # zero out the pdf at the current cutoff velocity
+        #    pdf[ivpa_min,1,end,ir] = 0.0
+        #    # update wpa * pdf
+        #    vpa.scratch3[ivpa_min] = 0.0
+        #    # calculate the updated first moment of the normalised pdf
+        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #    if first_vspace_moment < upar_over_vth
+        #        ivpa_min += 1
+        #    end
+        #end
+        upar_end = upar[end,ir]
+        while upar_integral < upar_end
+            ivpa -= 1
             ivpa_min += 1
-            # end
+            # zero out the reversed pdf at the current cutoff velocity
+            reversed_pdf[ivpa_min] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
         end
+        vmin = vpa_unnorm[ivpa_min]
 
         # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
         # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
@@ -570,7 +609,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         # println(io_pdf_stages,"")
 
         # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
+        #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
+        phi[end,ir] = me_over_mi * vmin^2
+        iv0 = findlast(x -> x<0.0, vpa_unnorm)
+        pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments
         # of the modified pdf are 1, 0 and 1/2 respectively
         # will need vpa / vthe = wpa + upar/vthe

From 9ee380d8dcca5f479c1c5fa18c6a0675d587323a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:05:45 +0000
Subject: [PATCH 050/428] Reorder include()'s - use moment_constraints in
 electron_kinetic_eqution

---
 src/moment_constraints.jl | 1 -
 src/moment_kinetics.jl    | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/moment_constraints.jl b/src/moment_constraints.jl
index 81dd05cdf..fbc6b29f2 100644
--- a/src/moment_constraints.jl
+++ b/src/moment_constraints.jl
@@ -6,7 +6,6 @@ function.
 module moment_constraints
 
 using ..communication: _block_synchronize
-using ..initial_conditions: enforce_zero_incoming_bc!
 using ..looping
 using ..velocity_moments: integrate_over_vspace, update_qpar!
 
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 387dc71b0..2e29fb333 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -43,9 +43,9 @@ include("electron_vpa_advection.jl")
 include("neutral_r_advection.jl")
 include("neutral_z_advection.jl")
 include("neutral_vz_advection.jl")
+include("moment_constraints.jl")
 include("electron_kinetic_equation.jl")
 include("initial_conditions.jl")
-include("moment_constraints.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("continuity.jl")

From e63b1c37d0fae101adfecde52600f4a8749af01a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:07:21 +0000
Subject: [PATCH 051/428] Skip enforcing first-moment constraint in electron
 wall bc

The constraint should be approximately enforced by the way the cutoff is
calculated, so the distribution function cannot drift far from the
constraint being satisfied. The error will be small in spatial
discretization (related to the spacing of the vpa grid), and should also
be constrained by the constraints being applied at neighbouring grid
points. Trying to apply the constraint might interfere with the symmetry
of g(wpa) around vpa=0 (between -v_cutoff and +v_cutoff).
---
 src/electron_kinetic_equation.jl | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 4005b43a2..2e1c7e17f 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -498,12 +498,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             @. vpa.scratch3 *= vpa.grid
             vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
             # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+            #     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            #normalisation_constant_A = (1 + normalisation_constant_B 
+            #    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            #normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
             normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                                       (vpa2_wpa2_moment
+                                        - wpa2_moment / zeroth_moment * vpa2_moment)
+            normalisation_constant_A = (1 - normalisation_constant_B 
+                                            * vpa2_moment) / zeroth_moment
+            normalisation_constant_C = 0.0
             @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         else
@@ -696,12 +702,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             @. vpa.scratch3 *= vpa.grid
             vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
             # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+            #     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            #normalisation_constant_A = (1 + normalisation_constant_B 
+            #    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            #normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
             normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                                       (vpa2_wpa2_moment
+                                        - wpa2_moment / zeroth_moment * vpa2_moment)
+            normalisation_constant_A = (1 - normalisation_constant_B 
+                                            * vpa2_moment) / zeroth_moment
+            normalisation_constant_C = 0.0
             @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         else

From 2aee70f31761fa3379d6071a96871fe5dcbed992 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:11:11 +0000
Subject: [PATCH 052/428] Reduce tolerance for electron iteration convergence

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

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 2e1c7e17f..20511c25f 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -1379,7 +1379,7 @@ function check_electron_pdf_convergence(residual, norm)
     #     average_residual /= length(residual)
     # end
     average_residual = sum(abs.(residual)) / sum(norm)
-    electron_pdf_converged = (average_residual < 1e-2)
+    electron_pdf_converged = (average_residual < 1e-3)
     return average_residual, electron_pdf_converged
 end
 
@@ -1389,4 +1389,4 @@ function check_electron_pdf_convergence(residual)
     return average_residual, electron_pdf_converged
 end
 
-end
\ No newline at end of file
+end

From b4b08d72edb3cd84804c296570eb78873bd4a84e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:12:29 +0000
Subject: [PATCH 053/428] DEBUGGING! return some arrays from electron solve for
 debugging

This entirely breaks the code after the first solve of the electron
steady state during initialisation, but useful for debugging!
---
 src/electron_kinetic_equation.jl | 17 +++++++++++++++--
 src/initial_conditions.jl        |  4 ++--
 src/moment_kinetics.jl           |  5 +++--
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 20511c25f..b279ff010 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -56,7 +56,7 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe,
     #solution_method = "picard_iteration"
     # solve the electron kinetic equation using the specified method
     if solution_method == "artificial_time_derivative"
-        update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
+        return update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
             moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
             num_diss_params, max_electron_pdf_iterations)
@@ -200,6 +200,15 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
     #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
     # evolve (artificially) in time until the residual is less than the tolerance
+    output_interval = 50
+    result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
+    result_pdf[:,:,1] .= pdf[:,1,:,1]
+    result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+    result_ppar[:,1] .= ppar[:,1]
+    result_qpar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+    result_qpar[:,1] .= qpar[:,1]
+    result_phi = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+    result_phi[:,1] .= phi[:,1]
     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
         @. pdf -= dt_electron * residual
@@ -262,6 +271,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             println(io_qpar,"")
             println(io_ppar,"")
             println(io_vth,"")
+            result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
+            result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
+            result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
+            result_phi[:,iteration÷output_interval+1] .= phi[:,1]
         end
 
         dt_energy = dt_electron * 100.0
@@ -336,7 +349,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     close(io_vth)
     close(io_pdf)
     close(io_pdf_stages)
-    return nothing    
+    return result_pdf, result_ppar, result_qpar, result_phi, z, vpa
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 75248305a..e01381d7a 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -340,7 +340,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, com
         # get the initial electrostatic potential and parallel electric field
         update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
         # initialize the electron pdf that satisfies the electron kinetic equation
-        initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, z, vpa, vperp, z_spectral, vpa_spectral, 
+        return initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, z, vpa, vperp, z_spectral, vpa_spectral, 
                                  advection_structs.electron_z_advect, advection_structs.electron_vpa_advect,
                                  scratch_dummy, collisions, composition, 
                                  num_diss_params, t_input.dt)
@@ -500,7 +500,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
             # solution for the electron pdf
             #max_electron_pdf_iterations = 500000
             max_electron_pdf_iterations = 10000
-            @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
+            return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
                                         moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
                                         phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
                                         moments.electron.dqpar_dz, moments.electron.dvth_dz, z, vpa, z_spectral, 
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 2e29fb333..ce30ed3ee 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -96,7 +96,8 @@ function run_moment_kinetics(to::TimerOutput, input_dict=Dict())
     mk_state = nothing
     try
         # set up all the structs, etc. needed for a run
-        mk_state = setup_moment_kinetics(input_dict)
+        #mk_state = setup_moment_kinetics(input_dict)
+        return setup_moment_kinetics(input_dict)
 
         # solve the 1+1D kinetic equation to advance f in time by nstep time steps
         if run_type == performance_test
@@ -382,7 +383,7 @@ function setup_moment_kinetics(input_dict::Dict; restart_prefix_iblock=nothing,
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
-        init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, composition, r, z,
+        return init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, composition, r, z,
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, vz_spectral,
                               species, collisions, t_input.use_manufactured_solns_for_init,
                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs)

From 0624f39feb0be6d0c4c6d969343a7e7e2203b291 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:13:46 +0000
Subject: [PATCH 054/428] DEBUGGING! Break the electron iteration if NaN is
 found

Means arrays get returned (from changes in previous debug commit) so the
cause of the NaNs can be investigated.
---
 src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index b279ff010..27db977b2 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -328,7 +328,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
-        if electron_pdf_converged
+        if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
             break
         end
         iteration += 1

From 356ef05bab055bffab683ce8304b0fece691a639 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 20:14:57 +0000
Subject: [PATCH 055/428] DEBUGGING! Commented out debugging print statements

---
 src/electron_kinetic_equation.jl | 23 +++++++++++++++++++++++
 src/initial_conditions.jl        |  1 +
 2 files changed, 24 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 27db977b2..f3657d77a 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -215,11 +215,16 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         # enforce the boundary condition(s) on the electron pdf
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
+        #println("A pdf 1 ", pdf[:,1,1,1])
+        #println("A pdf end ", pdf[:,1,end,1])
         pdf = max.(pdf, 0.0)
         for ir ∈ 1:size(ppar, 2), iz ∈ 2:size(ppar,1)-1
             #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
             @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
         end
+        #println("B pdf 1 ", pdf[:,1,1,1])
+        #println("B pdf end ", pdf[:,1,end,1])
+        #error("foo")
         
         if (mod(iteration,50)==1)
             @loop_vpa ivpa begin
@@ -400,6 +405,8 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
         # if the first moment is positive, then the cutoff speed needs to be reduced
         upar_over_vth = upar[1,ir] / vthe[1,ir]
+        #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
+        #println("$first_vspace_moment, u/vth=$upar_over_vth")
         vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
         upar_integral = 0.0
         #while first_vspace_moment > upar_over_vth # > 0.0
@@ -409,11 +416,13 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         #    vpa.scratch3[ivpa_max] = 0.0
         #    # calculate the updated first moment of the normalised pdf
         #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
         #    if first_vspace_moment > upar_over_vth #0.0
         #        ivpa_max -= 1
         #    end
         #end
         upar0 = upar[1,ir]
+        #println("before pdf left ", pdf[:,1,1,ir])
         while upar_integral > upar0
             ivpa += 1
             ivpa_max -= 1
@@ -421,13 +430,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             reversed_pdf[ivpa_max] = 0.0
             # calculate the updated first moment of the normalised pdf
             upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
+            #println("left ", ivpa, " ", upar_integral, " ", upar0)
         end
         vmax = vpa_unnorm[ivpa_max]
+        #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
+        #println("done first cutoff loop")
         # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
         #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
         phi[1,ir] = me_over_mi * vmax^2
         iv0 = findfirst(x -> x>0.0, vpa_unnorm)
         pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
+        #println("reversed_pdf ", reversed_pdf)
+        #println("after pdf left ", pdf[:,1,1,ir])
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments
         # of the modified pdf are 1, 0 and 1/2 respectively
         # will need vpa / vthe = wpa + upar/vthe
@@ -592,6 +606,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
         # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
         upar_over_vth = upar[end,ir] / vthe[end,ir]
+        #println("$first_vspace_moment, u/vth=$upar_over_vth")
         vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
         upar_integral = 0.0
         #while first_vspace_moment < upar_over_vth # < 0.0
@@ -606,6 +621,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         #    end
         #end
         upar_end = upar[end,ir]
+        #println("before pdf ", pdf[:,1,end,ir])
         while upar_integral < upar_end
             ivpa -= 1
             ivpa_min += 1
@@ -613,8 +629,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             reversed_pdf[ivpa_min] = 0.0
             # calculate the updated first moment of the normalised pdf
             upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
+            #println("right ", ivpa, " ", upar_integral, " ", upar_end)
         end
         vmin = vpa_unnorm[ivpa_min]
+        #println("done second cutoff loop")
 
         # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
         # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
@@ -632,6 +650,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         phi[end,ir] = me_over_mi * vmin^2
         iv0 = findlast(x -> x<0.0, vpa_unnorm)
         pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
+        #println("after pdf ", pdf[:,1,end,ir])
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments
         # of the modified pdf are 1, 0 and 1/2 respectively
         # will need vpa / vthe = wpa + upar/vthe
@@ -1006,6 +1025,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
     single_term .= residual
     max_term .= abs.(residual)
+    #println("z_adv residual = ", maximum(abs.(single_term)))
     #println("z_advection: ", sum(residual), " dqpar_dz: ", sum(abs.(dqpar_dz)))
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
     # add in the contribution to the residual from the wpa advection term
@@ -1015,12 +1035,14 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     @. single_term = residual - single_term
     max_term .= max.(max_term, abs.(single_term))
     @. single_term = residual
+    #println("v_adv residual = ", maximum(abs.(single_term)))
     #add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
     # add in the contribution to the residual from the term proportional to the pdf
     add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
     @. single_term = residual - single_term
     max_term .= max.(max_term, abs.(single_term))
     @. single_term = residual
+    #println("pdf_term residual = ", maximum(abs.(single_term)))
     # @loop_vpa ivpa begin
     #     @loop_z iz begin
     #         println("LHS: ", residual[ivpa,1,iz,1], " vpa: ", vpa.grid[ivpa], " z: ", z.grid[iz], " dvth_dz: ", dvth_dz[iz,1], " type: ", 1) 
@@ -1031,6 +1053,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     # add in numerical dissipation terms
     add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
     @. single_term = residual - single_term
+    #println("dissipation residual = ", maximum(abs.(single_term)))
     max_term .= max.(max_term, abs.(single_term))
     # add in particle and heat source term(s)
     #@. single_term = residual
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index e01381d7a..2e9dc133e 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -1060,6 +1060,7 @@ function init_electron_pdf_over_density!(pdf, density, upar, vth, phi, z, vpa, v
         @loop_z_vperp iz ivperp begin
             #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
             #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
+            #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
             @. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2) * (
                                   (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
                                   tanh(sharp_fac*(vpa.grid[:]-vpa_crit_zmin))) *

From 78abb26156bf3aa49319ac4e637b57b9f36ef462 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Feb 2024 23:40:45 +0000
Subject: [PATCH 056/428] Change 'charged' -> 'ion' in all remaining places

Some of the missed places are due to the merge, some may have just been
overlooked.
---
 debug_test/README.md                          |   2 +-
 src/analysis.jl                               |   6 +-
 src/continuity.jl                             |   8 +-
 src/external_sources.jl                       |  40 ++---
 src/file_io.jl                                |  34 ++--
 src/initial_conditions.jl                     |  10 +-
 src/krook_collisions.jl                       |  10 +-
 src/load_data.jl                              | 116 +++++++-------
 src/makie_post_processing.jl                  |  34 ++--
 src/manufactured_solns.jl                     |  28 ++--
 src/post_processing.jl                        |  42 ++---
 src/source_terms.jl                           |   4 +-
 src/time_advance.jl                           |  44 +++---
 src/velocity_moments.jl                       |  32 ++--
 src/vpa_advection.jl                          |  24 +--
 test/Krook_collisions_tests.jl                | 146 +++++++++---------
 test/fokker_planck_time_evolution_tests.jl    |  96 ++++++------
 ...ear_sound_wave_inputs_and_expected_data.jl |  56 +++----
 test/nonlinear_sound_wave_tests.jl            |  30 ++--
 test/restart_interpolation_tests.jl           |  74 ++++-----
 20 files changed, 417 insertions(+), 419 deletions(-)

diff --git a/debug_test/README.md b/debug_test/README.md
index fc2c93626..5f05055e7 100644
--- a/debug_test/README.md
+++ b/debug_test/README.md
@@ -81,7 +81,7 @@ Suggested debugging strategy for race conditions is:
   failure. Usually a failure should indicate where there is a missing
   `begin_*_region()` call. There may be places though where synchronization is
   required even though the type of loop macros used does not change (for
-  example when `phi` is calculated contributions from all charged species need
+  example when `phi` is calculated contributions from all ion species need
   to be summed, resulting in an unusual pattern of array accesses); in this
   case `_block_synchronize()` can be called directly.
     * The function `debug_check_shared_memory()` can be inserted between
diff --git a/src/analysis.jl b/src/analysis.jl
index bfc21d5a6..6dded4529 100644
--- a/src/analysis.jl
+++ b/src/analysis.jl
@@ -13,7 +13,7 @@ using ..coordinates: coordinate
 using ..initial_conditions: vpagrid_to_dzdt
 using ..interpolation: interpolate_to_grid_1d
 using ..load_data: open_readonly_output_file, get_nranks, load_pdf_data, load_rank_data
-using ..load_data: load_distributed_charged_pdf_slice
+using ..load_data: load_distributed_ion_pdf_slice
 using ..looping
 using ..type_definitions: mk_int
 using ..velocity_moments: integrate_over_vspace
@@ -139,12 +139,12 @@ function check_Chodura_condition(r, z, vperp, vpa, dens, upar, vth, composition,
         end
     end
     if f_lower === nothing
-        f_lower = load_distributed_charged_pdf_slice(run_name, nblocks, t_range,
+        f_lower = load_distributed_ion_pdf_slice(run_name, nblocks, t_range,
                                                      composition.n_ion_species, r, z,
                                                      vperp, vpa; iz=1, ir=ir0)
     end
     if f_upper === nothing
-        f_upper = load_distributed_charged_pdf_slice(run_name, nblocks, t_range,
+        f_upper = load_distributed_ion_pdf_slice(run_name, nblocks, t_range,
                                                      composition.n_ion_species, r, z,
                                                      vperp, vpa; iz=z.n_global, ir=ir0)
     end
diff --git a/src/continuity.jl b/src/continuity.jl
index 31302a2dd..1c3600f07 100644
--- a/src/continuity.jl
+++ b/src/continuity.jl
@@ -18,8 +18,8 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     @loop_s_r_z is ir iz begin
         # Use ddens_dz is upwinded using upar
         dens_out[iz,ir,is] -=
-            dt*(fvec_in.upar[iz,ir,is]*moments.charged.ddens_dz_upwind[iz,ir,is] +
-                fvec_in.density[iz,ir,is]*moments.charged.dupar_dz[iz,ir,is])
+            dt*(fvec_in.upar[iz,ir,is]*moments.ion.ddens_dz_upwind[iz,ir,is] +
+                fvec_in.density[iz,ir,is]*moments.ion.dupar_dz[iz,ir,is])
     end
 
     # update the density to account for ionization collisions;
@@ -31,7 +31,7 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     end
 
     if ion_source_settings.active
-        source_amplitude = moments.charged.external_source_density_amplitude
+        source_amplitude = moments.ion.external_source_density_amplitude
         @loop_s_r_z is ir iz begin
             dens_out[iz,ir,is] +=
                 dt * source_amplitude[iz,ir]
@@ -42,7 +42,7 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
-            dens_out[iz,ir,is] += dt*diffusion_coefficient*moments.charged.d2dens_dz2[iz,ir,is]
+            dens_out[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2dens_dz2[iz,ir,is]
         end
     end
 end
diff --git a/src/external_sources.jl b/src/external_sources.jl
index 641f33a1d..8ee1ae980 100644
--- a/src/external_sources.jl
+++ b/src/external_sources.jl
@@ -230,10 +230,10 @@ end
     initialize_external_source_amplitude!(moments, external_source_settings, vperp,
                                           vzeta, vr, n_neutral_species)
 
-Initialize the arrays `moments.charged.external_source_amplitude`,
-`moments.charged.external_source_density_amplitude`,
-`moments.charged.external_source_momentum_amplitude`,
-`moments.charged.external_source_pressure_amplitude`,
+Initialize the arrays `moments.ion.external_source_amplitude`,
+`moments.ion.external_source_density_amplitude`,
+`moments.ion.external_source_momentum_amplitude`,
+`moments.ion.external_source_pressure_amplitude`,
 `moments.neutral.external_source_amplitude`,
 `moments.neutral.external_source_density_amplitude`,
 `moments.neutral.external_source_momentum_amplitude`, and
@@ -246,20 +246,20 @@ function initialize_external_source_amplitude!(moments, external_source_settings
     if ion_source_settings.active
         if ion_source_settings.source_type == "energy"
             @loop_r_z ir iz begin
-                moments.charged.external_source_amplitude[iz,ir] =
+                moments.ion.external_source_amplitude[iz,ir] =
                     ion_source_settings.source_strength *
                     ion_source_settings.r_amplitude[ir] *
                     ion_source_settings.z_amplitude[iz]
             end
             if moments.evolve_density
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_density_amplitude[iz,ir] = 0.0
+                    moments.ion.external_source_density_amplitude[iz,ir] = 0.0
                 end
             end
             if moments.evolve_upar
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_momentum_amplitude[iz,ir] =
-                        - moments.charged.dens[iz,ir] * moments.charged.upar[iz,ir] *
+                    moments.ion.external_source_momentum_amplitude[iz,ir] =
+                        - moments.ion.dens[iz,ir] * moments.ion.upar[iz,ir] *
                           ion_source_settings.source_strength *
                           ion_source_settings.r_amplitude[ir] *
                           ion_source_settings.z_amplitude[iz]
@@ -267,9 +267,9 @@ function initialize_external_source_amplitude!(moments, external_source_settings
             end
             if moments.evolve_ppar
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_pressure_amplitude[iz,ir] =
+                    moments.ion.external_source_pressure_amplitude[iz,ir] =
                         (0.5 * ion_source_settings.source_T +
-                         moments.charged.upar[iz,ir]^2 - moments.charged.ppar[iz,ir]) *
+                         moments.ion.upar[iz,ir]^2 - moments.ion.ppar[iz,ir]) *
                         ion_source_settings.source_strength *
                         ion_source_settings.r_amplitude[ir] *
                         ion_source_settings.z_amplitude[iz]
@@ -277,14 +277,14 @@ function initialize_external_source_amplitude!(moments, external_source_settings
             end
         else
             @loop_r_z ir iz begin
-                moments.charged.external_source_amplitude[iz,ir] =
+                moments.ion.external_source_amplitude[iz,ir] =
                     ion_source_settings.source_strength *
                     ion_source_settings.r_amplitude[ir] *
                     ion_source_settings.z_amplitude[iz]
             end
             if moments.evolve_density
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_density_amplitude[iz,ir] =
+                    moments.ion.external_source_density_amplitude[iz,ir] =
                         ion_source_settings.source_strength *
                         ion_source_settings.r_amplitude[ir] *
                         ion_source_settings.z_amplitude[iz]
@@ -292,14 +292,14 @@ function initialize_external_source_amplitude!(moments, external_source_settings
             end
             if moments.evolve_upar
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_momentum_amplitude[iz,ir] = 0.0
+                    moments.ion.external_source_momentum_amplitude[iz,ir] = 0.0
                 end
             end
             if moments.evolve_ppar
                 @loop_r_z ir iz begin
-                    moments.charged.external_source_pressure_amplitude[iz,ir] =
+                    moments.ion.external_source_pressure_amplitude[iz,ir] =
                         (0.5 * ion_source_settings.source_T +
-                         moments.charged.upar[iz,ir]^2) *
+                         moments.ion.upar[iz,ir]^2) *
                         ion_source_settings.source_strength *
                         ion_source_settings.r_amplitude[ir] *
                         ion_source_settings.z_amplitude[iz]
@@ -384,7 +384,7 @@ end
 function initialize_external_source_controller_integral!(
              moments, external_source_settings, n_neutral_species)
 
-Initialize the arrays `moments.charged.external_source_controller_integral` and
+Initialize the arrays `moments.ion.external_source_controller_integral` and
 `moments.neutral.external_source_controller_integral`, using the settings in
 `external_source_settings`
 """
@@ -395,7 +395,7 @@ function initialize_external_source_controller_integral!(
         if ion_source_settings.PI_density_controller_I != 0.0 &&
             ion_source_settings.source_type ∈ ("density_profile_control",
                                                "density_midpoint_control")
-            moments.charged.external_source_controller_integral .= 0.0
+            moments.ion.external_source_controller_integral .= 0.0
         end
     end
 
@@ -421,7 +421,7 @@ Add external source term to the ion kinetic equation.
 function external_ion_source!(pdf, fvec, moments, ion_source_settings, vperp, vpa, dt)
     begin_s_r_z_vperp_region()
 
-    source_amplitude = moments.charged.external_source_amplitude
+    source_amplitude = moments.ion.external_source_amplitude
     source_T = ion_source_settings.source_T
     if vperp.n == 1
         vth_factor = 1.0 / sqrt(source_T)
@@ -432,7 +432,7 @@ function external_ion_source!(pdf, fvec, moments, ion_source_settings, vperp, vp
     vperp_grid = vperp.grid
 
     if moments.evolve_ppar && moments.evolve_upar && moments.evolve_density
-        vth = moments.charged.vth
+        vth = moments.ion.vth
         density = fvec.density
         upar = fvec.upar
         @loop_s_r_z is ir iz begin
@@ -609,7 +609,7 @@ source amplitude.
 function external_ion_source_controller!(fvec_in, moments, ion_source_settings, dt)
 
     is = 1
-    ion_moments = moments.charged
+    ion_moments = moments.ion
 
     if ion_source_settings.source_type == "Maxwellian"
         if moments.evolve_ppar
diff --git a/src/file_io.jl b/src/file_io.jl
index 06e767ea0..85de97a15 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -387,14 +387,14 @@ function write_boundary_distributions!(fid, boundary_distributions, parallel_io,
     @serial_region begin
         boundary_distributions_io = create_io_group(fid, "boundary_distributions")
 
-        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_left",
-            boundary_distributions.pdf_rboundary_charged[:,:,:,1,:], vpa, vperp, z,
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_ion_left",
+            boundary_distributions.pdf_rboundary_ion[:,:,:,1,:], vpa, vperp, z,
             parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
-            description="Initial charged-particle pdf at left radial boundary")
-        write_single_value!(boundary_distributions_io, "pdf_rboundary_charged_right",
-            boundary_distributions.pdf_rboundary_charged[:,:,:,2,:], vpa, vperp, z,
+            description="Initial ion-particle pdf at left radial boundary")
+        write_single_value!(boundary_distributions_io, "pdf_rboundary_ion_right",
+            boundary_distributions.pdf_rboundary_ion[:,:,:,2,:], vpa, vperp, z,
             parallel_io=parallel_io, n_ion_species=composition.n_ion_species,
-            description="Initial charged-particle pdf at right radial boundary")
+            description="Initial ion-particle pdf at right radial boundary")
         write_single_value!(boundary_distributions_io, "pdf_rboundary_neutral_left",
             boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:], vz, vr, vzeta, z,
             parallel_io=parallel_io, n_neutral_species=composition.n_neutral_species,
@@ -644,7 +644,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_pperp = create_dynamic_variable!(dynamic, "perpendicular_pressure", mk_float, z, r;
                                            n_ion_species=n_ion_species,
                                            parallel_io=parallel_io,
-                                           description="charged species perpendicular pressure",
+                                           description="ion species perpendicular pressure",
                                            units="n_ref*T_ref")
 
         # io_qpar is the handle for the ion parallel heat flux
@@ -695,7 +695,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_dSdt = create_dynamic_variable!(dynamic, "entropy_production", mk_float, z, r;
                                           n_ion_species=n_ion_species,
                                           parallel_io=parallel_io,
-                                          description="charged species entropy production",
+                                          description="ion species entropy production",
                                           units="")
 
         if parallel_io || z.irank == 0
@@ -1222,41 +1222,41 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
         end
         if z.irank == z.nrank - 1 # upper wall
             append_to_dynamic_var(io_moments.chodura_integral_upper,
-                                  moments.charged.chodura_integral_upper, t_idx,
+                                  moments.ion.chodura_integral_upper, t_idx,
                                   parallel_io, r, n_ion_species)
         elseif io_moments.chodura_integral_upper !== nothing
             append_to_dynamic_var(io_moments.chodura_integral_upper,
-                                  moments.charged.chodura_integral_upper, t_idx,
+                                  moments.ion.chodura_integral_upper, t_idx,
                                   parallel_io, 0, n_ion_species)
         end
         if io_moments.external_source_amplitude !== nothing
             append_to_dynamic_var(io_moments.external_source_amplitude,
-                                  moments.charged.external_source_amplitude, t_idx,
+                                  moments.ion.external_source_amplitude, t_idx,
                                   parallel_io, z, r)
             if moments.evolve_density
                 append_to_dynamic_var(io_moments.external_source_density_amplitude,
-                                      moments.charged.external_source_density_amplitude,
+                                      moments.ion.external_source_density_amplitude,
                                       t_idx, parallel_io, z, r)
             end
             if moments.evolve_upar
                 append_to_dynamic_var(io_moments.external_source_momentum_amplitude,
-                                      moments.charged.external_source_momentum_amplitude,
+                                      moments.ion.external_source_momentum_amplitude,
                                       t_idx, parallel_io, z, r)
             end
             if moments.evolve_ppar
                 append_to_dynamic_var(io_moments.external_source_pressure_amplitude,
-                                      moments.charged.external_source_pressure_amplitude,
+                                      moments.ion.external_source_pressure_amplitude,
                                       t_idx, parallel_io, z, r)
             end
         end
         if io_moments.external_source_controller_integral !== nothing
-            if size(moments.charged.external_source_controller_integral) == (1,1)
+            if size(moments.ion.external_source_controller_integral) == (1,1)
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
-                                      moments.charged.external_source_controller_integral[1,1],
+                                      moments.ion.external_source_controller_integral[1,1],
                                       t_idx, parallel_io)
             else
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
-                                      moments.charged.external_source_controller_integral,
+                                      moments.ion.external_source_controller_integral,
                                       t_idx, parallel_io, z, r)
             end
         end
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index fe1c794ee..cbc77c592 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -231,7 +231,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
             init_vth!(moments.ion.vth, z, r, species.ion, n_ion_species)
             @. moments.ion.ppar = 0.5 * moments.ion.dens * moments.ion.vth^2
             # initialise pressures assuming isotropic distribution
-            @. moments.ion.ppar = 0.5 * moments.charged.dens * moments.charged.vth^2
+            @. moments.ion.ppar = 0.5 * moments.ion.dens * moments.ion.vth^2
             @. moments.ion.pperp = moments.ion.ppar
             if n_neutral_species > 0
                 # initialise the neutral density profile
@@ -294,7 +294,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     # collision operator will not be calculated before the initial values are written to
     # file.
     @serial_region begin
-        moments.charged.dSdt .= 0.0
+        moments.ion.dSdt .= 0.0
     end
 
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
@@ -407,8 +407,8 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
         for is ∈ 1:composition.n_ion_species, ir ∈ 1:r.n
             # Add ion contributions to wall flux here. Neutral contributions will be
             # added in init_neutral_pdf_over_density!()
-            if species.charged[is].z_IC.initialization_option == "bgk" || species.charged[is].vpa_IC.initialization_option == "bgk"
-                @views init_bgk_pdf!(pdf.charged.norm[:,1,:,ir,is], 0.0, species.charged[is].initial_temperature, z.grid, z.L, vpa.grid)
+            if species.ion[is].z_IC.initialization_option == "bgk" || species.ion[is].vpa_IC.initialization_option == "bgk"
+                @views init_bgk_pdf!(pdf.ion.norm[:,1,:,ir,is], 0.0, species.ion[is].initial_temperature, z.grid, z.L, vpa.grid)
             else
                 # updates pdf_norm to contain pdf / density, so that ∫dvpa pdf.norm = 1,
                 # ∫dwpa wpa * pdf.norm = 0, and ∫dwpa m_s (wpa/vths)^2 pdf.norm = 1/2
@@ -1561,7 +1561,7 @@ end
 """
 enforce boundary conditions on neutral particle distribution function
 """
-function enforce_neutral_boundary_conditions!(f_neutral, f_charged,
+function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
         boundary_distributions, density_neutral, uz_neutral, pz_neutral, moments,
         density_ion, upar_ion, Er, vzeta_spectral, vr_spectral, vz_spectral, r_adv, z_adv,
         vzeta_adv, vr_adv, vz_adv, r, z, vzeta, vr, vz, composition, geometry,
diff --git a/src/krook_collisions.jl b/src/krook_collisions.jl
index fe6477098..a9223933e 100644
--- a/src/krook_collisions.jl
+++ b/src/krook_collisions.jl
@@ -85,7 +85,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         # through by vth, remembering pdf is already multiplied by vth
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
-            vth = moments.charged.vth[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
             nu_ii = get_collision_frequency(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
@@ -98,7 +98,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         # is already multiplied by vth, and grid is already normalized by vth
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
-            vth = moments.charged.vth[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
             nu_ii = get_collision_frequency(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
@@ -111,7 +111,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         # Compared to evolve_density version, grid is already shifted by upar
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
-            vth = moments.charged.vth[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
             nu_ii = get_collision_frequency(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
@@ -125,7 +125,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         # that pdf is already normalized by density
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
-            vth = moments.charged.vth[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
             nu_ii = get_collision_frequency(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
@@ -138,7 +138,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
     else
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
-            vth = moments.charged.vth[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
             if vperp.n == 1
                 vth_prefactor = 1.0 / vth
             else
diff --git a/src/load_data.jl b/src/load_data.jl
index 5f2585e05..b73ca82ac 100644
--- a/src/load_data.jl
+++ b/src/load_data.jl
@@ -442,10 +442,10 @@ function load_ion_moments_data(fid; printout=false, extended_moments = false)
     thermal_speed = load_variable(group, "thermal_speed")
 
     if extended_moments
-        # Read charged species perpendicular pressure
+        # Read ion species perpendicular pressure
         perpendicular_pressure = load_variable(group, "perpendicular_pressure")
 
-        # Read charged species entropy_production
+        # Read ion species entropy_production
         entropy_production = load_variable(group, "entropy_production")
     end
 
@@ -714,7 +714,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     load_moment("external_source_controller_integral")
             end
 
-            function load_charged_pdf()
+            function load_ion_pdf()
                 this_pdf = load_slice(dynamic, "f", vpa_range, vperp_range, z_range,
                                       r_range, :, time_index)
                 orig_nvpa, orig_nvperp, orig_nz, orig_nr, nspecies = size(this_pdf)
@@ -780,7 +780,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa - upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .- moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -793,7 +793,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -807,8 +807,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
-                            @. (vpa.grid - moments.charged.upar[iz,ir,is]) /
-                               moments.charged.vth[iz,ir,is]
+                            @. (vpa.grid - moments.ion.upar[iz,ir,is]) /
+                               moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -821,7 +821,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa + upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .+ moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -835,8 +835,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
-                            @. (vpa.grid + moments.charged.upar[iz,ir,is]) /
-                               moments.charged.vth
+                            @. (vpa.grid + moments.ion.upar[iz,ir,is]) /
+                               moments.ion.vth
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -849,7 +849,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -862,7 +862,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -875,8 +875,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth - upar/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.charged.vth[iz,ir,is] -
-                                              moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] -
+                                              moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -891,7 +891,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
                             @. vpa.grid -
-                               moments.charged.upar[iz,ir,is]/moments.charged.vth[iz,ir,is]
+                               moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -904,8 +904,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth + upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.charged.vth[iz,ir,is] +
-                                              moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] +
+                                              moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -918,7 +918,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -933,7 +933,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
                             @. vpa.grid +
-                               moments.charged.upar[iz,ir,is] / moments.charged.vth[iz,ir,is]
+                               moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,ir,is], restart_vpa,
@@ -954,45 +954,45 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                 if moments.evolve_density && !restart_evolve_density
                     # Need to normalise by density
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] ./= moments.charged.dens[iz,ir,is]
+                        this_pdf[:,:,iz,ir,is] ./= moments.ion.dens[iz,ir,is]
                     end
                 elseif !moments.evolve_density && restart_evolve_density
                     # Need to unnormalise by density
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] .*= moments.charged.dens[iz,ir,is]
+                        this_pdf[:,:,iz,ir,is] .*= moments.ion.dens[iz,ir,is]
                     end
                 end
                 if moments.evolve_ppar && !restart_evolve_ppar
                     # Need to normalise by vth
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] .*= moments.charged.vth[iz,ir,is]
+                        this_pdf[:,:,iz,ir,is] .*= moments.ion.vth[iz,ir,is]
                     end
                 elseif !moments.evolve_ppar && restart_evolve_ppar
                     # Need to unnormalise by vth
                     for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] ./= moments.charged.vth[iz,ir,is]
+                        this_pdf[:,:,iz,ir,is] ./= moments.ion.vth[iz,ir,is]
                     end
                 end
 
                 return this_pdf
             end
 
-            pdf.charged.norm .= load_charged_pdf()
+            pdf.ion.norm .= load_ion_pdf()
                 if z.irank == 0
-                    moments.charged.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower", :, :,
+                    moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower", :, :,
                                                   time_index)
                 else
-                    moments.charged.chodura_integral_lower .= 0.0
+                    moments.ion.chodura_integral_lower .= 0.0
                 end
                 if z.irank == z.nrank - 1
-                    moments.charged.chodura_integral_upper .= load_slice(dynamic, "chodura_integral_upper", :, :,
+                    moments.ion.chodura_integral_upper .= load_slice(dynamic, "chodura_integral_upper", :, :,
                                                   time_index)
                 else
-                    moments.charged.chodura_integral_upper .= 0.0
+                    moments.ion.chodura_integral_upper .= 0.0
                 end
             boundary_distributions_io = get_group(fid, "boundary_distributions")
 
-            function load_charged_boundary_pdf(var_name, ir)
+            function load_ion_boundary_pdf(var_name, ir)
                 this_pdf = load_slice(boundary_distributions_io, var_name, vpa_range,
                                       vperp_range, z_range, :)
                 orig_nvpa, orig_nvperp, orig_nz, nspecies = size(this_pdf)
@@ -1047,7 +1047,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa - upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .- moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1059,7 +1059,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1072,8 +1072,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
-                            @. (vpa.grid - moments.charged.upar[iz,ir,is]) /
-                               moments.charged.vth[iz,ir,is]
+                            @. (vpa.grid - moments.ion.upar[iz,ir,is]) /
+                               moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1085,7 +1085,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa + upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .+ moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1098,8 +1098,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
-                            @. (vpa.grid + moments.charged.upar[iz,ir,is]) /
-                               moments.charged.vth
+                            @. (vpa.grid + moments.ion.upar[iz,ir,is]) /
+                               moments.ion.vth
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1111,7 +1111,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1123,7 +1123,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1135,8 +1135,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth - upar/vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.charged.vth[iz,ir,is] -
-                                              moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] -
+                                              moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1150,7 +1150,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
                             @. vpa.grid -
-                               moments.charged.upar[iz,ir,is]/moments.charged.vth[iz,ir,is]
+                               moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1162,8 +1162,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth + upar
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.charged.vth[iz,ir,is] +
-                                              moments.charged.upar[iz,ir,is]
+                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] +
+                                              moments.ion.upar[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1175,7 +1175,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     # => old_wpa = new_wpa*vth
                     new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.charged.vth[iz,ir,is]
+                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1189,7 +1189,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
                         restart_vpa_vals =
                             @. vpa.grid +
-                               moments.charged.upar[iz,ir,is] / moments.charged.vth[iz,ir,is]
+                               moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
                         @views interpolate_to_grid_1d!(
                             new_pdf[:,ivperp,iz,is], restart_vpa_vals,
                             this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
@@ -1209,33 +1209,33 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                 if moments.evolve_density && !restart_evolve_density
                     # Need to normalise by density
                     for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] ./= moments.charged.dens[iz,ir,is]
+                        this_pdf[:,:,iz,is] ./= moments.ion.dens[iz,ir,is]
                     end
                 elseif !moments.evolve_density && restart_evolve_density
                     # Need to unnormalise by density
                     for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] .*= moments.charged.dens[iz,ir,is]
+                        this_pdf[:,:,iz,is] .*= moments.ion.dens[iz,ir,is]
                     end
                 end
                 if moments.evolve_ppar && !restart_evolve_ppar
                     # Need to normalise by vth
                     for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] .*= moments.charged.vth[iz,ir,is]
+                        this_pdf[:,:,iz,is] .*= moments.ion.vth[iz,ir,is]
                     end
                 elseif !moments.evolve_ppar && restart_evolve_ppar
                     # Need to unnormalise by vth
                     for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] ./= moments.charged.vth[iz,ir,is]
+                        this_pdf[:,:,iz,is] ./= moments.ion.vth[iz,ir,is]
                     end
                 end
 
                 return this_pdf
             end
 
-            boundary_distributions.pdf_rboundary_charged[:,:,:,1,:] .=
-                load_charged_boundary_pdf("pdf_rboundary_charged_left", 1)
-            boundary_distributions.pdf_rboundary_charged[:,:,:,2,:] .=
-                load_charged_boundary_pdf("pdf_rboundary_charged_right", r.n)
+            boundary_distributions.pdf_rboundary_ion[:,:,:,1,:] .=
+                load_ion_boundary_pdf("pdf_rboundary_ion_left", 1)
+            boundary_distributions.pdf_rboundary_ion[:,:,:,2,:] .=
+                load_ion_boundary_pdf("pdf_rboundary_ion_right", r.n)
 
             if composition.n_neutral_species > 0
                 moments.neutral.dens .= load_moment("density_neutral")
@@ -1846,11 +1846,11 @@ restarts (which are sequential in time), so concatenate the data from each entry
 
 The slice to take is specified by the keyword arguments.
 """
-function load_distributed_charged_pdf_slice(run_names::Tuple, nblocks::Tuple, t_range,
-                                            n_species::mk_int, r::coordinate,
-                                            z::coordinate, vperp::coordinate,
-                                            vpa::coordinate; is=nothing, ir=nothing,
-                                            iz=nothing, ivperp=nothing, ivpa=nothing)
+function load_distributed_ion_pdf_slice(run_names::Tuple, nblocks::Tuple, t_range,
+                                        n_species::mk_int, r::coordinate, z::coordinate,
+                                        vperp::coordinate, vpa::coordinate; is=nothing,
+                                        ir=nothing, iz=nothing, ivperp=nothing,
+                                        ivpa=nothing)
     # dimension of pdf is [vpa,vperp,z,r,species,t]
 
     result_dims = mk_int[]
diff --git a/src/makie_post_processing.jl b/src/makie_post_processing.jl
index 6cacb3a30..7d22ea184 100644
--- a/src/makie_post_processing.jl
+++ b/src/makie_post_processing.jl
@@ -30,7 +30,7 @@ using ..looping: all_dimensions, ion_dimensions, neutral_dimensions
 using ..manufactured_solns: manufactured_solutions, manufactured_electric_fields
 using ..moment_kinetics_input: mk_input
 using ..load_data: open_readonly_output_file, get_group, load_block_data,
-                   load_coordinate_data, load_distributed_charged_pdf_slice,
+                   load_coordinate_data, load_distributed_ion_pdf_slice,
                    load_distributed_neutral_pdf_slice, load_input, load_mk_options,
                    load_species_data, load_time_data
 using ..initial_conditions: vpagrid_to_dzdt
@@ -301,7 +301,7 @@ function makie_post_process(run_dir::Union{String,Tuple},
         end
     end
 
-    plot_charged_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
+    plot_ion_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
     if has_neutrals
         plot_neutral_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
     end
@@ -1287,14 +1287,12 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                                       run_info.r_local.n, run_info.itime_skip)
             result = result[iz,ir,is,it]
         elseif nd === 6
-            result = load_distributed_charged_pdf_slice(run_info.files, run_info.nblocks,
-                                                        it, run_info.n_ion_species,
-                                                        run_info.r_local,
-                                                        run_info.z_local, run_info.vperp,
-                                                        run_info.vpa;
-                                                        is=(is === (:) ? nothing : is),
-                                                        ir=ir, iz=iz, ivperp=ivperp,
-                                                        ivpa=ivpa)
+            result = load_distributed_ion_pdf_slice(run_info.files, run_info.nblocks, it,
+                                                    run_info.n_ion_species,
+                                                    run_info.r_local, run_info.z_local,
+                                                    run_info.vperp, run_info.vpa;
+                                                    is=(is === (:) ? nothing : is), ir=ir,
+                                                    iz=iz, ivperp=ivperp, ivpa=ivpa)
         elseif nd === 7
             result = load_distributed_neutral_pdf_slice(run_info.files, run_info.nblocks,
                                                         it, run_info.n_ion_species,
@@ -4796,9 +4794,9 @@ function animate_f_unnorm_vs_vpa_z(run_info; input=nothing, neutral=false, is=1,
 end
 
 """
-    plot_charged_pdf_2D_at_wall(run_info; plot_prefix)
+    plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
 
-Make plots/animations of the charged particle distribution function at wall boundaries.
+Make plots/animations of the ion distribution function at wall boundaries.
 
 The information for the runs to plot is passed in `run_info` (as returned by
 [`get_run_info`](@ref)). If `run_info` is a Tuple, comparison plots are made where line
@@ -4812,7 +4810,7 @@ will be saved with the format `plot_prefix<some_identifying_string>.pdf`. When `
 is not a Tuple, `plot_prefix` is optional - plots/animations will be saved only if it is
 passed.
 """
-function plot_charged_pdf_2D_at_wall(run_info; plot_prefix)
+function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
     input = Dict_to_NamedTuple(input_dict_dfns["wall_pdf"])
     if !(input.plot || input.animate)
         # nothing to do
@@ -4826,7 +4824,7 @@ function plot_charged_pdf_2D_at_wall(run_info; plot_prefix)
     z_lower = 1
     z_upper = run_info[1].z.n
     if !all(ri.z.n == z_upper for ri ∈ run_info)
-        println("Cannot run plot_charged_pdf_2D_at_wall() for runs with different "
+        println("Cannot run plot_ion_pdf_2D_at_wall() for runs with different "
                 * "z-grid sizes. Got $(Tuple(ri.z.n for ri ∈ run_info))")
         return nothing
     end
@@ -6238,7 +6236,7 @@ end
                    field_sym_label, norm_label, plot_dims, animate_dims)
 
 Utility function for making plots to avoid duplicated code in
-[`compare_charged_pdf_symbolic_test`](@ref) and
+[`compare_ion_pdf_symbolic_test`](@ref) and
 [`compare_neutral_pdf_symbolic_test`](@ref).
 
 The information for the run to analyse is passed in `run_info` (as returned by
@@ -6388,7 +6386,7 @@ function _MMS_pdf_plots(run_info, input, variable_name, plot_prefix, field_label
 end
 
 """
-    compare_charged_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
+    compare_ion_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
                                       input=nothing)
 
 Compare the computed and manufactured solutions for the ion distribution function.
@@ -6409,7 +6407,7 @@ Note: when calculating error norms, data is loaded only for 1 time point and for
 chunk that is the same size as computed by 1 block of the simulation at run time. This
 should prevent excessive memory requirements for this function.
 """
-function compare_charged_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
+function compare_ion_pdf_symbolic_test(run_info, plot_prefix; io=nothing,
                                            input=nothing)
 
     field_label = L"\tilde{f}_i"
@@ -6928,7 +6926,7 @@ function manufactured_solutions_analysis_dfns(run_info; plot_prefix)
         println_to_stdout_and_file(io, "# ", run_info.run_name)
         println_to_stdout_and_file(io, join(run_info.time, " "), " # time / (Lref/cref): ")
 
-        compare_charged_pdf_symbolic_test(run_info, plot_prefix; io=io, input=input)
+        compare_ion_pdf_symbolic_test(run_info, plot_prefix; io=io, input=input)
 
         if run_info.n_neutral_species > 0
             compare_neutral_pdf_symbolic_test(run_info, plot_prefix; io=io, input=input)
diff --git a/src/manufactured_solns.jl b/src/manufactured_solns.jl
index 3130bd674..21574a4c4 100644
--- a/src/manufactured_solns.jl
+++ b/src/manufactured_solns.jl
@@ -377,7 +377,7 @@ using IfElse
 
     function manufactured_solutions(manufactured_solns_input, Lr, Lz, r_bc, z_bc,
                                     geometry, composition, species, nr, nvperp)
-        charged_species = species.charged[1]
+        ion_species = species.ion[1]
         if composition.n_neutral_species > 0
             neutral_species = species.neutral[1]
         else
@@ -385,17 +385,17 @@ using IfElse
         end
 
         densi = densi_sym(Lr, Lz, r_bc, z_bc, composition, manufactured_solns_input,
-                          charged_species)
+                          ion_species)
         upari = upari_sym(Lr, Lz, r_bc, z_bc, composition, geometry, nr, manufactured_solns_input,
-                          charged_species)
+                          ion_species)
         ppari = ppari_sym(Lr, Lz, r_bc, z_bc, composition, manufactured_solns_input,
-                          charged_species)
+                          ion_species)
         pperpi = pperpi_sym(Lr, Lz, r_bc, z_bc, composition, manufactured_solns_input,
-                          charged_species, nvperp)
+                          ion_species, nvperp)
         vthi = vthi_sym(Lr, Lz, r_bc, z_bc, composition, manufactured_solns_input,
-                          charged_species, nvperp)
+                          ion_species, nvperp)
         dfni = dfni_sym(Lr, Lz, r_bc, z_bc, composition, geometry, nr,
-                        manufactured_solns_input, charged_species)
+                        manufactured_solns_input, ion_species)
 
         densn = densn_sym(Lr, Lz, r_bc, z_bc, geometry,composition,
                           manufactured_solns_input, neutral_species)
@@ -447,7 +447,7 @@ using IfElse
             vpa_coord, vzeta_coord, vr_coord, vz_coord, composition, geometry, collisions,
             num_diss_params, species)
 
-        charged_species = species.charged[1]
+        ion_species = species.ion[1]
         if composition.n_neutral_species > 0
             neutral_species = species.neutral[1]
         else
@@ -456,16 +456,16 @@ using IfElse
 
         # ion manufactured solutions
         densi = densi_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc, composition,
-                          manufactured_solns_input, charged_species)
-        upari = upari_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc, composition, geometry, r_coord.n, manufactured_solns_input, charged_species)
+                          manufactured_solns_input, ion_species)
+        upari = upari_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc, composition, geometry, r_coord.n, manufactured_solns_input, ion_species)
         vthi = vthi_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc, composition, manufactured_solns_input,
-                          charged_species, vperp_coord.n)
+                          ion_species, vperp_coord.n)
         dfni = dfni_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc, composition,
-                        geometry, r_coord.n, manufactured_solns_input, charged_species)
+                        geometry, r_coord.n, manufactured_solns_input, ion_species)
         #dfni in vr vz vzeta coordinates
         vrvzvzeta_dfni = cartesian_dfni_sym(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc,
                                             composition, manufactured_solns_input,
-                                            charged_species)
+                                            ion_species)
 
         # neutral manufactured solutions
         densn = densn_sym(r_coord.L,z_coord.L, r_coord.bc, z_coord.bc, geometry,
@@ -508,7 +508,7 @@ using IfElse
         # calculate the electric fields and the potential
         Er, Ez, phi = electric_fields(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc,
                                       composition, r_coord.n, manufactured_solns_input,
-                                      charged_species)
+                                      ion_species)
 
         # the ion source to maintain the manufactured solution
         Si = ( Dt(dfni) + ( vpa * (Bzed/Bmag) - 0.5*rhostar*Er ) * Dz(dfni) + ( 0.5*rhostar*Ez*rfac ) * Dr(dfni) + ( 0.5*Ez*Bzed/Bmag ) * Dvpa(dfni)
diff --git a/src/post_processing.jl b/src/post_processing.jl
index dc34391f6..bd1721445 100644
--- a/src/post_processing.jl
+++ b/src/post_processing.jl
@@ -40,7 +40,7 @@ using ..type_definitions: mk_float, mk_int
 using ..load_data: open_readonly_output_file, get_group, load_input, load_time_data
 using ..load_data: get_nranks
 using ..load_data: load_fields_data, load_pdf_data
-using ..load_data: load_distributed_charged_pdf_slice, load_distributed_neutral_pdf_slice, iglobal_func
+using ..load_data: load_distributed_ion_pdf_slice, load_distributed_neutral_pdf_slice, iglobal_func
 using ..load_data: load_neutral_pdf_data
 using ..load_data: load_variable
 using ..load_data: load_coordinate_data, load_block_data, load_rank_data,
@@ -303,7 +303,7 @@ function allocate_global_zr_ion_moments(nz_global,nr_global,n_ion_species,ntime)
     return density, parallel_flow, parallel_pressure, perpendicular_pressure, parallel_heat_flux, thermal_speed, entropy_production, chodura_integral_lower, chodura_integral_upper
 end
 
-function allocate_global_zr_charged_dfns(nvpa_global, nvperp_global, nz_global, nr_global,
+function allocate_global_zr_ion_dfns(nvpa_global, nvperp_global, nz_global, nr_global,
                                          n_ion_species, ntime)
     f = allocate_float(nvpa_global, nvperp_global, nz_global, nr_global, n_ion_species,
                        ntime)
@@ -735,7 +735,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
     perpendicular_pressure_at_pdf_timse, parallel_heat_flux_at_pdf_times,
     thermal_speed_at_pdf_times, entropy_production_at_pdf_times,
     chodura_integral_lower_at_pdf_times, chodura_integral_upper_at_pdf_times =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+        get_tuple_of_return_values(allocate_global_zr_ion_moments,
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r), n_ion_species,
                                    ntime_pdfs)
@@ -760,7 +760,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r), iskip_pdfs)
-    # charged particle moments
+    # ion particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density_at_pdf_times, "density",
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
@@ -841,7 +841,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
                    (vr === nothing ? true : (this_vr.n == 1 for this_vr ∈ vr))...])
     if is_1D1V
         # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-        ff = get_tuple_of_return_values(load_distributed_charged_pdf_slice, run_names,
+        ff = get_tuple_of_return_values(load_distributed_ion_pdf_slice, run_names,
                                         nblocks, itime_min_pdfs:iskip_pdfs:itime_max_pdfs,
                                         n_ion_species, r, z, vperp, vpa)
         if has_neutrals
@@ -1233,7 +1233,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
         if Maxwellian_diagnostic
             pressure = copy(parallel_pressure)
             @. pressure = (2.0*perpendicular_pressure + parallel_pressure)/3.0
-            ff = load_distributed_charged_pdf_slice(run_name, nblocks, 1:ntime, n_ion_species, r,
+            ff = load_distributed_ion_pdf_slice(run_name, nblocks, 1:ntime, n_ion_species, r,
                                                    z, vperp, vpa; iz=iz0, ir=ir0)
             plot_Maxwellian_diagnostic(ff[:,:,:,:], density[iz0,ir0,:,:],
              parallel_flow[iz0,ir0,:,:], thermal_speed[iz0,ir0,:,:], vpa.grid, vpa.wgts, 
@@ -1339,7 +1339,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
         compare_moments_symbolic_test(run_name_label,thermal_speed,vthi_sym,"ion",z_global.grid,r_global.grid,time,z_global.n,r_global.n,ntime,
          L"\widetilde{v}_{th,i}",L"\widetilde{v}_{th,i}^{sym}",L"\varepsilon(\widetilde{v}_{th,i})","vthi")
 
-        compare_charged_pdf_symbolic_test(run_name_label,manufactured_solns_list,"ion",
+        compare_ion_pdf_symbolic_test(run_name_label,manufactured_solns_list,"ion",
           L"\widetilde{f}_i",L"\widetilde{f}^{sym}_i",L"\varepsilon(\widetilde{f}_i)","pdf")
         if n_neutral_species > 0
             # neutral test
@@ -1435,7 +1435,7 @@ function calculate_differences(prefix...)
                                              ntime)
     density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed,
     chodura_integral_lower, chodura_integral_upper =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+        get_tuple_of_return_values(allocate_global_zr_ion_moments,
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r),
                                    n_ion_species, ntime)
@@ -1464,7 +1464,7 @@ function calculate_differences(prefix...)
                                nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r))
-    # charged particle moments
+    # ion particle moments
     get_tuple_of_return_values(read_distributed_zr_data!, density, "density", run_names,
                                "moments", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
@@ -1540,7 +1540,7 @@ function calculate_differences(prefix...)
     density_at_pdf_times, parallel_flow_at_pdf_times, parallel_pressure_at_pdf_times,
     parallel_heat_flux_at_pdf_times, thermal_speed_at_pdf_times, chodura_integral_lower_at_pdf_times,
     chodura_integral_upper_at_pdf_times =
-        get_tuple_of_return_values(allocate_global_zr_charged_moments,
+        get_tuple_of_return_values(allocate_global_zr_ion_moments,
                                    Tuple(this_z.n_global for this_z ∈ z),
                                    Tuple(this_r.n_global for this_r ∈ r), n_ion_species,
                                    ntime_pdfs)
@@ -1565,7 +1565,7 @@ function calculate_differences(prefix...)
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
                                Tuple(this_r.n for this_r ∈ r))
-    # charged particle moments
+    # ion moments
     get_tuple_of_return_values(read_distributed_zr_data!, density_at_pdf_times, "density",
                                run_names, "dfns", nblocks,
                                Tuple(this_z.n for this_z ∈ z),
@@ -1613,7 +1613,7 @@ function calculate_differences(prefix...)
     end
 
     # load full (vpa,z,r,species,t) particle distribution function (pdf) data
-    ff = get_tuple_of_return_values(load_distributed_charged_pdf_slice, run_names,
+    ff = get_tuple_of_return_values(load_distributed_ion_pdf_slice, run_names,
                                     nblocks, Tuple(1:nt for nt ∈ ntime_pdfs),
                                     n_ion_species, r, z, vperp, vpa)
     if maximum(n_neutral_species) > 0
@@ -3203,7 +3203,7 @@ function compare_ion_pdf_symbolic_test(run_name,manufactured_solns_list,spec_str
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time, _ = load_time_data(fid, printout=false)
     close(fid)
-    # get the charged particle pdf
+    # get the ion pdf
     dfni_func = manufactured_solns_list.dfni_func
     is = 1 # only one species supported currently
 
@@ -3285,7 +3285,7 @@ function compare_neutral_pdf_symbolic_test(run_name,manufactured_solns_list,spec
     # load time data (unique to pdf, may differ to moment values depending on user nwrite_dfns value)
     ntime, time, _ = load_time_data(fid, printout=false)
     close(fid)
-    # get the charged particle pdf
+    # get the ion pdf
     dfnn_func = manufactured_solns_list.dfnn_func
     is = 1 # only one species supported currently
 
@@ -3427,7 +3427,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
     end
     # make a gif animation of f(vperp,z,t) at a given (vpa,r) location
     if pp.animate_f_vs_vperp_z
-        pdf = load_distributed_charged_pdf_slice(run_name, nblocks,
+        pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa; ivpa=ivpa0,
                                                  ir=ir0)
@@ -3441,7 +3441,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
     end
     # make a gif animation of f(vperp,r,t) at a given (vpa,z) location
     if pp.animate_f_vs_vperp_r
-        pdf = load_distributed_charged_pdf_slice(run_name, nblocks,
+        pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa; ivpa=ivpa0,
                                                  iz=iz0)
@@ -3455,7 +3455,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
     end
     # make a gif animation of f(vpa,vperp,t) at a given (z,r) location
     if pp.animate_f_vs_vperp_vpa
-        pdf = load_distributed_charged_pdf_slice(run_name, nblocks,
+        pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa; iz=iz0,
                                                  ir=ir0)
@@ -3953,10 +3953,10 @@ function plot_ion_pdf_2D_at_wall(run_name, run_name_label, r_global, z_global,
     # Note only need final time point, so only load the one time point
     itime0 = 1
     # pdf at lower wall
-    pdf_lower = load_distributed_charged_pdf_slice(run_name, nblocks, ntime:ntime, n_ion_species, r,
+    pdf_lower = load_distributed_ion_pdf_slice(run_name, nblocks, ntime:ntime, n_ion_species, r,
                                                    z, vperp, vpa; iz=1)
     # pdf at upper wall
-    pdf_upper = load_distributed_charged_pdf_slice(run_name, nblocks, ntime:ntime, n_ion_species, r,
+    pdf_upper = load_distributed_ion_pdf_slice(run_name, nblocks, ntime:ntime, n_ion_species, r,
                                                    z, vperp, vpa; iz=z.n_global)
     for (pdf, zlabel) ∈ ((pdf_lower, "wall-"), (pdf_upper, "wall+"))
         for is in 1:n_ion_species
@@ -3973,7 +3973,7 @@ function plot_ion_pdf_2D_at_wall(run_name, run_name_label, r_global, z_global,
             outfile = string(run_name_label, "_pdf(vpa,vperp,iz_"*zlabel*",ir0)"*description*"vs_vperp_vpa.pdf")
             trysavefig(outfile)
 
-            # Skip this because load_distributed_charged_pdf_slice() currently only
+            # Skip this because load_distributed_ion_pdf_slice() currently only
             # handles selecting a single value like `iz=1`, not a sub-slice like
             # `iz=1:n_local`, so we only have data for one point in z here, so we can't
             # plot vs z.
@@ -3985,7 +3985,7 @@ function plot_ion_pdf_2D_at_wall(run_name, run_name_label, r_global, z_global,
 
             # plot f(ivpa0,ivperp0,z,r,is,itime) near the wall
             if r.n > 1
-                # Skip this because load_distributed_charged_pdf_slice() currently only
+                # Skip this because load_distributed_ion_pdf_slice() currently only
                 # handles selecting a single value like `iz=1`, not a sub-slice like
                 # `iz=1:n_local`, so we only have data for one point in z here, so we
                 # can't plot vs z.
diff --git a/src/source_terms.jl b/src/source_terms.jl
index 08fd91179..790603f34 100644
--- a/src/source_terms.jl
+++ b/src/source_terms.jl
@@ -42,7 +42,7 @@ function source_terms!(pdf_out, fvec_in, moments, vpa, z, r, dt, spectral, compo
             @views source_terms_evolve_density!(
                 pdf_out[:,:,:,:,is], fvec_in.pdf[:,:,:,:,is], fvec_in.density[:,:,is],
                 fvec_in.upar[:,:,is], moments.ion.ddens_dz[:,:,is],
-                moments.charged.dupar_dz[:,:,is], moments, z, r, dt, spectral,
+                moments.ion.dupar_dz[:,:,is], moments, z, r, dt, spectral,
                 ion_source_settings)
         end
     end
@@ -65,7 +65,7 @@ function source_terms_evolve_density!(pdf_out, pdf_in, dens, upar, ddens_dz, dup
     end
 
     if ion_source_settings.active
-        source_amplitude = moments.charged.external_source_amplitude
+        source_amplitude = moments.ion.external_source_amplitude
         @loop_r_z ir iz begin
             term = dt * source_amplitude[iz,ir] / dens[iz,ir]
             @loop_vperp_vpa ivperp ivpa begin
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 46914f31b..d7f079a03 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -292,7 +292,7 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
     # that may be evolved separately via fluid equations
-    scratch = setup_scratch_arrays(moments, pdf.charged.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    scratch = setup_scratch_arrays(moments, pdf.ion.norm, pdf.neutral.norm, t_input.n_rk_stages)
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
@@ -340,8 +340,8 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         # initialise the r advection speed
         begin_s_z_vperp_vpa_region()
         @loop_s is begin
-            @views update_speed_r!(r_advect[is], moments.charged.upar[:,:,is],
-                                   moments.charged.vth[:,:,is], fields, moments.evolve_upar,
+            @views update_speed_r!(r_advect[is], moments.ion.upar[:,:,is],
+                                   moments.ion.vth[:,:,is], fields, moments.evolve_upar,
                                    moments.evolve_ppar, vpa, vperp, z, r, geometry)
         end
         # enforce prescribed boundary condition in r on the distribution function f
@@ -356,8 +356,8 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         # initialise the z advection speed
         begin_s_r_vperp_vpa_region()
         @loop_s is begin
-            @views update_speed_z!(z_advect[is], moments.charged.upar[:,:,is],
-                                   moments.charged.vth[:,:,is], moments.evolve_upar,
+            @views update_speed_z!(z_advect[is], moments.ion.upar[:,:,is],
+                                   moments.ion.vth[:,:,is], moments.evolve_upar,
                                    moments.evolve_ppar, fields, vpa, vperp, z, r, 0.0,
                                    geometry)
         end
@@ -469,24 +469,24 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         # Ensure normalised pdf exactly obeys integral constraints if evolving moments
         begin_s_r_z_region()
         @loop_s_r_z is ir iz begin
-            @views hard_force_moment_constraints!(pdf.charged.norm[:,:,iz,ir,is], moments,
+            @views hard_force_moment_constraints!(pdf.ion.norm[:,:,iz,ir,is], moments,
                                                   vpa)
         end
         # update moments in case they were affected by applying boundary conditions or
         # constraints to the pdf
         reset_moments_status!(moments)
-        update_moments!(moments, pdf.charged.norm, vpa, vperp, z, r, composition)
+        update_moments!(moments, pdf.ion.norm, vpa, vperp, z, r, composition)
         # update the Chodura diagnostic -- note that the pdf should be the unnormalised one
         # so this will break for the split moments cases
-        update_chodura!(moments,pdf.charged.norm,vpa,vperp,z,r,r_spectral,composition,geometry,scratch_dummy,z_advect)
+        update_chodura!(moments,pdf.ion.norm,vpa,vperp,z,r,r_spectral,composition,geometry,scratch_dummy,z_advect)
         # enforce boundary conditions in r and z on the neutral particle distribution function
         if n_neutral_species > 0
             # Note, so far vr and vzeta do not need advect objects, so pass `nothing` for
             # those as a placeholder
             enforce_neutral_boundary_conditions!(
-                pdf.neutral.norm, pdf.charged.norm, boundary_distributions,
+                pdf.neutral.norm, pdf.ion.norm, boundary_distributions,
                 moments.neutral.dens, moments.neutral.uz, moments.neutral.pz, moments,
-                moments.charged.dens, moments.charged.upar, fields.Er, vzeta_spectral,
+                moments.ion.dens, moments.ion.upar, fields.Er, vzeta_spectral,
                 vr_spectral, vz_spectral, neutral_r_advect, neutral_z_advect, nothing,
                 nothing, neutral_vz_advect, r, z, vzeta, vr, vz, composition, geometry,
                 scratch_dummy, advance.r_diffusion, advance.vz_diffusion)
@@ -1133,7 +1133,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="vpa", ylabel="z", title="f", width=ax_width, height=ax_height)
-                    plot_2d(vpa.grid, z.grid, pdf.charged.norm[:,1,:,1,1]; ax=ax, colorbar_place=layout[1,2])
+                    plot_2d(vpa.grid, z.grid, pdf.ion.norm[:,1,:,1,1]; ax=ax, colorbar_place=layout[1,2])
                     if composition.n_neutral_species > 0
                         ax = Axis(layout[1,3], xlabel="vz", ylabel="z", title="f_neutral", width=ax_width, height=ax_height)
                         plot_2d(vz.grid, z.grid, pdf.neutral.norm[:,1,1,:,1,1]; ax=ax, colorbar_place=layout[1,4])
@@ -1142,7 +1142,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="vpa", ylabel="z", title="f", width=ax_width, height=ax_height)
-                    plot_2d(vpa.grid, z.grid, pdf.charged.norm[:,1,:,1,1]; ax=ax,
+                    plot_2d(vpa.grid, z.grid, pdf.ion.norm[:,1,:,1,1]; ax=ax,
                             colorbar_place=layout[1,2], colorscale=log10,
                             transform=x->positive_or_nan(x, epsilon=1.e-20))
                     if composition.n_neutral_species > 0
@@ -1155,7 +1155,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="vpa", ylabel="f0", width=ax_width, height=ax_height)
-                    plot_1d(vpa.grid, pdf.charged.norm[:,1,1,1,1]; ax=ax)
+                    plot_1d(vpa.grid, pdf.ion.norm[:,1,1,1,1]; ax=ax)
                     if composition.n_neutral_species > 0
                         ax = Axis(layout[1,2], xlabel="vz", ylabel="f0_neutral", width=ax_width, height=ax_height)
                         plot_1d(vz.grid, pdf.neutral.norm[:,1,1,1,1,1]; ax=ax)
@@ -1164,7 +1164,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="vpa", ylabel="fL", width=ax_width, height=ax_height)
-                    plot_1d(vpa.grid, pdf.charged.norm[:,1,end,1,1]; ax=ax)
+                    plot_1d(vpa.grid, pdf.ion.norm[:,1,end,1,1]; ax=ax)
                     if composition.n_neutral_species > 0
                         ax = Axis(layout[1,2], xlabel="vz", ylabel="fL_neutral", width=ax_width, height=ax_height)
                         plot_1d(vz.grid, pdf.neutral.norm[:,1,1,end,1,1]; ax=ax)
@@ -1173,13 +1173,13 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="z", ylabel="density", width=ax_width, height=ax_height)
-                    plot_1d(z.grid, moments.charged.dens[:,1,1]; ax=ax, label="ion")
+                    plot_1d(z.grid, moments.ion.dens[:,1,1]; ax=ax, label="ion")
                     if composition.n_neutral_species > 0
                         plot_1d(z.grid, moments.neutral.dens[:,1,1]; ax=ax, label="neutral")
                     end
                     #axislegend(ax)
                     ax = Axis(layout[1,2], xlabel="z", ylabel="upar", width=ax_width, height=ax_height)
-                    plot_1d(z.grid, moments.charged.upar[:,1,1]; ax=ax, label="ion")
+                    plot_1d(z.grid, moments.ion.upar[:,1,1]; ax=ax, label="ion")
                     if composition.n_neutral_species > 0
                         plot_1d(z.grid, moments.neutral.uz[:,1,1]; ax=ax, label="neutral")
                     end
@@ -1188,13 +1188,13 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="z", ylabel="ppar", width=ax_width, height=ax_height)
-                    plot_1d(z.grid, moments.charged.ppar[:,1,1]; ax=ax, label="ion")
+                    plot_1d(z.grid, moments.ion.ppar[:,1,1]; ax=ax, label="ion")
                     if composition.n_neutral_species > 0
                         plot_1d(z.grid, moments.neutral.pz[:,1,1]; ax=ax, label="neutral")
                     end
                     #axislegend(ax)
                     ax = Axis(layout[1,2], xlabel="z", ylabel="vth", width=ax_width, height=ax_height)
-                    plot_1d(z.grid, moments.charged.vth[:,1,1]; ax=ax, label="ion")
+                    plot_1d(z.grid, moments.ion.vth[:,1,1]; ax=ax, label="ion")
                     if composition.n_neutral_species > 0
                         plot_1d(z.grid, moments.neutral.vth[:,1,1]; ax=ax, label="neutral")
                     end
@@ -1203,7 +1203,7 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     irow += 1
                     layout = fig[irow,1] = GridLayout()
                     ax = Axis(layout[1,1], xlabel="z", ylabel="qpar", width=ax_width, height=ax_height)
-                    plot_1d(z.grid, moments.charged.qpar[:,1,1]; ax=ax, label="ion")
+                    plot_1d(z.grid, moments.ion.qpar[:,1,1]; ax=ax, label="ion")
                     if composition.n_neutral_species > 0
                         plot_1d(z.grid, moments.neutral.qz[:,1,1]; ax=ax, label="neutral")
                     end
@@ -1528,8 +1528,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         rethrow(e)
     end
     # update the parallel heat flux
-    update_qpar!(moments.charged.qpar, moments.charged.qpar_updated, new_scratch.density,
-                 new_scratch.upar, moments.charged.vth, new_scratch.pdf, vpa, vperp, z, r,
+    update_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
+                 new_scratch.upar, moments.ion.vth, new_scratch.pdf, vpa, vperp, z, r,
                  composition, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
 
@@ -2099,7 +2099,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     # advance with the Fokker-Planck self-collision operator
     if advance.explicit_weakform_fp_collisions
         update_entropy_diagnostic = (istage == 1)
-        explicit_fokker_planck_collisions_weak_form!(fvec_out.pdf,fvec_in.pdf,moments.charged.dSdt,composition,collisions,dt,
+        explicit_fokker_planck_collisions_weak_form!(fvec_out.pdf,fvec_in.pdf,moments.ion.dSdt,composition,collisions,dt,
                                              fp_arrays,r,z,vperp,vpa,vperp_spectral,vpa_spectral,scratch_dummy,
                                              diagnose_entropy_production = update_entropy_diagnostic)
     end
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index 0c6a5eca4..8de992fda 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -393,7 +393,7 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
     end
 
     # return struct containing arrays needed to update moments
-    return moments_charged_substruct(density, density_updated, parallel_flow,
+    return moments_ion_substruct(density, density_updated, parallel_flow,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,perpendicular_pressure,
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
         chodura_integral_lower, chodura_integral_upper, v_norm_fac,
@@ -636,7 +636,7 @@ function update_moments!(moments, ff, vpa, vperp, z, r, composition)
                                         moments.evolve_upar)
             moments.ion.ppar_updated[is] = true
         end
-        @views update_pperp_species!(moments.charged.pperp[:,:,is], ff[:,:,:,:,is], vpa, vperp, z, r)
+        @views update_pperp_species!(moments.ion.pperp[:,:,is], ff[:,:,:,:,is], vpa, vperp, z, r)
         if moments.ion.qpar_updated[is] == false
             @views update_qpar_species!(moments.ion.qpar[:,:,is],
                                         moments.ion.dens[:,:,is],
@@ -644,10 +644,10 @@ function update_moments!(moments, ff, vpa, vperp, z, r, composition)
                                         moments.ion.vth[:,:,is], ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
                                         moments.evolve_upar, moments.evolve_ppar)
-            moments.charged.qpar_updated[is] = true
+            moments.ion.qpar_updated[is] = true
         end
     end
-    update_vth!(moments.charged.vth, moments.charged.ppar, moments.charged.pperp, moments.charged.dens, vperp, z, r, composition)
+    update_vth!(moments.ion.vth, moments.ion.ppar, moments.ion.pperp, moments.ion.dens, vperp, z, r, composition)
     return nothing
 end
 
@@ -987,22 +987,22 @@ function update_chodura!(moments,ff,vpa,vperp,z,r,r_spectral,composition,geometr
     begin_s_r_region()
     if z.irank == 0
         @loop_s_r is ir begin
-            @views moments.charged.chodura_integral_lower[ir,is] = update_chodura_integral_species!(ff[:,:,1,ir,is],dffdr[:,:,1,ir,is],
-            ff_dummy[:,:,1,ir,is],vpa,vperp,z,r,composition,geometry,z_advect[is].speed[1,:,:,ir],moments.charged.dens[1,ir,is],del_vpa)
+            @views moments.ion.chodura_integral_lower[ir,is] = update_chodura_integral_species!(ff[:,:,1,ir,is],dffdr[:,:,1,ir,is],
+            ff_dummy[:,:,1,ir,is],vpa,vperp,z,r,composition,geometry,z_advect[is].speed[1,:,:,ir],moments.ion.dens[1,ir,is],del_vpa)
         end
     else # we do not save this Chodura integral to the output file
         @loop_s_r is ir begin
-            moments.charged.chodura_integral_lower[ir,is] = 0.0
+            moments.ion.chodura_integral_lower[ir,is] = 0.0
         end
     end
     if z.irank == z.nrank - 1
         @loop_s_r is ir begin
-            @views moments.charged.chodura_integral_upper[ir,is] = update_chodura_integral_species!(ff[:,:,end,ir,is],dffdr[:,:,end,ir,is],
-            ff_dummy[:,:,end,ir,is],vpa,vperp,z,r,composition,geometry,z_advect[is].speed[end,:,:,ir],moments.charged.dens[end,ir,is],del_vpa)
+            @views moments.ion.chodura_integral_upper[ir,is] = update_chodura_integral_species!(ff[:,:,end,ir,is],dffdr[:,:,end,ir,is],
+            ff_dummy[:,:,end,ir,is],vpa,vperp,z,r,composition,geometry,z_advect[is].speed[end,:,:,ir],moments.ion.dens[end,ir,is],del_vpa)
         end
     else # we do not save this Chodura integral to the output file
         @loop_s_r is ir begin
-            moments.charged.chodura_integral_upper[ir,is] =  0.0
+            moments.ion.chodura_integral_upper[ir,is] =  0.0
         end
     end
 end
@@ -1058,8 +1058,8 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
         density = @view scratch.density[:,:,is]
         upar = @view scratch.upar[:,:,is]
         ppar = @view scratch.ppar[:,:,is]
-        qpar = @view moments.charged.qpar[:,:,is]
-        vth = @view moments.charged.vth[:,:,is]
+        qpar = @view moments.ion.qpar[:,:,is]
+        vth = @view moments.ion.vth[:,:,is]
         dummy_zr = @view scratch_dummy.dummy_zrs[:,:,is]
         buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,is]
         buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,is]
@@ -1068,13 +1068,13 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
         buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,is]
         buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,is]
         if moments.evolve_density
-            @views derivative_z!(moments.charged.ddens_dz[:,:,is], density, buffer_r_1,
+            @views derivative_z!(moments.ion.ddens_dz[:,:,is], density, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
             # Upwinded using upar as advection velocity, to be used in continuity equation
             @loop_r_z ir iz begin
                 dummy_zr[iz,ir] = -upar[iz,ir]
             end
-            @views derivative_z!(moments.charged.ddens_dz_upwind[:,:,is], density,
+            @views derivative_z!(moments.ion.ddens_dz_upwind[:,:,is], density,
                                  dummy_zr, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                                  buffer_r_5, buffer_r_6, z_spectral, z)
         end
@@ -1084,11 +1084,11 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
             # centred second derivative for dissipation
             @views derivative_z!(dummy_zr, density, buffer_r_1, buffer_r_2, buffer_r_3,
                                  buffer_r_4, z_spectral, z)
-            @views derivative_z!(moments.charged.d2dens_dz2[:,:,is], dummy_zr, buffer_r_1,
+            @views derivative_z!(moments.ion.d2dens_dz2[:,:,is], dummy_zr, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
-            @views derivative_z!(moments.charged.dupar_dz[:,:,is], upar, buffer_r_1,
+            @views derivative_z!(moments.ion.dupar_dz[:,:,is], upar, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
         if moments.evolve_upar
diff --git a/src/vpa_advection.jl b/src/vpa_advection.jl
index fe3afa8b0..7ae70840d 100644
--- a/src/vpa_advection.jl
+++ b/src/vpa_advection.jl
@@ -107,9 +107,9 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
             # • -wpar^2 * d(vth)/dz term
             @loop_z_vperp iz ivperp begin
                 @views @. advect[is].speed[:,ivperp,iz,ir] =
-                    moments.charged.dppar_dz[iz,ir,is]/(fvec.density[iz,ir,is]*moments.charged.vth[iz,ir,is]) +
-                    0.5*vpa.grid*moments.charged.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] -
-                    vpa.grid^2*moments.charged.dvth_dz[iz,ir,is]
+                    moments.ion.dppar_dz[iz,ir,is]/(fvec.density[iz,ir,is]*moments.ion.vth[iz,ir,is]) +
+                    0.5*vpa.grid*moments.ion.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] -
+                    vpa.grid^2*moments.ion.dvth_dz[iz,ir,is]
             end
         end
     end
@@ -139,17 +139,17 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
                             / fvec.ppar[iz,ir,is])
                      - fvec.density_neutral[iz,ir,is]
                        * (fvec.uz_neutral[iz,ir,is] - fvec.upar[iz,ir,is])
-                       / moments.charged.vth[iz,ir,is])
+                       / moments.ion.vth[iz,ir,is])
             end
         end
     end
     if ion_source_settings.active
-        source_amplitude = moments.charged.external_source_amplitude
+        source_amplitude = moments.ion.external_source_amplitude
         source_T = ion_source_settings.source_T
         density = fvec.density
         upar = fvec.upar
         ppar = fvec.ppar
-        vth = moments.charged.vth
+        vth = moments.ion.vth
         vpa_grid = vpa.grid
         @loop_s_r_z is ir iz begin
             prefactor = source_amplitude[iz,ir]
@@ -181,9 +181,9 @@ function update_speed_n_p_evolution!(advect, fields, fvec, moments, vpa, z, r,
             # • vpahat*d(upar)/dz
             # • -(1/2)*(dphi/dz)/vthi
             @loop_z_vperp iz ivperp begin
-                @views @. advect[is].speed[:,ivperp,iz,ir] = 0.5*vpa.grid*moments.charged.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] +
-                                                             vpa.grid*moments.charged.dvth_dz[iz] * (fvec.upar[iz,ir,is]/moments.vth[iz,ir,is] - vpa.grid) +
-                                                             vpa.grid*moments.charged.dupar_dz[iz,ir,is] +
+                @views @. advect[is].speed[:,ivperp,iz,ir] = 0.5*vpa.grid*moments.ion.dqpar_dz[iz,ir,is]/fvec.ppar[iz,ir,is] +
+                                                             vpa.grid*moments.ion.dvth_dz[iz] * (fvec.upar[iz,ir,is]/moments.vth[iz,ir,is] - vpa.grid) +
+                                                             vpa.grid*moments.ion.dupar_dz[iz,ir,is] +
                                                              0.5*fields.Ez[iz,ir]/moments.vth[iz,ir,is]
             end
         end
@@ -223,7 +223,7 @@ function update_speed_n_u_evolution!(advect, fvec, moments, vpa, z, r, compositi
             @loop_z_vperp iz ivperp begin
                 @views @. advect[is].speed[:,ivperp,iz,ir] =
                     moments.ion.dppar_dz[iz,ir,is]/fvec.density[iz,ir,is] -
-                    vpa.grid*moments.charged.dupar_dz[iz,ir,is]
+                    vpa.grid*moments.ion.dupar_dz[iz,ir,is]
             end
         end
     end
@@ -247,14 +247,14 @@ function update_speed_n_u_evolution!(advect, fvec, moments, vpa, z, r, compositi
         end
     end
     if ion_source_settings.active
-        source_amplitude = moments.charged.external_source_amplitude
+        source_amplitude = moments.ion.external_source_amplitude
         source_strength = ion_source_settings.source_strength
         source_T = ion_source_settings.source_T
         r_amplitude = ion_source_settings.r_amplitude
         z_amplitude = ion_source_settings.z_amplitude
         density = fvec.density
         upar = fvec.upar
-        vth = moments.charged.vth
+        vth = moments.ion.vth
         @loop_s_r_z is ir iz begin
             term = source_amplitude[iz,ir] * upar[iz,ir,is] / density[iz,ir,is]
             @loop_vperp_vpa ivperp ivpa begin
diff --git a/test/Krook_collisions_tests.jl b/test/Krook_collisions_tests.jl
index bccdf539a..ec247545a 100644
--- a/test/Krook_collisions_tests.jl
+++ b/test/Krook_collisions_tests.jl
@@ -10,7 +10,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_charged_particle_moments_data, load_pdf_data,
+                                 load_ion_particle_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -35,48 +35,48 @@ const expected =
           -0.353459763648872 -0.3755795658392631; -0.5494724768120175 -0.5904511161957432;
           -0.8609860698164502 -0.8726509017405427; -1.2115129555832849 -1.1306145497660032;
           -1.3862820803244258 -1.2382381134436695],
-     n_charged=[0.2500030702177186 0.28989452952580286; 0.2977473631375158 0.32283464906590775;
-                0.42274585818529853 0.41784232850006636; 0.5772542465450629 0.5540775204094593;
-                0.7022542481909738 0.6868914177534788; 0.7499999999999394 0.7468272160708606;
-                0.7022542481909738 0.6868914177534787; 0.577254246545063 0.554077520409459;
-                0.42274585818529864 0.4178423285000665; 0.2977473631375159 0.3228346490659078;
-                0.2500030702177185 0.2898945295258028],
+     n_ion=[0.2500030702177186 0.28989452952580286; 0.2977473631375158 0.32283464906590775;
+            0.42274585818529853 0.41784232850006636; 0.5772542465450629 0.5540775204094593;
+            0.7022542481909738 0.6868914177534788; 0.7499999999999394 0.7468272160708606;
+            0.7022542481909738 0.6868914177534787; 0.577254246545063 0.554077520409459;
+            0.42274585818529864 0.4178423285000665; 0.2977473631375159 0.3228346490659078;
+            0.2500030702177185 0.2898945295258028],
      n_neutral=[0.7499999999999382 0.7736770941648626; 0.7022542481909748 0.7056867075427516;
                 0.5772542465450632 0.5582975660019874; 0.4227458581852985 0.4096913953484598;
                 0.29774736313751604 0.3053964124252619; 0.2500030702177186 0.2681998023548167;
                 0.29774736313751604 0.3053964124252619; 0.42274585818529836 0.4096913953484599;
                 0.5772542465450631 0.5582975660019875; 0.7022542481909745 0.7056867075427524;
                 0.7499999999999383 0.7736770941648626],
-     upar_charged=[-2.7135787559953277e-17 -1.6845791254993525e-16; -9.321028970172899e-18 -0.18245939812953485;
-                   -2.8374879811351724e-18 -0.19666454846377826; 1.2124327390522635e-17 -0.11128043369942339;
-                   3.6525788403693063e-17 -0.03317985705380149; -2.0930856430671915e-17 4.720175801869314e-17;
-                   8.753545920086251e-18 0.033179857053801595; 1.1293771270243255e-17 0.11128043369942343;
-                   1.3739171132886587e-17 0.19666454846377784; -6.840453743089351e-18 0.18245939812953468;
-                   -2.7135787559953277e-17 -1.9129596434811267e-16],
+     upar_ion=[-2.7135787559953277e-17 -1.6845791254993525e-16; -9.321028970172899e-18 -0.18245939812953485;
+               -2.8374879811351724e-18 -0.19666454846377826; 1.2124327390522635e-17 -0.11128043369942339;
+               3.6525788403693063e-17 -0.03317985705380149; -2.0930856430671915e-17 4.720175801869314e-17;
+               8.753545920086251e-18 0.033179857053801595; 1.1293771270243255e-17 0.11128043369942343;
+               1.3739171132886587e-17 0.19666454846377784; -6.840453743089351e-18 0.18245939812953468;
+               -2.7135787559953277e-17 -1.9129596434811267e-16],
      upar_neutral=[6.5569385065066925e-18 8.08747058038406e-18; 1.1054500872839027e-17 -0.03620988455458174;
                    -3.241833393685864e-17 -0.009156078199383568; -3.617637280460899e-17 0.05452623197292568;
                    4.417578961284041e-17 0.07607875911384775; 4.9354467746194965e-17 1.635044638743921e-16;
                    6.573091229872379e-18 -0.0760787591138477; 2.989662686945165e-17 -0.05452623197292564;
                    -3.1951996361666834e-17 0.009156078199383685; -4.395464518158184e-18 0.03620988455458165;
                    6.5569385065066925e-18 1.8232586069007834e-18],
-     ppar_charged=[0.18749999999999992 0.23302732230115558; 0.20909325514551116 0.21936799130257528;
-                   0.24403180771238264 0.20856296024163393; 0.24403180771238278 0.2154266357557397;
-                   0.2090932551455113 0.2206183912107678; 0.1875 0.21979739387340663;
-                   0.20909325514551128 0.22061839121076784; 0.2440318077123828 0.21542663575573945;
-                   0.24403180771238256 0.20856296024163395; 0.20909325514551116 0.2193679913025754;
-                   0.18749999999999992 0.23302732230115553],
+     ppar_ion=[0.18749999999999992 0.23302732230115558; 0.20909325514551116 0.21936799130257528;
+               0.24403180771238264 0.20856296024163393; 0.24403180771238278 0.2154266357557397;
+               0.2090932551455113 0.2206183912107678; 0.1875 0.21979739387340663;
+               0.20909325514551128 0.22061839121076784; 0.2440318077123828 0.21542663575573945;
+               0.24403180771238256 0.20856296024163395; 0.20909325514551116 0.2193679913025754;
+               0.18749999999999992 0.23302732230115553],
      ppar_neutral=[0.18750000000000003 0.2480292382671593; 0.20909325514551122 0.24401255100297964;
                    0.24403180771238286 0.22861763406831279; 0.24403180771238278 0.2058922545451891;
                    0.20909325514551144 0.1926313699453636; 0.18749999999999992 0.19090651730415983;
                    0.20909325514551141 0.19263136994536365; 0.2440318077123828 0.20589225454518903;
                    0.24403180771238286 0.2286176340683127; 0.20909325514551114 0.24401255100297964;
                    0.18750000000000006 0.24802923826715936],
-     f_charged=[0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;
-                0.20411991941198782 0.25123395823993105 0.3934413727192304 0.6277900619432855 0.9100364506661008 1.0606601717796504 0.910036450666101 0.6277900619432859 0.39344137271923046 0.25123395823993094 0.20411991941198776;
-                0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;;;
-                0.0538996852594264 0.06066864433237418 0.03746866696438989 0.014783440166032301 0.010917691665145668 0.018422971878502774 0.027170953411068444 0.027269146560166702 0.026567569739750264 0.035612674100528624 0.05389968525942639;
-                0.2118369019176154 0.24917436308523389 0.37345448114678914 0.5972219245577428 0.8859681860177208 1.0485988935814787 0.8859681860177204 0.5972219245577435 0.37345448114678825 0.24917436308523389 0.2118369019176155;
-                0.05389968525942635 0.03561267410052869 0.02656756973975021 0.02726914656016675 0.027170953411068514 0.018422971878502753 0.01091769166514568 0.014783440166032254 0.037468666964389795 0.060668644332374164 0.05389968525942635],
+     f_ion=[0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;
+            0.20411991941198782 0.25123395823993105 0.3934413727192304 0.6277900619432855 0.9100364506661008 1.0606601717796504 0.910036450666101 0.6277900619432859 0.39344137271923046 0.25123395823993094 0.20411991941198776;
+            0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;;;
+            0.0538996852594264 0.06066864433237418 0.03746866696438989 0.014783440166032301 0.010917691665145668 0.018422971878502774 0.027170953411068444 0.027269146560166702 0.026567569739750264 0.035612674100528624 0.05389968525942639;
+            0.2118369019176154 0.24917436308523389 0.37345448114678914 0.5972219245577428 0.8859681860177208 1.0485988935814787 0.8859681860177204 0.5972219245577435 0.37345448114678825 0.24917436308523389 0.2118369019176155;
+            0.05389968525942635 0.03561267410052869 0.02656756973975021 0.02726914656016675 0.027170953411068514 0.018422971878502753 0.01091769166514568 0.014783440166032254 0.037468666964389795 0.060668644332374164 0.05389968525942635],
      f_neutral=[0.0063385294703834595 0.012360459026988546 0.030398267195914108 0.04284314198717859 0.040599270638920985 0.03704623609948259 0.040599270638920965 0.0428431419871786 0.030398267195914094 0.012360459026988546 0.006338529470383456;
                 1.0606601717796493 0.9100364506661016 0.6277900619432857 0.3934413727192303 0.2512339582399308 0.20411991941198754 0.2512339582399307 0.3934413727192301 0.6277900619432853 0.9100364506661016 1.0606601717796487;
                 0.0063385294703834595 0.012360459026988546 0.030398267195914108 0.04284314198717859 0.040599270638920985 0.03704623609948259 0.040599270638920965 0.0428431419871786 0.030398267195914094 0.012360459026988546 0.006338529470383456;;;
@@ -193,10 +193,10 @@ function run_test(test_input, rtol, atol; args...)
     end
 
     phi = nothing
-    n_charged = nothing
-    upar_charged = nothing
-    ppar_charged = nothing
-    f_charged = nothing
+    n_ion = nothing
+    upar_ion = nothing
+    ppar_ion = nothing
+    f_ion = nothing
     n_neutral = nothing
     upar_neutral = nothing
     ppar_neutral = nothing
@@ -224,7 +224,7 @@ function run_test(test_input, rtol, atol; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, qpar_charged_zrst, v_t_charged_zrst = load_charged_particle_moments_data(fid)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_particle_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
             z, z_spectral = load_coordinate_data(fid, "z")
 
@@ -234,19 +234,19 @@ function run_test(test_input, rtol, atol; args...)
             fid = open_readonly_output_file(path, "dfns")
             
             # load particle distribution function (pdf) data
-            f_charged_vpavperpzrst = load_pdf_data(fid)
+            f_ion_vpavperpzrst = load_pdf_data(fid)
             f_neutral_vzvrvzetazrst = load_neutral_pdf_data(fid)
             vpa, vpa_spectral = load_coordinate_data(fid, "vpa")
 
             close(fid)
             
             phi = phi_zrt[:,1,:]
-            n_charged = n_charged_zrst[:,1,:,:]
-            upar_charged = upar_charged_zrst[:,1,:,:]
-            ppar_charged = ppar_charged_zrst[:,1,:,:]
-            qpar_charged = qpar_charged_zrst[:,1,:,:]
-            v_t_charged = v_t_charged_zrst[:,1,:,:]
-            f_charged = f_charged_vpavperpzrst[:,1,:,1,:,:]
+            n_ion = n_ion_zrst[:,1,:,:]
+            upar_ion = upar_ion_zrst[:,1,:,:]
+            ppar_ion = ppar_ion_zrst[:,1,:,:]
+            qpar_ion = qpar_ion_zrst[:,1,:,:]
+            v_t_ion = v_t_ion_zrst[:,1,:,:]
+            f_ion = f_ion_vpavperpzrst[:,1,:,1,:,:]
             n_neutral = n_neutral_zrst[:,1,:,:]
             upar_neutral = upar_neutral_zrst[:,1,:,:]
             ppar_neutral = ppar_neutral_zrst[:,1,:,:]
@@ -257,7 +257,7 @@ function run_test(test_input, rtol, atol; args...)
             # Unnormalize f
             if input["evolve_moments_density"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] .*= n_charged[iz,is,it]
+                    f_ion[:,iz,is,it] .*= n_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] .*= n_neutral[iz,isn,it]
@@ -265,7 +265,7 @@ function run_test(test_input, rtol, atol; args...)
             end
             if input["evolve_moments_parallel_pressure"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] ./= v_t_charged[iz,is,it]
+                    f_ion[:,iz,is,it] ./= v_t_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] ./= v_t_neutral[iz,isn,it]
@@ -283,11 +283,11 @@ function run_test(test_input, rtol, atol; args...)
         #println("phi ", size(newgrid_phi))
         #println(newgrid_phi)
         #println()
-        #newgrid_n_charged = cat(interpolate_to_grid_z(expected.z, n_charged[:, :, 1], z, z_spectral)[:,1],
-        #                        interpolate_to_grid_z(expected.z, n_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_n_ion = cat(interpolate_to_grid_z(expected.z, n_ion[:, :, 1], z, z_spectral)[:,1],
+        #                        interpolate_to_grid_z(expected.z, n_ion[:, :, 2], z, z_spectral)[:,1];
         #                        dims=2)
-        #println("n_charged ", size(newgrid_n_charged))
-        #println(newgrid_n_charged)
+        #println("n_ion ", size(newgrid_n_ion))
+        #println(newgrid_n_ion)
         #println()
         #newgrid_n_neutral = cat(interpolate_to_grid_z(expected.z, n_neutral[:, :, 1], z, z_spectral)[:,1],
         #                        interpolate_to_grid_z(expected.z, n_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -295,11 +295,11 @@ function run_test(test_input, rtol, atol; args...)
         #println("n_neutral ", size(newgrid_n_neutral))
         #println(newgrid_n_neutral)
         #println()
-        #newgrid_upar_charged = cat(interpolate_to_grid_z(expected.z, upar_charged[:, :, 1], z, z_spectral)[:,1],
-        #                           interpolate_to_grid_z(expected.z, upar_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_upar_ion = cat(interpolate_to_grid_z(expected.z, upar_ion[:, :, 1], z, z_spectral)[:,1],
+        #                           interpolate_to_grid_z(expected.z, upar_ion[:, :, 2], z, z_spectral)[:,1];
         #                           dims=2)
-        #println("upar_charged ", size(newgrid_upar_charged))
-        #println(newgrid_upar_charged)
+        #println("upar_ion ", size(newgrid_upar_ion))
+        #println(newgrid_upar_ion)
         #println()
         #newgrid_upar_neutral = cat(interpolate_to_grid_z(expected.z, upar_neutral[:, :, 1], z, z_spectral)[:,1],
         #                           interpolate_to_grid_z(expected.z, upar_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -307,11 +307,11 @@ function run_test(test_input, rtol, atol; args...)
         #println("upar_neutral ", size(newgrid_upar_neutral))
         #println(newgrid_upar_neutral)
         #println()
-        #newgrid_ppar_charged = cat(interpolate_to_grid_z(expected.z, ppar_charged[:, :, 1], z, z_spectral)[:,1],
-        #                           interpolate_to_grid_z(expected.z, ppar_charged[:, :, 2], z, z_spectral)[:,1];
+        #newgrid_ppar_ion = cat(interpolate_to_grid_z(expected.z, ppar_ion[:, :, 1], z, z_spectral)[:,1],
+        #                           interpolate_to_grid_z(expected.z, ppar_ion[:, :, 2], z, z_spectral)[:,1];
         #                           dims=2)
-        #println("ppar_charged ", size(newgrid_ppar_charged))
-        #println(newgrid_ppar_charged)
+        #println("ppar_ion ", size(newgrid_ppar_ion))
+        #println(newgrid_ppar_ion)
         #println()
         #newgrid_ppar_neutral = cat(interpolate_to_grid_z(expected.z, ppar_neutral[:, :, 1], z, z_spectral)[:,1],
         #                           interpolate_to_grid_z(expected.z, ppar_neutral[:, :, 2], z, z_spectral)[:,1];
@@ -319,11 +319,11 @@ function run_test(test_input, rtol, atol; args...)
         #println("ppar_neutral ", size(newgrid_ppar_neutral))
         #println(newgrid_ppar_neutral)
         #println()
-        #newgrid_f_charged = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_charged[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
-        #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_charged[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
+        #newgrid_f_ion = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_ion[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
+        #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_ion[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
         #                        dims=4)
-        #println("f_charged ", size(newgrid_f_charged))
-        #println(newgrid_f_charged)
+        #println("f_ion ", size(newgrid_f_ion))
+        #println(newgrid_f_ion)
         #println()
         #newgrid_f_neutral = cat(interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_neutral[:, :, :, 1], z, z_spectral), vpa, vpa_spectral)[:,:,1],
         #                        interpolate_to_grid_vpa(expected.vpa, interpolate_to_grid_z(expected.z, f_neutral[:, :, :, 2], z, z_spectral), vpa, vpa_spectral)[:,:,1];
@@ -336,36 +336,36 @@ function run_test(test_input, rtol, atol; args...)
                 newgrid_phi = interpolate_to_grid_z(expected.z, phi[:, tind], z, z_spectral)
                 @test isapprox(expected.phi[:, tind], newgrid_phi, rtol=rtol)
 
-                # Check charged particle moments and f
+                # Check ion particle moments and f
                 ######################################
 
-                newgrid_n_charged = interpolate_to_grid_z(expected.z, n_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.n_charged[:, tind], newgrid_n_charged[:,1], rtol=rtol)
+                newgrid_n_ion = interpolate_to_grid_z(expected.z, n_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.n_ion[:, tind], newgrid_n_ion[:,1], rtol=rtol)
 
-                newgrid_upar_charged = interpolate_to_grid_z(expected.z, upar_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.upar_charged[:, tind], newgrid_upar_charged[:,1], rtol=rtol, atol=atol)
+                newgrid_upar_ion = interpolate_to_grid_z(expected.z, upar_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.upar_ion[:, tind], newgrid_upar_ion[:,1], rtol=rtol, atol=atol)
 
-                newgrid_ppar_charged = interpolate_to_grid_z(expected.z, ppar_charged[:, :, tind], z, z_spectral)
-                @test isapprox(expected.ppar_charged[:, tind], newgrid_ppar_charged[:,1], rtol=rtol)
+                newgrid_ppar_ion = interpolate_to_grid_z(expected.z, ppar_ion[:, :, tind], z, z_spectral)
+                @test isapprox(expected.ppar_ion[:, tind], newgrid_ppar_ion[:,1], rtol=rtol)
 
-                newgrid_vth_charged = @. sqrt(2.0*newgrid_ppar_charged/newgrid_n_charged)
-                newgrid_f_charged = interpolate_to_grid_z(expected.z, f_charged[:, :, :, tind], z, z_spectral)
-                temp = newgrid_f_charged
-                newgrid_f_charged = fill(NaN, length(expected.vpa),
-                                         size(newgrid_f_charged, 2),
-                                         size(newgrid_f_charged, 3),
-                                         size(newgrid_f_charged, 4))
+                newgrid_vth_ion = @. sqrt(2.0*newgrid_ppar_ion/newgrid_n_ion)
+                newgrid_f_ion = interpolate_to_grid_z(expected.z, f_ion[:, :, :, tind], z, z_spectral)
+                temp = newgrid_f_ion
+                newgrid_f_ion = fill(NaN, length(expected.vpa),
+                                         size(newgrid_f_ion, 2),
+                                         size(newgrid_f_ion, 3),
+                                         size(newgrid_f_ion, 4))
                 for iz ∈ 1:length(expected.z)
                     wpa = copy(expected.vpa)
                     if input["evolve_moments_parallel_flow"]
-                        wpa .-= newgrid_upar_charged[iz,1]
+                        wpa .-= newgrid_upar_ion[iz,1]
                     end
                     if input["evolve_moments_parallel_pressure"]
-                        wpa ./= newgrid_vth_charged[iz,1]
+                        wpa ./= newgrid_vth_ion[iz,1]
                     end
-                    newgrid_f_charged[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
+                    newgrid_f_ion[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
                 end
-                @test isapprox(expected.f_charged[:, :, tind], newgrid_f_charged[:,:,1], rtol=rtol)
+                @test isapprox(expected.f_ion[:, :, tind], newgrid_f_ion[:,:,1], rtol=rtol)
 
                 # Check neutral particle moments and f
                 ######################################
diff --git a/test/fokker_planck_time_evolution_tests.jl b/test/fokker_planck_time_evolution_tests.jl
index a022170cb..3d066c67f 100644
--- a/test/fokker_planck_time_evolution_tests.jl
+++ b/test/fokker_planck_time_evolution_tests.jl
@@ -8,7 +8,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_charged_particle_moments_data, load_pdf_data,
+                                 load_ion_particle_moments_data, load_pdf_data,
                                  load_time_data, load_species_data
 using moment_kinetics.type_definitions: mk_float
 
@@ -24,14 +24,14 @@ struct expected_data
     vpa::Array{mk_float, 1}
     vperp::Array{mk_float, 1}
     phi::Array{mk_float, 1} #time
-    n_charged::Array{mk_float, 1} #time
-    upar_charged::Array{mk_float, 1} # time
-    ppar_charged::Array{mk_float, 1} # time
-    pperp_charged::Array{mk_float, 1} # time
-    qpar_charged::Array{mk_float, 1} # time
-    v_t_charged::Array{mk_float, 1} # time
+    n_ion::Array{mk_float, 1} #time
+    upar_ion::Array{mk_float, 1} # time
+    ppar_ion::Array{mk_float, 1} # time
+    pperp_ion::Array{mk_float, 1} # time
+    qpar_ion::Array{mk_float, 1} # time
+    v_t_ion::Array{mk_float, 1} # time
     dSdt::Array{mk_float, 1} # time
-    f_charged::Array{mk_float, 3} # vpa, vperp, time
+    f_ion::Array{mk_float, 3} # vpa, vperp, time
 end
 
 const expected =
@@ -40,21 +40,21 @@ const expected =
    [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
    # Expected phi:
    [-1.267505494648937, -1.275683298550937],
-   # Expected n_charged:
+   # Expected n_ion:
    [0.2815330322340072, 0.2792400986636072],
-   # Expected upar_charged:
+   # Expected upar_ion:
    [0.0, 0.0],
-   # Expected ppar_charged:
+   # Expected ppar_ion:
    [0.17982280248048935, 0.14891126175332367],
-   # Expected pperp_charged
+   # Expected pperp_ion
    [0.14340146667506784, 0.1581377822859991],
-   # Expected qpar_charged
+   # Expected qpar_ion
    [0.0, 0.0],
-   # Expected v_t_charged
+   # Expected v_t_ion
    [1.0511726083010418, 1.0538509291794658],
    # Expected dSdt
    [0.0, 1.1853081348031516e-5],
-   # Expected f_charged:
+   # Expected f_ion:
    [0.0 0.0 0.0 0.0 0.0 0.0 0.0;
     0.0006199600161806666 0.00047805300997075977 0.0002665817112117718 7.637693901737056e-5 1.3272321881722645e-5 1.3988924344690309e-6 0.0;
     0.005882016862626724 0.0045356406743786385 0.002529256854781707 0.0007246442213864763 0.00012592428394890537 1.3272321881722645e-5 0.0;
@@ -87,25 +87,25 @@ const expected =
 ########################################################################################## 
 """
 fid = open_readonly_output_file(path, "dfns")
-f_charged_vpavperpzrst = load_pdf_data(fid)
-f_charged = f_charged_vpavperpzrst[:,:,1,1,1,:]
+f_ion_vpavperpzrst = load_pdf_data(fid)
+f_ion = f_ion_vpavperpzrst[:,:,1,1,1,:]
 ntind = 2
 nvpa = 13  #subject to grid choices
 nvperp = 7 #subject to grid choices
 for k in 1:ntind
   for j in 1:nvperp-1
       for i in 1:nvpa-1
-         @printf("%.15f ", f_charged[i,j,k])
+         @printf("%.15f ", f_ion[i,j,k])
          print("; ")
       end
-      @printf("%.15f ", f_charged[nvpa,j,k])
+      @printf("%.15f ", f_ion[nvpa,j,k])
       print(";;\n")
   end
   for i in 1:nvpa-1
-    @printf("%.15f ", f_charged[i,nvperp,k])
+    @printf("%.15f ", f_ion[i,nvperp,k])
     print("; ")
   end
-  @printf("%.15f ", f_charged[nvpa,nvperp,k])
+  @printf("%.15f ", f_ion[nvpa,nvperp,k])
   if k < ntind
       print(";;;\n")
   end  
@@ -216,14 +216,14 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
     end
 
     phi = nothing
-    n_charged = nothing
-    upar_charged = nothing
-    ppar_charged = nothing
-    pperp_charged = nothing
-    qpar_charged = nothing
-    v_t_charged = nothing
+    n_ion = nothing
+    upar_ion = nothing
+    ppar_ion = nothing
+    pperp_ion = nothing
+    qpar_ion = nothing
+    v_t_ion = nothing
     dSdt = nothing
-    f_charged = nothing
+    f_ion = nothing
     f_err = nothing
     vpa, vpa_spectral = nothing, nothing
     vperp, vperp_spectral = nothing, nothing
@@ -248,8 +248,8 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, 
-            pperp_charged_zrst, qpar_charged_zrst, v_t_charged_zrst, dSdt_zrst = load_charged_particle_moments_data(fid,extended_moments=true)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, 
+            pperp_ion_zrst, qpar_ion_zrst, v_t_ion_zrst, dSdt_zrst = load_ion_particle_moments_data(fid,extended_moments=true)
             
             close(fid)
             
@@ -260,21 +260,21 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             vperp, vperp_spectral = load_coordinate_data(fid, "vperp")
 
             # load particle distribution function (pdf) data
-            f_charged_vpavperpzrst = load_pdf_data(fid)
+            f_ion_vpavperpzrst = load_pdf_data(fid)
             
             close(fid)
             # select the single z, r, s point
             # keep the two time points in the arrays
             phi = phi_zrt[1,1,:]
-            n_charged = n_charged_zrst[1,1,1,:]
-            upar_charged = upar_charged_zrst[1,1,1,:]
-            ppar_charged = ppar_charged_zrst[1,1,1,:]
-            pperp_charged = pperp_charged_zrst[1,1,1,:]
-            qpar_charged = qpar_charged_zrst[1,1,1,:]
-            v_t_charged = v_t_charged_zrst[1,1,1,:]
+            n_ion = n_ion_zrst[1,1,1,:]
+            upar_ion = upar_ion_zrst[1,1,1,:]
+            ppar_ion = ppar_ion_zrst[1,1,1,:]
+            pperp_ion = pperp_ion_zrst[1,1,1,:]
+            qpar_ion = qpar_ion_zrst[1,1,1,:]
+            v_t_ion = v_t_ion_zrst[1,1,1,:]
             dSdt = dSdt_zrst[1,1,1,:]
-            f_charged = f_charged_vpavperpzrst[:,:,1,1,1,:]
-            f_err = copy(f_charged)
+            f_ion = f_ion_vpavperpzrst[:,:,1,1,1,:]
+            f_err = copy(f_ion)
             # Unnormalize f
             # NEED TO UPGRADE TO 2V MOMENT KINETICS HERE
             
@@ -293,20 +293,20 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
                 
                 @test isapprox(expected.phi[tind], phi[tind], rtol=rtol)
 
-                # Check charged particle moments and f
+                # Check ion particle moments and f
                 ######################################
 
-                @test isapprox(expected.n_charged[tind], n_charged[tind], atol=atol)
-                @test isapprox(expected.upar_charged[tind], upar_charged[tind], atol=atol)
-                @test isapprox(expected.ppar_charged[tind], ppar_charged[tind], atol=atol)
-                @test isapprox(expected.pperp_charged[tind], pperp_charged[tind], atol=atol)
-                @test isapprox(expected.qpar_charged[tind], qpar_charged[tind], atol=atol)
-                @test isapprox(expected.v_t_charged[tind], v_t_charged[tind], atol=atol)
+                @test isapprox(expected.n_ion[tind], n_ion[tind], atol=atol)
+                @test isapprox(expected.upar_ion[tind], upar_ion[tind], atol=atol)
+                @test isapprox(expected.ppar_ion[tind], ppar_ion[tind], atol=atol)
+                @test isapprox(expected.pperp_ion[tind], pperp_ion[tind], atol=atol)
+                @test isapprox(expected.qpar_ion[tind], qpar_ion[tind], atol=atol)
+                @test isapprox(expected.v_t_ion[tind], v_t_ion[tind], atol=atol)
                 @test isapprox(expected.dSdt[tind], dSdt[tind], atol=atol)
-                @. f_err = abs(expected.f_charged - f_charged)
+                @. f_err = abs(expected.f_ion - f_ion)
                 max_f_err = maximum(f_err)
                 @test isapprox(max_f_err, 0.0, atol=atol)
-                @test isapprox(expected.f_charged[:,:,tind], f_charged[:,:,tind], atol=atol)
+                @test isapprox(expected.f_ion[:,:,tind], f_ion[:,:,tind], atol=atol)
             end
         end
 
diff --git a/test/nonlinear_sound_wave_inputs_and_expected_data.jl b/test/nonlinear_sound_wave_inputs_and_expected_data.jl
index f3f7948da..be3201666 100644
--- a/test/nonlinear_sound_wave_inputs_and_expected_data.jl
+++ b/test/nonlinear_sound_wave_inputs_and_expected_data.jl
@@ -7,13 +7,13 @@ struct expected_data
     z::Array{mk_float, 1}
     vpa::Array{mk_float, 1}
     phi::Array{mk_float, 2}
-    n_charged::Array{mk_float, 2}
+    n_ion::Array{mk_float, 2}
     n_neutral::Array{mk_float, 2}
-    upar_charged::Array{mk_float, 2}
+    upar_ion::Array{mk_float, 2}
     upar_neutral::Array{mk_float, 2}
-    ppar_charged::Array{mk_float, 2}
+    ppar_ion::Array{mk_float, 2}
     ppar_neutral::Array{mk_float, 2}
-    f_charged::Array{mk_float, 3}
+    f_ion::Array{mk_float, 3}
     f_neutral::Array{mk_float, 3}
 end
 
@@ -32,48 +32,48 @@ const expected =
         -0.353459763648872 -0.3755265897483555; -0.5494724768120175 -0.5903597644911376;
         -0.8609860698164502 -0.8726370464896476; -1.2115129555832849 -1.1306355658313922;
         -1.3862820803244258 -1.2382641646968997],
-   n_charged=[0.2500030702177186 0.2898869775083742; 0.2977473631375158 0.3228278662412625;
-              0.42274585818529853 0.417848119539277; 0.5772542465450629 0.5541281150892785;
-              0.7022542481909738 0.6869277664245242; 0.7499999999999394 0.7466605958319346;
-              0.7022542481909738 0.6869277664245237; 0.577254246545063 0.5541281150892783;
-              0.42274585818529864 0.41784811953927686; 0.2977473631375159 0.32282786624126253;
-              0.2500030702177185 0.2898869775083743],
+   n_ion=[0.2500030702177186 0.2898869775083742; 0.2977473631375158 0.3228278662412625;
+          0.42274585818529853 0.417848119539277; 0.5772542465450629 0.5541281150892785;
+          0.7022542481909738 0.6869277664245242; 0.7499999999999394 0.7466605958319346;
+          0.7022542481909738 0.6869277664245237; 0.577254246545063 0.5541281150892783;
+          0.42274585818529864 0.41784811953927686; 0.2977473631375159 0.32282786624126253;
+          0.2500030702177185 0.2898869775083743],
    n_neutral=[0.7499999999999382 0.7736769553678673; 0.7022542481909748 0.7056866352169496;
               0.5772542465450632 0.5582977481633454; 0.4227458581852985 0.40969188756651037;
               0.29774736313751604 0.30539644783353687; 0.2500030702177186 0.268198658560817;
               0.29774736313751604 0.305396447833537; 0.42274585818529836 0.4096918875665103;
               0.5772542465450631 0.5582977481633457; 0.7022542481909745 0.7056866352169494;
               0.7499999999999383 0.7736769553678673],
-   upar_charged=[-2.7135787559953277e-17 -6.299214622140781e-17; -9.321028970172899e-18 -0.1823721921091055;
-                 -2.8374879811351724e-18 -0.19657035490893093; 1.2124327390522635e-17 -0.11139486685283827;
-                 3.6525788403693063e-17 -0.033691837771623996; -2.0930856430671915e-17 4.84147091991613e-17;
-                 8.753545920086251e-18 0.033691837771624024; 1.1293771270243255e-17 0.11139486685283813;
-                 1.3739171132886587e-17 0.19657035490893102; -6.840453743089351e-18 0.18237219210910513;
-                 -2.7135787559953277e-17 -4.656897959900552e-17],
+   upar_ion=[-2.7135787559953277e-17 -6.299214622140781e-17; -9.321028970172899e-18 -0.1823721921091055;
+             -2.8374879811351724e-18 -0.19657035490893093; 1.2124327390522635e-17 -0.11139486685283827;
+             3.6525788403693063e-17 -0.033691837771623996; -2.0930856430671915e-17 4.84147091991613e-17;
+             8.753545920086251e-18 0.033691837771624024; 1.1293771270243255e-17 0.11139486685283813;
+             1.3739171132886587e-17 0.19657035490893102; -6.840453743089351e-18 0.18237219210910513;
+             -2.7135787559953277e-17 -4.656897959900552e-17],
    upar_neutral=[6.5569385065066925e-18 7.469475342067322e-17; 1.1054500872839027e-17 -0.036209130454625794;
                  -3.241833393685864e-17 -0.00915544640981337; -3.617637280460899e-17 0.05452268209340691;
                  4.417578961284041e-17 0.07606644718003618; 4.9354467746194965e-17 4.452343983947504e-17;
                  6.573091229872379e-18 -0.07606644718003616; 2.989662686945165e-17 -0.05452268209340687;
                  -3.1951996361666834e-17 0.009155446409813412; -4.395464518158184e-18 0.03620913045462582;
                  6.5569385065066925e-18 7.150569974151354e-17],
-   ppar_charged=[0.18749999999999992 0.2328164829490338; 0.20909325514551116 0.21912575009260987;
-                 0.24403180771238264 0.20822611102296495; 0.24403180771238278 0.21506741942934832;
-                 0.2090932551455113 0.22097085045011763; 0.1875 0.22119050467096843;
-                 0.20909325514551128 0.2209708504501176; 0.2440318077123828 0.2150674194293483;
-                 0.24403180771238256 0.20822611102296476; 0.20909325514551116 0.21912575009260982;
-                 0.18749999999999992 0.2328164829490338],
+   ppar_ion=[0.18749999999999992 0.2328164829490338; 0.20909325514551116 0.21912575009260987;
+             0.24403180771238264 0.20822611102296495; 0.24403180771238278 0.21506741942934832;
+             0.2090932551455113 0.22097085045011763; 0.1875 0.22119050467096843;
+             0.20909325514551128 0.2209708504501176; 0.2440318077123828 0.2150674194293483;
+             0.24403180771238256 0.20822611102296476; 0.20909325514551116 0.21912575009260982;
+             0.18749999999999992 0.2328164829490338],
    ppar_neutral=[0.18750000000000003 0.2480244331470989; 0.20909325514551122 0.2440075646485762;
                  0.24403180771238286 0.22861256884534023; 0.24403180771238278 0.20588932618946498;
                  0.20909325514551144 0.19263633346696638; 0.18749999999999992 0.19091848744561835;
                  0.20909325514551141 0.19263633346696654; 0.2440318077123828 0.20588932618946482;
                  0.24403180771238286 0.22861256884534029; 0.20909325514551114 0.24400756464857642;
                  0.18750000000000006 0.24802443314709893],
-   f_charged=[0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;
-              0.20411991941198782 0.25123395823993105 0.3934413727192304 0.6277900619432855 0.9100364506661008 1.0606601717796504 0.910036450666101 0.6277900619432859 0.39344137271923046 0.25123395823993094 0.20411991941198776;
-              0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;;;
-              0.05392403019146985 0.06057819609646438 0.03676744157455075 0.013740507879552622 0.010777319583092297 0.019330359159894384 0.027982173790396116 0.027603104735767332 0.02667986700464528 0.035654512254837005 0.05392403019146984;
-              0.21177720235387912 0.24902901234066305 0.3729377138332225 0.596281539172339 0.8870867512643452 1.0533860567375264 0.887086751264345 0.5962815391723388 0.3729377138332225 0.24902901234066285 0.21177720235387912;
-              0.053924030191469796 0.035654512254837074 0.02667986700464531 0.02760310473576733 0.02798217379039615 0.019330359159894287 0.010777319583092311 0.013740507879552624 0.03676744157455069 0.060578196096464365 0.05392403019146979],
+   f_ion=[0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;
+          0.20411991941198782 0.25123395823993105 0.3934413727192304 0.6277900619432855 0.9100364506661008 1.0606601717796504 0.910036450666101 0.6277900619432859 0.39344137271923046 0.25123395823993094 0.20411991941198776;
+          0.0370462360994826 0.04059927063892091 0.0428431419871786 0.030398267195914062 0.01236045902698859 0.006338529470383425 0.012360459026988587 0.030398267195914028 0.04284314198717859 0.0405992706389209 0.0370462360994826;;;
+          0.05392403019146985 0.06057819609646438 0.03676744157455075 0.013740507879552622 0.010777319583092297 0.019330359159894384 0.027982173790396116 0.027603104735767332 0.02667986700464528 0.035654512254837005 0.05392403019146984;
+          0.21177720235387912 0.24902901234066305 0.3729377138332225 0.596281539172339 0.8870867512643452 1.0533860567375264 0.887086751264345 0.5962815391723388 0.3729377138332225 0.24902901234066285 0.21177720235387912;
+          0.053924030191469796 0.035654512254837074 0.02667986700464531 0.02760310473576733 0.02798217379039615 0.019330359159894287 0.010777319583092311 0.013740507879552624 0.03676744157455069 0.060578196096464365 0.05392403019146979],
    f_neutral=[0.0063385294703834595 0.012360459026988546 0.030398267195914108 0.04284314198717859 0.040599270638920985 0.03704623609948259 0.040599270638920965 0.0428431419871786 0.030398267195914094 0.012360459026988546 0.006338529470383456;
               1.0606601717796493 0.9100364506661016 0.6277900619432857 0.3934413727192303 0.2512339582399308 0.20411991941198754 0.2512339582399307 0.3934413727192301 0.6277900619432853 0.9100364506661016 1.0606601717796487;
               0.0063385294703834595 0.012360459026988546 0.030398267195914108 0.04284314198717859 0.040599270638920985 0.03704623609948259 0.040599270638920965 0.0428431419871786 0.030398267195914094 0.012360459026988546 0.006338529470383456;;;
diff --git a/test/nonlinear_sound_wave_tests.jl b/test/nonlinear_sound_wave_tests.jl
index 691c70b80..71f32dd1d 100644
--- a/test/nonlinear_sound_wave_tests.jl
+++ b/test/nonlinear_sound_wave_tests.jl
@@ -8,7 +8,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_charged_particle_moments_data, load_pdf_data,
+                                 load_ion_particle_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -58,10 +58,10 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
     end
 
     phi = nothing
-    n_charged = nothing
-    upar_charged = nothing
-    ppar_charged = nothing
-    f_charged = nothing
+    n_ion = nothing
+    upar_ion = nothing
+    ppar_ion = nothing
+    f_ion = nothing
     n_neutral = nothing
     upar_neutral = nothing
     ppar_neutral = nothing
@@ -89,7 +89,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_charged_zrst, upar_charged_zrst, ppar_charged_zrst, qpar_charged_zrst, v_t_charged_zrst = load_charged_particle_moments_data(fid)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_particle_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
             z, z_spectral = load_coordinate_data(fid, "z")
 
@@ -106,12 +106,12 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             close(fid)
             
             phi = phi_zrt[:,1,:]
-            n_charged = n_charged_zrst[:,1,:,:]
-            upar_charged = upar_charged_zrst[:,1,:,:]
-            ppar_charged = ppar_charged_zrst[:,1,:,:]
-            qpar_charged = qpar_charged_zrst[:,1,:,:]
-            v_t_charged = v_t_charged_zrst[:,1,:,:]
-            f_charged = f_charged_vpavperpzrst[:,1,:,1,:,:]
+            n_ion = n_ion_zrst[:,1,:,:]
+            upar_ion = upar_ion_zrst[:,1,:,:]
+            ppar_ion = ppar_ion_zrst[:,1,:,:]
+            qpar_ion = qpar_ion_zrst[:,1,:,:]
+            v_t_ion = v_t_ion_zrst[:,1,:,:]
+            f_ion = f_ion_vpavperpzrst[:,1,:,1,:,:]
             n_neutral = n_neutral_zrst[:,1,:,:]
             upar_neutral = upar_neutral_zrst[:,1,:,:]
             ppar_neutral = ppar_neutral_zrst[:,1,:,:]
@@ -122,7 +122,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # Unnormalize f
             if input["evolve_moments_density"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] .*= n_charged[iz,is,it]
+                    f_ion[:,iz,is,it] .*= n_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] .*= n_neutral[iz,isn,it]
@@ -130,7 +130,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             end
             if input["evolve_moments_parallel_pressure"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] ./= v_t_charged[iz,is,it]
+                    f_ion[:,iz,is,it] ./= v_t_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] ./= v_t_neutral[iz,isn,it]
@@ -228,7 +228,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
                     if input["evolve_moments_parallel_pressure"]
                         wpa ./= newgrid_vth_ion[iz,1]
                     end
-                    newgrid_f_charged[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
+                    newgrid_f_ion[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
                 end
                 @test isapprox(expected.f_ion[:, :, tind], newgrid_f_ion[:,:,1], rtol=rtol)
 
diff --git a/test/restart_interpolation_tests.jl b/test/restart_interpolation_tests.jl
index 251e6d4a9..507642aeb 100644
--- a/test/restart_interpolation_tests.jl
+++ b/test/restart_interpolation_tests.jl
@@ -12,7 +12,7 @@ using moment_kinetics.file_io: io_has_parallel
 using moment_kinetics.input_structs: grid_input, advection_input, hdf5
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_charged_particle_moments_data, load_pdf_data,
+                                 load_ion_particle_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -115,10 +115,10 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
     end
 
     phi = nothing
-    n_charged = nothing
-    upar_charged = nothing
-    ppar_charged = nothing
-    f_charged = nothing
+    n_ion = nothing
+    upar_ion = nothing
+    ppar_ion = nothing
+    f_ion = nothing
     n_neutral = nothing
     upar_neutral = nothing
     ppar_neutral = nothing
@@ -150,12 +150,12 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
             time = run_info.time
             n_ion_species = run_info.n_ion_species
             n_neutral_species = run_info.n_neutral_species
-            n_charged_zrst = postproc_load_variable(run_info, "density")
-            upar_charged_zrst = postproc_load_variable(run_info, "parallel_flow")
-            ppar_charged_zrst = postproc_load_variable(run_info, "parallel_pressure")
-            qpar_charged_zrst = postproc_load_variable(run_info, "parallel_heat_flux")
-            v_t_charged_zrst = postproc_load_variable(run_info, "thermal_speed")
-            f_charged_vpavperpzrst  = postproc_load_variable(run_info, "f")
+            n_ion_zrst = postproc_load_variable(run_info, "density")
+            upar_ion_zrst = postproc_load_variable(run_info, "parallel_flow")
+            ppar_ion_zrst = postproc_load_variable(run_info, "parallel_pressure")
+            qpar_ion_zrst = postproc_load_variable(run_info, "parallel_heat_flux")
+            v_t_ion_zrst = postproc_load_variable(run_info, "thermal_speed")
+            f_ion_vpavperpzrst  = postproc_load_variable(run_info, "f")
             n_neutral_zrst = postproc_load_variable(run_info, "density_neutral")
             upar_neutral_zrst = postproc_load_variable(run_info, "uz_neutral")
             ppar_neutral_zrst = postproc_load_variable(run_info, "pz_neutral")
@@ -176,12 +176,12 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
             rm(joinpath(realpath(input["base_directory"]), name); recursive=true)
 
             phi = phi_zrt[:,1,:]
-            n_charged = n_charged_zrst[:,1,:,:]
-            upar_charged = upar_charged_zrst[:,1,:,:]
-            ppar_charged = ppar_charged_zrst[:,1,:,:]
-            qpar_charged = qpar_charged_zrst[:,1,:,:]
-            v_t_charged = v_t_charged_zrst[:,1,:,:]
-            f_charged = f_charged_vpavperpzrst[:,1,:,1,:,:]
+            n_ion = n_ion_zrst[:,1,:,:]
+            upar_ion = upar_ion_zrst[:,1,:,:]
+            ppar_ion = ppar_ion_zrst[:,1,:,:]
+            qpar_ion = qpar_ion_zrst[:,1,:,:]
+            v_t_ion = v_t_ion_zrst[:,1,:,:]
+            f_ion = f_ion_vpavperpzrst[:,1,:,1,:,:]
             n_neutral = n_neutral_zrst[:,1,:,:]
             upar_neutral = upar_neutral_zrst[:,1,:,:]
             ppar_neutral = ppar_neutral_zrst[:,1,:,:]
@@ -192,7 +192,7 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
             # Unnormalize f
             if input["evolve_moments_density"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] .*= n_charged[iz,is,it]
+                    f_ion[:,iz,is,it] .*= n_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] .*= n_neutral[iz,isn,it]
@@ -200,7 +200,7 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
             end
             if input["evolve_moments_parallel_pressure"]
                 for it ∈ 1:length(time), is ∈ 1:n_ion_species, iz ∈ 1:z.n
-                    f_charged[:,iz,is,it] ./= v_t_charged[iz,is,it]
+                    f_ion[:,iz,is,it] ./= v_t_ion[iz,is,it]
                 end
                 for it ∈ 1:length(time), isn ∈ 1:n_neutral_species, iz ∈ 1:z.n
                     f_neutral[:,iz,isn,it] ./= v_t_neutral[iz,isn,it]
@@ -211,36 +211,36 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
         newgrid_phi = interpolate_to_grid_z(expected.z, phi[:, end], z, z_spectral)
         @test isapprox(expected.phi[:, end], newgrid_phi, rtol=rtol)
 
-        # Check charged particle moments and f
+        # Check ion particle moments and f
         ######################################
 
-        newgrid_n_charged = interpolate_to_grid_z(expected.z, n_charged[:, :, end], z, z_spectral)
-        @test isapprox(expected.n_charged[:, end], newgrid_n_charged[:,1], rtol=rtol)
+        newgrid_n_ion = interpolate_to_grid_z(expected.z, n_ion[:, :, end], z, z_spectral)
+        @test isapprox(expected.n_ion[:, end], newgrid_n_ion[:,1], rtol=rtol)
 
-        newgrid_upar_charged = interpolate_to_grid_z(expected.z, upar_charged[:, :, end], z, z_spectral)
-        @test isapprox(expected.upar_charged[:, end], newgrid_upar_charged[:,1], rtol=rtol, atol=atol_3V)
+        newgrid_upar_ion = interpolate_to_grid_z(expected.z, upar_ion[:, :, end], z, z_spectral)
+        @test isapprox(expected.upar_ion[:, end], newgrid_upar_ion[:,1], rtol=rtol, atol=atol_3V)
 
-        newgrid_ppar_charged = interpolate_to_grid_z(expected.z, ppar_charged[:, :, end], z, z_spectral)
-        @test isapprox(expected.ppar_charged[:, end], newgrid_ppar_charged[:,1], rtol=rtol)
+        newgrid_ppar_ion = interpolate_to_grid_z(expected.z, ppar_ion[:, :, end], z, z_spectral)
+        @test isapprox(expected.ppar_ion[:, end], newgrid_ppar_ion[:,1], rtol=rtol)
 
-        newgrid_vth_charged = @. sqrt(2.0*newgrid_ppar_charged/newgrid_n_charged)
-        newgrid_f_charged = interpolate_to_grid_z(expected.z, f_charged[:, :, :, end], z, z_spectral)
-        temp = newgrid_f_charged
-        newgrid_f_charged = fill(NaN, length(expected.vpa),
-                                 size(newgrid_f_charged, 2),
-                                 size(newgrid_f_charged, 3),
-                                 size(newgrid_f_charged, 4))
+        newgrid_vth_ion = @. sqrt(2.0*newgrid_ppar_ion/newgrid_n_ion)
+        newgrid_f_ion = interpolate_to_grid_z(expected.z, f_ion[:, :, :, end], z, z_spectral)
+        temp = newgrid_f_ion
+        newgrid_f_ion = fill(NaN, length(expected.vpa),
+                                 size(newgrid_f_ion, 2),
+                                 size(newgrid_f_ion, 3),
+                                 size(newgrid_f_ion, 4))
         for iz ∈ 1:length(expected.z)
             wpa = copy(expected.vpa)
             if input["evolve_moments_parallel_flow"]
-                wpa .-= newgrid_upar_charged[iz,1]
+                wpa .-= newgrid_upar_ion[iz,1]
             end
             if input["evolve_moments_parallel_pressure"]
-                wpa ./= newgrid_vth_charged[iz,1]
+                wpa ./= newgrid_vth_ion[iz,1]
             end
-            newgrid_f_charged[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
+            newgrid_f_ion[:,iz,1] = interpolate_to_grid_vpa(wpa, temp[:,iz,1], vpa, vpa_spectral)
         end
-        @test isapprox(expected.f_charged[:, :, end], newgrid_f_charged[:,:,1], rtol=rtol_3V)
+        @test isapprox(expected.f_ion[:, :, end], newgrid_f_ion[:,:,1], rtol=rtol_3V)
 
         # Check neutral particle moments and f
         ######################################

From 42aa0f1a35ea35c1cc0659be875c7c3b37b49631 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:00:45 +0000
Subject: [PATCH 057/428] Use Deuterium instead of Hydrogen to calculate
 electron mass ratio

This increases the mass ratio by a factor of 2, so makes the neglect of
mass-ratio small terms a bit better.
---
 src/moment_kinetics_input.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 9bf53079d..615e02694 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -911,7 +911,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0
     # ratio of the electron particle mass to the ion particle mass
-    me_over_mi = 1.0/1836.0
+    me_over_mi = 1.0/1836.0/2.0
     # The ion flux reaching the wall that is recycled as neutrals is reduced by
     # `recycling_fraction` to account for ions absorbed by the wall.
     recycling_fraction = 1.0

From 8ae43975db9ab5e1269bbfc4a88d0ec6f50d377f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:18:04 +0000
Subject: [PATCH 058/428] Calculate error from residual before hack to speed up
 convergence

Want the stopping criterion for the iteration to use the actual dg/dt,
not the hacked version.
---
 src/electron_kinetic_equation.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index e30d6ff65..d1d8d2e7e 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -325,14 +325,15 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                             dppar_dz, dqpar_dz, dvth_dz, 
                                             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                             num_diss_params, dt_max)
+        # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
+        #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
+
         # Divide by wpa to relax CFL condition at large wpa - only looking for steady
         # state here, so does not matter that this makes time evolution incorrect.
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
             residual[ivpa,ivperp,iz,ir] /= sqrt(1.0 + vpa.grid[ivpa]^2)
         end
-        # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
-        #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
         if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
             break
         end

From 3385943a9b46eba6d6f82221b0522437cd4d2229 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:03:01 +0000
Subject: [PATCH 059/428] DEBUGGING! Return more stuff from electron solve

---
 src/electron_kinetic_equation.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index d1d8d2e7e..adb32e85a 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -205,10 +205,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     result_pdf[:,:,1] .= pdf[:,1,:,1]
     result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
     result_ppar[:,1] .= ppar[:,1]
+    result_vth = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+    result_vth[:,1] .= vthe[:,1]
     result_qpar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
     result_qpar[:,1] .= qpar[:,1]
     result_phi = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
     result_phi[:,1] .= phi[:,1]
+    result_residual = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
         @. pdf -= dt_electron * residual
@@ -278,6 +281,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             println(io_vth,"")
             result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
             result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
+            result_vth[:,iteration÷output_interval+1] .= vthe[:,1]
             result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
             result_phi[:,iteration÷output_interval+1] .= phi[:,1]
         end
@@ -328,6 +332,9 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
+        if (mod(iteration,output_interval) == 0)
+            result_residual[:,:,iteration÷output_interval+1] .= residual[:,1,:,1]
+        end
 
         # Divide by wpa to relax CFL condition at large wpa - only looking for steady
         # state here, so does not matter that this makes time evolution incorrect.
@@ -355,7 +362,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     close(io_vth)
     close(io_pdf)
     close(io_pdf_stages)
-    return result_pdf, result_ppar, result_qpar, result_phi, z, vpa
+    return result_pdf, dens, moments.electron.upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)

From dc62cf5e879bcdca2bea2a955f7f804334e10e1a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:03:34 +0000
Subject: [PATCH 060/428] Don't zero out negative values of g_e

Zeroing these out does not seem to help, although it doesn't seem to
hurt either.
---
 src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index adb32e85a..869be118f 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -220,7 +220,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
-        pdf = max.(pdf, 0.0)
+        #pdf = max.(pdf, 0.0)
         for ir ∈ 1:size(ppar, 2), iz ∈ 2:size(ppar,1)-1
             #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
             @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)

From 271d0d557074a2f3369c5da0015ee3f5e8008271 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:24:09 +0000
Subject: [PATCH 061/428] Extra hack to speed up electron evolution away from
 walls

Increase the residual by some factor (up to 21x for the moment) in the
centre of the z-domain. Timestep seems to be limited by positions near
the sheath entrance, so this effectively increases the timestep away
from there, aiming to speed up the global convergence.
---
 src/electron_kinetic_equation.jl | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 869be118f..68f1cf494 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -338,8 +338,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         # Divide by wpa to relax CFL condition at large wpa - only looking for steady
         # state here, so does not matter that this makes time evolution incorrect.
+        Lz = z.L
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-            residual[ivpa,ivperp,iz,ir] /= sqrt(1.0 + vpa.grid[ivpa]^2)
+            zval = z.grid[iz]
+            znorm = 2.0*zval/Lz
+            residual[ivpa,ivperp,iz,ir] *=
+                (1.0 + 20.0*(1.0 - znorm^2)) /
+                sqrt(1.0 + vpa.grid[ivpa]^2)
         end
         if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
             break

From 3b50e4484c73bf55b166a5fec25e9e14ee0805e1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:38:14 +0000
Subject: [PATCH 062/428] DEBUGGING! Fix return-electron-arrays debug hack

One return statement went missing in merge of master
---
 src/moment_kinetics.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index b765ce8ed..abc0a5ae4 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -338,7 +338,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
-        init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry,
+        return init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry,
                               composition, r, z, vperp, vpa, vzeta, vr, vz,
                               z_spectral, r_spectral, vpa_spectral, vz_spectral, species,
                               collisions, external_source_settings,

From 410301519b8de1ae194964d56339926319dc2aac Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 4 Feb 2024 21:39:28 +0000
Subject: [PATCH 063/428] Cut off g_e at sheath with tanh() instead of sharp
 cut-off

One significant advantage of this update is that the exact position of
the cut-off is allowed to be in between grid points, so the position
does not jump from grid point to grid point (it is interpolated linearly
in between). The jumps from grid point to grid point prevent (or at
least significantly slow down) convergence to a steady state at long-ish
times.

Not clear if there is a large advantage to the tanh-cutoff approach, as
there is still a lot of ringing in the tails of the distribution
function, but if this change is reverted, the replacement should be a
cutoff that rather than just setting points to zero, scales the last
point so that the integral over the tail gives u_i exactly (to avoid the
jumping-from-grid-point-to-grid-point thing).
---
 src/electron_kinetic_equation.jl | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 68f1cf494..40c47cb7b 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -383,6 +383,8 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
     # corrections to the pdf needed to ensure moment constraints are satisfied
     pdf_adjustment_option = "vpa4_gaussian"
 
+    cutoff_step_width = 0.1
+
     # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
     #wpa_values = vpa.scratch
     # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
@@ -436,16 +438,23 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         #end
         upar0 = upar[1,ir]
         #println("before pdf left ", pdf[:,1,1,ir])
-        while upar_integral > upar0
+        while upar_integral > upar0 && ivpa_max > 1
             ivpa += 1
             ivpa_max -= 1
             # zero out the reversed pdf at the current cutoff velocity
-            reversed_pdf[ivpa_max] = 0.0
+            #reversed_pdf[ivpa_max] = 0.0
             # calculate the updated first moment of the normalised pdf
             upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
             #println("left ", ivpa, " ", upar_integral, " ", upar0)
         end
-        vmax = vpa_unnorm[ivpa_max]
+        integral_excess = upar_integral - upar0
+        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
+        vmax = 0.5*(vpa_unnorm[ivpa+1] + vpa_unnorm[ivpa]) +
+               fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa+1])
+        wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
+        @loop_vpa ivpa begin
+            reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
+        end
         #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
         #println("done first cutoff loop")
         # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
@@ -635,16 +644,23 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         #end
         upar_end = upar[end,ir]
         #println("before pdf ", pdf[:,1,end,ir])
-        while upar_integral < upar_end
+        while upar_integral < upar_end && ivpa > 1
             ivpa -= 1
             ivpa_min += 1
             # zero out the reversed pdf at the current cutoff velocity
-            reversed_pdf[ivpa_min] = 0.0
+            #reversed_pdf[ivpa_min] = 0.0
             # calculate the updated first moment of the normalised pdf
             upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
             #println("right ", ivpa, " ", upar_integral, " ", upar_end)
         end
-        vmin = vpa_unnorm[ivpa_min]
+        integral_excess = upar_integral - upar_end
+        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
+        vmin = 0.5*(vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
+               fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa-1])
+        wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
+        @loop_vpa ivpa begin
+            reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
+        end
         #println("done second cutoff loop")
 
         # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)

From c8713de1bc884f61d62b67df5c0cd65df0c27904 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:14:54 +0000
Subject: [PATCH 064/428] Undo removal of wpar^1 moment constraint enforcement
 at boundaries

A previous commit (e63b1c37d0fae101adfecde52600f4a8749af01ao0) removed
enforcement of the `\int d(wpa) wpa g_e = 0` constraint at the sheath
entrance boundaries. Applying this constraint does not seem to be
detrimental to numerical stability, so restore it.
---
 src/electron_kinetic_equation.jl | 44 ++++++++++++++++----------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 40c47cb7b..127da5fb7 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -547,18 +547,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             @. vpa.scratch3 *= vpa.grid
             vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
             # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-            #     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            #normalisation_constant_A = (1 + normalisation_constant_B 
-            #    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            #normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
             normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                                       (vpa2_wpa2_moment
-                                        - wpa2_moment / zeroth_moment * vpa2_moment)
-            normalisation_constant_A = (1 - normalisation_constant_B 
-                                            * vpa2_moment) / zeroth_moment
-            normalisation_constant_C = 0.0
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #                           (vpa2_wpa2_moment
+            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+            #normalisation_constant_A = (1 - normalisation_constant_B 
+            #                                * vpa2_moment) / zeroth_moment
+            #normalisation_constant_C = 0.0
             @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         else
@@ -763,18 +763,18 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             @. vpa.scratch3 *= vpa.grid
             vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
             # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-            #     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            #normalisation_constant_A = (1 + normalisation_constant_B 
-            #    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            #normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
             normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                                       (vpa2_wpa2_moment
-                                        - wpa2_moment / zeroth_moment * vpa2_moment)
-            normalisation_constant_A = (1 - normalisation_constant_B 
-                                            * vpa2_moment) / zeroth_moment
-            normalisation_constant_C = 0.0
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #                           (vpa2_wpa2_moment
+            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+            #normalisation_constant_A = (1 - normalisation_constant_B 
+            #                                * vpa2_moment) / zeroth_moment
+            #normalisation_constant_C = 0.0
             @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         else

From a3a54dfde8624e293c1374c35262b5ea0156c399 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:17:57 +0000
Subject: [PATCH 065/428] DEBUGGING! More commented-out debug prints

---
 src/electron_kinetic_equation.jl | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 127da5fb7..2ab133e81 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -144,6 +144,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
     #dt_electron = dt * sqrt(composition.me_over_mi)
     dt_max = 1.0e-8 #1.0
+    #dt_max = 2.5e-9 #1.0
     dt_energy = 0.001
     time = 0.0
 
@@ -449,9 +450,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         end
         integral_excess = upar_integral - upar0
         fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
+        #println("fraction_of_pdf=", fraction_of_pdf)
         vmax = 0.5*(vpa_unnorm[ivpa+1] + vpa_unnorm[ivpa]) +
                fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa+1])
+        #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
         wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
+        #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
         @loop_vpa ivpa begin
             reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
         end
@@ -462,6 +466,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         phi[1,ir] = me_over_mi * vmax^2
         iv0 = findfirst(x -> x>0.0, vpa_unnorm)
         pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
+        #println("check reversed change ", reversed_pdf[iv0:end])
         #println("reversed_pdf ", reversed_pdf)
         #println("after pdf left ", pdf[:,1,1,ir])
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments
@@ -655,8 +660,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         end
         integral_excess = upar_integral - upar_end
         fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
+        #println("B fraction_of_pdf=", fraction_of_pdf)
         vmin = 0.5*(vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
                fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa-1])
+        #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
         wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
         @loop_vpa ivpa begin
             reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))

From c8f46ee0c15c04f57c8788a07976579f1a54f84d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:21:27 +0000
Subject: [PATCH 066/428] DEBUGGING! Increase output_interval

Now that electron time-advance runs for quite a while, can take output
every 1000 timesteps instead of every 50, to avoid having too much
output to look at.
---
 src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 2ab133e81..7d9abcb07 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -201,7 +201,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
     #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
     # evolve (artificially) in time until the residual is less than the tolerance
-    output_interval = 50
+    output_interval = 1000
     result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
     result_pdf[:,:,1] .= pdf[:,1,:,1]
     result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)

From 32e2a94682182916092ee98ac7947118e1a5be25 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:23:54 +0000
Subject: [PATCH 067/428] Actually increase timestep for ppar subcycling

This change was supposed to be introduced in "Fix electron ppar
evolution", but in that commit the timestep was being reset inside the
subcycling loop, so the change had no effect. With the change actually
implemented, can only increase the timestep by 10x not 100x.
---
 src/electron_kinetic_equation.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 7d9abcb07..1ebd5b94f 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -287,7 +287,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             result_phi[:,iteration÷output_interval+1] .= phi[:,1]
         end
 
-        dt_energy = dt_electron * 100.0
+        dt_energy = dt_electron * 10.0
 
         # get an updated iterate of the electron parallel pressure
         #ppar .= ppar_old
@@ -306,7 +306,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
             end
 
-            dt_energy = dt_electron
+            #dt_energy = dt_electron
             electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z.grid)
             fvec.electron_ppar .= ppar
         

From f114130eb46c56c7240b78f6ccf81d29c255ee97 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:25:33 +0000
Subject: [PATCH 068/428] Reduce number of subcycles for ppar evolution

The ppar subcycle was taking most of the execution time (on 1 core...).
Not sure if taking that many subcycles was speeding up convergence of
the overall loop, so reduce from 1000 subcycles to 100 to make each
outer-loop iteration take less time.
---
 src/electron_kinetic_equation.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 1ebd5b94f..0d01d7653 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -292,7 +292,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # get an updated iterate of the electron parallel pressure
         #ppar .= ppar_old
         wpa3_moment = @. qpar / vthe^3
-        for i in 1:1000
+        #for i in 1:1000
+        for i in 1:100
             @. qpar = vthe^3 * wpa3_moment
             # Compute the upwinded z-derivative of the electron parallel pressure for the 
             # electron energy equation

From be521d2a3567f5d3a84a577b1904b9cba82c237a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:47:07 +0000
Subject: [PATCH 069/428] DEBUGGING! Output pdf to ascii every
 `output_interval` steps

...instead of hard-coded 50.
---
 src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 0d01d7653..fc0441986 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -230,7 +230,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         #println("B pdf end ", pdf[:,1,end,1])
         #error("foo")
         
-        if (mod(iteration,50)==1)
+        if (mod(iteration,output_interval)==1)
             @loop_vpa ivpa begin
                 println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " iteration: ", iteration, " flag: ", 1)
             end

From 1fb766d4442c5e317714bf5c6c722852d1a7238a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 09:54:36 +0000
Subject: [PATCH 070/428] Skip points set by boundary condition when checking
 convergence

The 'residual' at points that are set by the sheath boundary condition
is not used to update the distribution function, so it might not
converge to zero. Therefore skip these points when calculating the
convergence criterion from the 'residual'.
---
 src/electron_kinetic_equation.jl | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index fc0441986..c9ae827af 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -198,7 +198,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
     # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
     #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
+    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf), moments.electron.upar, vthe, vpa)
     #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
     # evolve (artificially) in time until the residual is less than the tolerance
     output_interval = 1000
@@ -333,7 +333,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                             num_diss_params, dt_max)
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf))
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf), moments.electron.upar, vthe, vpa)
         if (mod(iteration,output_interval) == 0)
             result_residual[:,:,iteration÷output_interval+1] .= residual[:,1,:,1]
         end
@@ -1441,7 +1441,7 @@ function get_electron_critical_velocities(phi, vthe, me_over_mi)
     return crit_speed_zmin, crit_speed_zmax
 end
 
-function check_electron_pdf_convergence(residual, norm)
+function check_electron_pdf_convergence(residual, norm, upar, vthe, vpa)
     #average_residual = 0.0
     # @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
     #     if norm[ivpa, ivperp, iz, ir] > eps(mk_float)
@@ -1451,8 +1451,26 @@ function check_electron_pdf_convergence(residual, norm)
     #     end
     #     average_residual /= length(residual)
     # end
-    average_residual = sum(abs.(residual)) / sum(norm)
+
+    # Only sum residual over points that are not set by the sheath boundary condition, as
+    # those that are set by the sheath boundary condition are not used by the time
+    # advance, and so might not converge to 0.
+    # First, sum the contributions from the bulk of the domain
+    sum_residual = sum(abs.(@view residual[:,:,2:end-1,:]))
+    # Then add the contributions from the evolved parts of the sheath boundary points
+    @loop_r ir begin
+        vpa_unnorm_lower = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
+        iv0_lower = findfirst(x -> x>0.0, vpa_unnorm_lower)
+        vpa_unnorm_upper = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
+        iv0_upper = findlast(x -> x>0.0, vpa_unnorm_upper)
+        sum_residual += sum(abs.(@view residual[1:iv0_lower-1,:,1,ir])) +
+                        sum(abs.(@view residual[iv0_upper+1:end,:,end,ir]))
+    end
+
+    average_residual = sum_residual / sum(norm)
+
     electron_pdf_converged = (average_residual < 1e-3)
+
     return average_residual, electron_pdf_converged
 end
 

From 2a9114e6bd2f972b55ccd9901bd3cbab812fe595 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 12:29:20 +0000
Subject: [PATCH 071/428] Go back to setting dt_energy explicitly, add var for
 number of subcycles

---
 src/electron_kinetic_equation.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index c9ae827af..935c2f856 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -145,7 +145,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #dt_electron = dt * sqrt(composition.me_over_mi)
     dt_max = 1.0e-8 #1.0
     #dt_max = 2.5e-9 #1.0
-    dt_energy = 0.001
+    #dt_energy = 0.001
+    dt_energy = 1.0e-7
+    #n_ppar_subcycles = 1000
+    n_ppar_subcycles = 100
     time = 0.0
 
     # define residual to point to a dummy array;
@@ -287,13 +290,12 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             result_phi[:,iteration÷output_interval+1] .= phi[:,1]
         end
 
-        dt_energy = dt_electron * 10.0
+        #dt_energy = dt_electron * 10.0
 
         # get an updated iterate of the electron parallel pressure
         #ppar .= ppar_old
         wpa3_moment = @. qpar / vthe^3
-        #for i in 1:1000
-        for i in 1:100
+        for i in 1:n_ppar_subcycles
             @. qpar = vthe^3 * wpa3_moment
             # Compute the upwinded z-derivative of the electron parallel pressure for the 
             # electron energy equation

From a897f1a01248f69bc36b4fc6dc991f85dc100cf5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 12:38:47 +0000
Subject: [PATCH 072/428] Settings for non-crashing simulation

Using the current code, and the input in
`examples/kinetic-electrons/wall+sheath-bc_kinetic.toml`, time advance
of the electron distribution function continues for a long time without
crashing. Convergence to steady state seems to be very slow, maybe not
helped by lots of ringing in the tails of the distribution function.
---
 .../wall+sheath-bc_kinetic.toml               | 85 +++++++++++++++++++
 src/electron_kinetic_equation.jl              |  4 +-
 src/initial_conditions.jl                     |  4 +-
 3 files changed, 89 insertions(+), 4 deletions(-)
 create mode 100644 examples/kinetic-electrons/wall+sheath-bc_kinetic.toml

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
new file mode 100644
index 000000000..e5a51d1f2
--- /dev/null
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -0,0 +1,85 @@
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+electron_physics = "kinetic_electrons"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+electron_ionization_frequency = 2.0
+ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 9
+#z_nelement = 16
+z_nelement = 32
+#z_nelement = 64
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_spacing_option = "sqrt"
+vpa_ngrid = 17
+#vpa_nelement = 10
+vpa_nelement = 20
+vpa_L = 12.0 #8.0
+vpa_bc = "periodic"
+vpa_discretization = "chebyshev_pseudospectral"
+#vpa_discretization = "gausslegendre_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "periodic"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+#moment_dissipation_coefficient = 0.0001
+#vpa_dissipation_coefficient = 0.002
+vpa_dissipation_coefficient = 0.2
diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 935c2f856..7cc864695 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -143,12 +143,12 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                          buffer_r_4, z_spectral, z)
 
     #dt_electron = dt * sqrt(composition.me_over_mi)
-    dt_max = 1.0e-8 #1.0
+    dt_max = 3.0e-8 #1.0
     #dt_max = 2.5e-9 #1.0
     #dt_energy = 0.001
     dt_energy = 1.0e-7
     #n_ppar_subcycles = 1000
-    n_ppar_subcycles = 100
+    n_ppar_subcycles = 200
     time = 0.0
 
     # define residual to point to a dummy array;
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index cbc77c592..dffe8ffbe 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -510,8 +510,8 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
             # now that we have our initial guess for the electron pdf, we iterate
             # using the time-independent electron kinetic equation to find a self-consistent
             # solution for the electron pdf
-            #max_electron_pdf_iterations = 500000
-            max_electron_pdf_iterations = 10000
+            max_electron_pdf_iterations = 500000
+            #max_electron_pdf_iterations = 10000
             return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
                                         moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
                                         phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 

From 5d0b8a174c70b1c2fc164beba15c2370dd439f9c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 5 Feb 2024 13:29:01 +0000
Subject: [PATCH 073/428] Print the wall-clock time when electron iteration
 terminates

---
 src/electron_kinetic_equation.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 7cc864695..4c313054d 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -1,5 +1,6 @@
 module electron_kinetic_equation
 
+using Dates
 using LinearAlgebra
 
 export get_electron_critical_velocities
@@ -358,6 +359,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     if !electron_pdf_converged
         # need to exit or handle this appropriately
         println("!!!max number of iterations for electron pdf update exceeded!!!")
+        println("Stopping at ", Dates.format(now(), dateformat"H:MM:SS"))
         @loop_vpa ivpa begin
             @loop_z iz begin
                 println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])

From 2c50fcb3aa15da5a72687f661c7472f575cd8c62 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Feb 2024 09:25:50 +0000
Subject: [PATCH 074/428] Use finite element derivative for electron dqpar_dz

---
 src/electron_kinetic_equation.jl | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 4c313054d..541dc3b80 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -249,16 +249,9 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         calculate_electron_qpar_from_pdf!(qpar, ppar, vthe, pdf, vpa)
         qpar_updated = true
         
-        # calculate the updated dqpar/dz
-        for iz ∈ 2:z.n-1
-            dqpar_dz[iz,1] = 0.5*(qpar[iz+1,1]-qpar[iz,1])/z.cell_width[iz] + 0.5*(qpar[iz]-qpar[iz-1,1])/z.cell_width[iz-1]
-        end
-        dqpar_dz[1,1] = (qpar[2,1]-qpar[1,1])/z.cell_width[1]
-        dqpar_dz[z.n,1] = (qpar[end,1]-qpar[end-1,1])/z.cell_width[end-1]
-
         # compute the z-derivative of the parallel electron heat flux
-        #@views derivative_z!(dqpar_dz, qpar, buffer_r_1,
-        #    buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+        @views derivative_z!(dqpar_dz, qpar, buffer_r_1,
+            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
         # Compute the upwinded z-derivative of the electron parallel pressure for the 
         # electron energy equation

From 718cdff1cf39cb655fcfd027475a3d4e9548eb3f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 10 Feb 2024 18:43:38 +0000
Subject: [PATCH 075/428] Don't subcycle ppar

Just use 1 subcycle, so ppar advances at the same rate as g_e.
---
 src/electron_kinetic_equation.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 541dc3b80..b6e9ecbab 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -149,7 +149,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #dt_energy = 0.001
     dt_energy = 1.0e-7
     #n_ppar_subcycles = 1000
-    n_ppar_subcycles = 200
+    #n_ppar_subcycles = 200
+    n_ppar_subcycles = 1
     time = 0.0
 
     # define residual to point to a dummy array;

From 59c84845d6479ded9526aedf51e34cd7be5d144b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Feb 2024 10:09:17 +0000
Subject: [PATCH 076/428] Include begin_*_region() calls in derivative
 functions

Ensures that a correct region type is being used, rather than relying on
the calling site of the derivative_*!() function to have already set a
correct region type.
---
 src/derivatives.jl | 62 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 45 insertions(+), 17 deletions(-)

diff --git a/src/derivatives.jl b/src/derivatives.jl
index d44a001ba..43c07e59c 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -31,6 +31,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,1},
         r_receive_buffer2::AbstractArray{mk_float,1}, r_spectral, r)
 
+        begin_z_region()
+
 	# differentiate f w.r.t r
 	@loop_z iz begin
 		@views derivative!(dfdr[iz,:], f[iz,:], r, r_spectral)
@@ -56,6 +58,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,4},
         r_receive_buffer2::AbstractArray{mk_float,4}, r_spectral, r)
 
+        begin_s_z_vperp_vpa_region()
+
 	# differentiate f w.r.t r
 	@loop_s_z_vperp_vpa is iz ivperp ivpa begin
 		@views derivative!(dfdr[ivpa,ivperp,iz,:,is], f[ivpa,ivperp,iz,:,is], r, r_spectral)
@@ -80,6 +84,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,6}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,5},
         r_receive_buffer2::AbstractArray{mk_float,5}, r_spectral, r)
 
+        begin_sn_z_vzeta_vr_vz_region()
+
 	# differentiate f w.r.t r
 	@loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
 		@views derivative!(dfdr[ivz,ivr,ivzeta,iz,:,isn], f[ivz,ivr,ivzeta,iz,:,isn], r, r_spectral)
@@ -113,6 +119,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,1},
         z_receive_buffer::AbstractArray{mk_float,1}, z_spectral, z)
 
+        begin_r_region()
+
 	# differentiate f w.r.t z
 	@loop_r ir begin
 		@views derivative!(dfdz[:,ir], f[:,ir], z, z_spectral)
@@ -137,6 +145,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,4},
         z_receive_buffer::AbstractArray{mk_float,4}, z_spectral, z)
 
+        begin_s_r_vperp_vpa_region()
+
 	# differentiate f w.r.t z
 	@loop_s_r_vperp_vpa is ir ivperp ivpa begin
 		@views derivative!(dfdz[ivpa,ivperp,:,ir,is], f[ivpa,ivperp,:,ir,is], z, z_spectral)
@@ -161,6 +171,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_floa
 	z_send_buffer::AbstractArray{mk_float,3},
 	z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z)
 
+        begin_r_vperp_vpa_region()
+
 	# differentiate f w.r.t z
 	@loop_r_vperp_vpa ir ivperp ivpa begin
 		@views derivative!(dfdz[ivpa,ivperp,:,ir], f[ivpa,ivperp,:,ir], z, z_spectral)
@@ -185,6 +197,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,5},
         z_receive_buffer::AbstractArray{mk_float,5}, z_spectral, z)
 
+        begin_sn_r_vzeta_vr_vz_region()
+
 	# differentiate f w.r.t z
 	@loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
 		@views derivative!(dfdz[ivz,ivr,ivzeta,:,ir,isn], f[ivz,ivr,ivzeta,:,ir,isn], z, z_spectral)
@@ -220,6 +234,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,1},
         r_receive_buffer2::AbstractArray{mk_float,1}, r_spectral, r)
 
+    begin_z_region()
+
     # differentiate f w.r.t r
     @loop_z iz begin
         @views derivative!(dfdr[iz,:], f[iz,:], r, adv_fac[:,iz], r_spectral)
@@ -248,6 +264,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,4},
         r_receive_buffer2::AbstractArray{mk_float,4}, r_spectral, r)
 
+        begin_s_z_vperp_vpa_region()
+
 	# differentiate f w.r.t r
 	@loop_s_z_vperp_vpa is iz ivperp ivpa begin
 		@views derivative!(dfdr[ivpa,ivperp,iz,:,is], f[ivpa,ivperp,iz,:,is], r, advect[is].adv_fac[:,ivpa,ivperp,iz], r_spectral)
@@ -277,6 +295,8 @@ function derivative_r!(dfdr::AbstractArray{mk_float,6}, f::AbstractArray{mk_floa
         r_receive_buffer1::AbstractArray{mk_float,5},
         r_receive_buffer2::AbstractArray{mk_float,5}, r_spectral, r)
 
+        begin_sn_z_vzeta_vr_vz_region()
+
 	# differentiate f w.r.t r
 	@loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
 		@views derivative!(dfdr[ivz,ivr,ivzeta,iz,:,isn], f[ivz,ivr,ivzeta,iz,:,isn],
@@ -315,6 +335,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,2}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,1},
         z_receive_buffer::AbstractArray{mk_float,1}, z_spectral, z)
 
+    begin_r_region()
+
     # differentiate f w.r.t z
     @loop_r ir begin
         @views derivative!(dfdz[:,ir], f[:,ir], z, adv_fac[:,ir], z_spectral)
@@ -342,6 +364,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,5}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,4},
         z_receive_buffer::AbstractArray{mk_float,4}, z_spectral, z)
 
+        begin_s_r_vperp_vpa_region()
+
 	# differentiate f w.r.t z
 	@loop_s_r_vperp_vpa is ir ivperp ivpa begin
 		@views derivative!(dfdz[ivpa,ivperp,:,ir,is], f[ivpa,ivperp,:,ir,is], z, advect[is].adv_fac[:,ivpa,ivperp,ir], z_spectral)
@@ -371,23 +395,25 @@ function derivative_z!(dfdz::AbstractArray{mk_float,4}, f::AbstractArray{mk_floa
 	z_send_buffer::AbstractArray{mk_float,3},
 	z_receive_buffer::AbstractArray{mk_float,3}, z_spectral, z)
 
-# differentiate the pdf f w.r.t z
-@loop_r_vperp_vpa ir ivperp ivpa begin
-	@views derivative!(dfdz[ivpa,ivperp,:,ir], f[ivpa,ivperp,:,ir], z, advect[1].adv_fac[:,ivpa,ivperp,ir], z_spectral)
-	# get external endpoints to reconcile via MPI
-	dfdz_lower_endpoints[ivpa,ivperp,ir] = z.scratch_2d[1,1]
-	dfdz_upper_endpoints[ivpa,ivperp,ir] = z.scratch_2d[end,end]
-	adv_fac_lower_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[1,ivpa,ivperp,ir]
-	adv_fac_upper_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[end,ivpa,ivperp,ir]
-end
-# now reconcile element boundaries across
-# processes with large message
-if z.nelement_local < z.nelement_global
-	reconcile_element_boundaries_MPI!(dfdz,
-	 adv_fac_lower_buffer, adv_fac_upper_buffer,
-	 dfdz_lower_endpoints,dfdz_upper_endpoints,
-	 z_send_buffer, z_receive_buffer, z)
-end
+    begin_r_vperp_vpa_region()
+
+    # differentiate the pdf f w.r.t z
+    @loop_r_vperp_vpa ir ivperp ivpa begin
+            @views derivative!(dfdz[ivpa,ivperp,:,ir], f[ivpa,ivperp,:,ir], z, advect[1].adv_fac[:,ivpa,ivperp,ir], z_spectral)
+            # get external endpoints to reconcile via MPI
+            dfdz_lower_endpoints[ivpa,ivperp,ir] = z.scratch_2d[1,1]
+            dfdz_upper_endpoints[ivpa,ivperp,ir] = z.scratch_2d[end,end]
+            adv_fac_lower_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[1,ivpa,ivperp,ir]
+            adv_fac_upper_buffer[ivpa,ivperp,ir] = advect[1].adv_fac[end,ivpa,ivperp,ir]
+    end
+    # now reconcile element boundaries across
+    # processes with large message
+    if z.nelement_local < z.nelement_global
+            reconcile_element_boundaries_MPI!(dfdz,
+             adv_fac_lower_buffer, adv_fac_upper_buffer,
+             dfdz_lower_endpoints,dfdz_upper_endpoints,
+             z_send_buffer, z_receive_buffer, z)
+    end
 
 end
 
@@ -400,6 +426,8 @@ function derivative_z!(dfdz::AbstractArray{mk_float,6}, f::AbstractArray{mk_floa
         z_send_buffer::AbstractArray{mk_float,5},
         z_receive_buffer::AbstractArray{mk_float,5}, z_spectral, z)
 
+        begin_sn_r_vzeta_vr_vz_region()
+
 	# differentiate f w.r.t z
 	@loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
 		@views derivative!(dfdz[ivz,ivr,ivzeta,:,ir,isn], f[ivz,ivr,ivzeta,:,ir,isn],

From 937b065ab2e4748449d39b002aeadabe403a339a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Feb 2024 10:27:49 +0000
Subject: [PATCH 077/428] Remove hacked-in z-diffusion and vpa-diffusion of
 electron pdf

---
 src/electron_vpa_advection.jl | 9 +++++----
 src/electron_z_advection.jl   | 9 +++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/electron_vpa_advection.jl b/src/electron_vpa_advection.jl
index c46686859..3d60b8945 100644
--- a/src/electron_vpa_advection.jl
+++ b/src/electron_vpa_advection.jl
@@ -27,12 +27,13 @@ function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar
     @loop_r_z_vperp ir iz ivperp begin
         @views derivative!(dpdf_dvpa[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, advect[1].adv_fac[:,ivperp,iz,ir], spectral)
     end
-    @loop_r_z_vperp ir iz ivperp begin
-        @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, spectral)
-    end
+    #@loop_r_z_vperp ir iz ivperp begin
+    #    @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, spectral)
+    #end
     # calculate the advection term
     @loop_vpa ivpa begin
-        @. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:] + 0.0001*d2pdf_dvpa2[ivpa,:,:,:]
+        @. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:]
+        #@. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:] + 0.0001*d2pdf_dvpa2[ivpa,:,:,:]
     end
     #@loop_vpa ivpa begin
     #    println("electron_vpa_advection: ", advection_term[ivpa,1,10,1], " vpa: ", vpa.grid[ivpa], " dpdf_dvpa: ", dpdf_dvpa[ivpa,1,10,1],
diff --git a/src/electron_z_advection.jl b/src/electron_z_advection.jl
index 57349cff5..e9cf1e43f 100644
--- a/src/electron_z_advection.jl
+++ b/src/electron_z_advection.jl
@@ -30,12 +30,13 @@ function electron_z_advection!(advection_term, pdf, vth, advect, z, vpa, spectra
                   scratch_dummy.buffer_vpavperpr_2, scratch_dummy.buffer_vpavperpr_3,
                   scratch_dummy.buffer_vpavperpr_4, scratch_dummy.buffer_vpavperpr_5,
                   scratch_dummy.buffer_vpavperpr_6, spectral, z)
-    @loop_r_vperp_vpa ir ivperp ivpa begin
-        @views second_derivative!(d2pdf_dz2[ivpa,ivperp,:,ir], pdf[ivpa,ivperp,:,ir], z, spectral)
-    end
+    #@loop_r_vperp_vpa ir ivperp ivpa begin
+    #    @views second_derivative!(d2pdf_dz2[ivpa,ivperp,:,ir], pdf[ivpa,ivperp,:,ir], z, spectral)
+    #end
     # calculate the advection term
     @loop_z iz begin
-        @. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:] + 0.0001*d2pdf_dz2[:,:,iz,:]
+        @. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:]
+        #@. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:] + 0.0001*d2pdf_dz2[:,:,iz,:]
     end
     return nothing
 end

From 3ae772e0f8c939fd52e66f9614ee8ce72b13f343 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Feb 2024 11:09:59 +0000
Subject: [PATCH 078/428] Add a couple of shared-memory scratch arrays to
 coordinate structs

---
 src/coordinates.jl | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/coordinates.jl b/src/coordinates.jl
index 8a8d520d6..a5224b14e 100644
--- a/src/coordinates.jl
+++ b/src/coordinates.jl
@@ -8,9 +8,10 @@ export set_element_boundaries
 
 using LinearAlgebra
 using ..type_definitions: mk_float, mk_int
-using ..array_allocation: allocate_float, allocate_int
+using ..array_allocation: allocate_float, allocate_shared_float, allocate_int
 using ..calculus: derivative!
 using ..chebyshev: scaled_chebyshev_grid, scaled_chebyshev_radau_grid, setup_chebyshev_pseudospectral
+using ..communication
 using ..finite_differences: finite_difference_info
 using ..gauss_legendre: scaled_gauss_legendre_lobatto_grid, scaled_gauss_legendre_radau_grid, setup_gausslegendre_pseudospectral
 using ..quadrature: composite_simpson_weights
@@ -77,6 +78,12 @@ struct coordinate
     scratch2::Array{mk_float,1}
     # scratch3 is an array used for intermediate calculations requiring n entries
     scratch3::Array{mk_float,1}
+    # scratch_shared is a shared-memory array used for intermediate calculations requiring
+    # n entries
+    scratch_shared::MPISharedArray{mk_float,1}
+    # scratch_shared2 is a shared-memory array used for intermediate calculations requiring
+    # n entries
+    scratch_shared2::MPISharedArray{mk_float,1}
     # scratch_2d and scratch2_2d are arrays used for intermediate calculations requiring
     # ngrid x nelement entries
     scratch_2d::Array{mk_float,2}
@@ -136,6 +143,8 @@ function define_coordinate(input, parallel_io::Bool=false; init_YY::Bool=true)
     duniform_dgrid = allocate_float(input.ngrid, input.nelement_local)
     # scratch is an array used for intermediate calculations requiring n entries
     scratch = allocate_float(n_local)
+    scratch_shared = allocate_shared_float(n_local)
+    scratch_shared2 = allocate_shared_float(n_local)
     # scratch_2d is an array used for intermediate calculations requiring ngrid x nelement entries
     scratch_2d = allocate_float(input.ngrid, input.nelement_local)
     # struct containing the advection speed options/inputs for this coordinate
@@ -165,7 +174,7 @@ function define_coordinate(input, parallel_io::Bool=false; init_YY::Bool=true)
     coord = coordinate(input.name, n_global, n_local, input.ngrid,
         input.nelement_global, input.nelement_local, input.nrank, input.irank, input.L, grid,
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
-        input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch),
+        input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch), copy(scratch), scratch_shared, scratch_shared2,
         scratch_2d, copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
         local_io_range, global_io_range, element_scale, element_shift, input.element_spacing_option)
 

From a11748bc54426cba67d1d18b3fa2772998623dae Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Feb 2024 11:25:21 +0000
Subject: [PATCH 079/428] Parallelise electron physics functions

---
 src/electron_fluid_equations.jl  | 133 ++++++----
 src/electron_kinetic_equation.jl | 415 ++++++++++++++++++++-----------
 src/electron_vpa_advection.jl    |  12 +-
 src/electron_z_advection.jl      |  12 +-
 src/em_fields.jl                 |  63 +++--
 src/initial_conditions.jl        | 177 ++++++-------
 src/time_advance.jl              |   2 +-
 7 files changed, 503 insertions(+), 311 deletions(-)

diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index 17d6fe8f0..9d3064dda 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -7,11 +7,14 @@ export calculate_electron_qpar!
 export calculate_electron_parallel_friction_force!
 export calculate_electron_qpar_from_pdf!
 
+using ..communication
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
 using ..input_structs: braginskii_fluid, kinetic_electrons
 using ..velocity_moments: integrate_over_vspace
 
+using MPI
+
 """
 use quasineutrality to obtain the electron density from the initial
 densities of the various ion species:
@@ -27,10 +30,13 @@ output:
 function calculate_electron_density!(dens_e, updated, dens_i)
     # only update the electron density if it has not already been updated
     if !updated
-        dens_e .= 0.0
+        begin_r_z_region()
         # enforce quasineutrality
-        @loop_s_r_z is ir iz begin
-            dens_e[iz,ir] += dens_i[iz,ir,is]
+        @loop_r_z ir iz begin
+            dens_e[iz,ir] = 0.0
+            @loop_s is begin
+                dens_e[iz,ir] += dens_i[iz,ir,is]
+            end
         end
     end
     # set flag indicating that the electron density has been updated
@@ -51,41 +57,52 @@ inputs:
 output:
     upar_e = contains the updated electron parallel flow
 """
-function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
+function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model, r, z)
     # only calculate the electron parallel flow if it is not already updated
     if !updated
+        begin_r_region()
         # get the number of zed grid points, nz
         nz = size(upar_e,1)
         # initialise the electron parallel flow density to zero
-        upar_e .= 0.0
+        @loop_r_z ir iz begin
+            upar_e[iz,ir] = 0.0
+        end
         # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
         # is equal and opposite to the ion parallel current
         if electron_model ∈ (boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons)
-            # loop over ion species, adding each species contribution to the
-            # ion parallel particle flux at the boundaries in zed
-            @loop_s_r is ir begin
-                # electron_upar at this intermediate stage is actually
-                # the electron particle flux
-                upar_e[1,ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
-                upar_e[end,ir] += dens_i[end,ir,is] * upar_i[end,ir,is]
-            end
-            @loop_r ir begin
-                # fix the boundary flux
-                boundary_flux = upar_e[1,ir]
-                for iz in 2:nz-1
-                    # initialise the electron particle flux to its value at the boundary in zed
-                    upar_e[iz,ir] = boundary_flux
-                    # add the contributions to the electron particle flux from the various ion species
-                    # particle fluxes
+            boundary_flux = r.scratch_shared
+            boundary_ion_flux = r.scratch_shared2
+            if z.irank == 0
+                @loop_r ir begin
+                    boundary_flux[ir] = 0.0
+                    boundary_ion_flux[ir] = 0.0
                     @loop_s is begin
-                        upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is] - dens_i[1,ir,is] * upar_i[1,ir,is]
+                        boundary_flux[ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
+                        boundary_ion_flux[ir] += dens_i[1,ir,is] * upar_i[1,ir,is]
                     end
-                    # convert from parallel particle flux to parallel particle density
-                    upar_e[iz,ir] /= dens_e[iz,ir]
                 end
-                # convert from parallel particle flux to parallel particle density for boundary points
-                upar_e[1,ir] /= dens_e[1,ir]
-                upar_e[end,ir] /= dens_e[end,ir]
+            end
+            begin_serial_region()
+            @serial_region begin
+                MPI.Bcast!(boundary_flux, 0, z.comm)
+                MPI.Bcast!(boundary_ion_flux, 0, z.comm)
+            end
+            # loop over ion species, adding each species contribution to the
+            # ion parallel particle flux at the boundaries in zed
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                # initialise the electron particle flux to its value at the boundary in
+                # zed and subtract the ion boundary flux - we want to calculate upar_e =
+                # boundary_flux + (ion_flux - boundary_ion_flux) as at this intermediate
+                # point upar is actually the electron particle flux
+                upar_e[iz,ir] = boundary_flux[ir] - boundary_ion_flux[ir]
+                # add the contributions to the electron particle flux from the various ion species
+                # particle fluxes
+                @loop_s is begin
+                    upar_e[iz,ir] += dens_i[iz,ir,is] * upar_i[iz,ir,is]
+                end
+                # convert from parallel particle flux to parallel particle density
+                upar_e[iz,ir] /= dens_e[iz,ir]
             end
         end
     end
@@ -100,7 +117,7 @@ NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
 function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt, composition,
-                                   num_diss_params, zgrid)
+                                   num_diss_params, z)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
     me_over_mi = composition.me_over_mi
@@ -137,13 +154,13 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     # add in contributions due to charge exchange/ionization collisions
     if composition.n_neutral_species > 0
         if abs(collisions.charge_exchange_electron) > 0.0
-            @loop_s_r_z is ir iz begin
+            @loop_sn_r_z isn ir iz begin
                 ppar[iz,ir] +=
                     dt * me_over_mi * collisions.charge_exchange_electron * (
-                    2*(fvec.electron_density[iz,ir]*fvec.pz_neutral[iz,ir,is] -
-                    fvec.density_neutral[iz,ir,is]*fvec.electron_ppar[iz,ir]) +
-                    (2/3)*fvec.electron_density[iz,ir]*fvec.density_neutral[iz,ir,is] *
-                    (fvec.uz_neutral[iz,ir,is] - fvec.electron_upar[iz,ir])^2)
+                    2*(fvec.electron_density[iz,ir]*fvec.pz_neutral[iz,ir,isn] -
+                    fvec.density_neutral[iz,ir,isn]*fvec.electron_ppar[iz,ir]) +
+                    (2/3)*fvec.electron_density[iz,ir]*fvec.density_neutral[iz,ir,isn] *
+                    (fvec.uz_neutral[iz,ir,isn] - fvec.electron_upar[iz,ir])^2)
             end
         end
         if abs(collisions.ionization_electron) > 0.0
@@ -153,9 +170,9 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
             #         fvec.electron_ppar[iz,ir] -
             #         (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
             # end
-            @loop_s_r_z is ir iz begin
+            @loop_sn_r_z isn ir iz begin
                 ppar[iz,ir] +=
-                    dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
+                    dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,isn] * (
                     fvec.electron_ppar[iz,ir] -
                     fvec.electron_density[iz,ir] * collisions.ionization_energy)
             end
@@ -165,8 +182,9 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     calculate_electron_heat_source!(moments.electron.heat_source, fvec.electron_ppar, moments.electron.dupar_dz,
                                     fvec.density_neutral, collisions.ionization, collisions.ionization_energy,
                                     fvec.electron_density, fvec.ppar, collisions.nu_ei, composition.me_over_mi,
-                                    composition.T_wall, zgrid)
+                                    composition.T_wall, z)
     # add the contribution from the electron heat source                                
+    begin_r_z_region()
     @loop_r_z ir iz begin
         ppar[iz,ir] += dt * moments.electron.heat_source[iz,ir]
     end
@@ -197,6 +215,7 @@ output:
 """
 function calculate_Epar_from_electron_force_balance!(Epar, dens_e, dppar_dz, nu_ei, friction, n_neutral_species,
                                                      charge_exchange, me_over_mi, dens_n, upar_n, upar_e)
+    begin_r_z_region()
     # get the contribution to Epar from the parallel pressure
     @loop_r_z ir iz begin
         Epar[iz,ir] = -(2/dens_e[iz,ir]) * dppar_dz[iz,ir]
@@ -209,8 +228,8 @@ function calculate_Epar_from_electron_force_balance!(Epar, dens_e, dppar_dz, nu_
     end
     # if there are neutral species evolved and accounting for charge exchange collisions with neutrals
     if n_neutral_species > 0 && charge_exchange > 0
-        @loop_s_r_z is ir iz begin
-            Epar[iz,ir] += 2 * me_over_mi * dens_n[iz,ir] * charge_exchange * (upar_n[iz,ir,is] - upar_e[iz,ir]) 
+        @loop_sn_r_z isn ir iz begin
+            Epar[iz,ir] += 2 * me_over_mi * dens_n[iz,ir] * charge_exchange * (upar_n[iz,ir,isn] - upar_e[iz,ir]) 
         end
     end
     return nothing
@@ -220,6 +239,7 @@ end
 """
 function calculate_electron_parallel_friction_force!(friction, dens_e, upar_e, upar_i, dTe_dz,
                                                      me_over_mi, nu_ei, electron_model)
+    begin_r_z_region()
     if electron_model == braginskii_fluid
         @loop_r_z ir iz begin
             friction[iz,ir] = -(1/2) * 0.71 * dens_e[iz,ir] * dTe_dz[iz,ir]
@@ -228,7 +248,9 @@ function calculate_electron_parallel_friction_force!(friction, dens_e, upar_e, u
             friction[iz,ir] += 0.51 * dens_e[iz,ir] * me_over_mi * nu_ei * (upar_i[iz,ir,is] - upar_e[iz,ir])
         end
     else
-        @. friction = 0.0
+        @loop_r_z ir iz begin
+            friction[iz,ir] = 0.0
+        end
     end
     return nothing
 end
@@ -259,11 +281,14 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth
                                   nu_ei, me_over_mi, electron_model, vpa)
     # only calculate qpar_e if needs updating
     if qpar_updated == false
-        @. qpar_e = 0.0
         if electron_model == braginskii_fluid
+            begin_r_z_region()
             # use the classical Braginskii expression for the electron heat flux
-            @loop_s_r_z is ir iz begin
-                qpar_e[iz,ir] -= 0.71 * ppar_e[iz,ir] * (upar_i[iz,ir,is] - upar_e[iz,ir])
+            @loop_r_z ir iz begin
+                qpar_e[iz,ir] = 0.0
+                @loop_s is begin
+                    qpar_e[iz,ir] -= 0.71 * ppar_e[iz,ir] * (upar_i[iz,ir,is] - upar_e[iz,ir])
+                end
             end
             if nu_ei > 0.0
                 @loop_r_z ir iz begin
@@ -285,6 +310,7 @@ calculate the parallel component of the electron heat flux,
 defined as qpar = 2 * ppar * vth * int dwpa (pdf * wpa^3)
 """
 function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
+    begin_r_z_region()
     # specialise to 1D for now
     ivperp = 1
     @loop_r_z ir iz begin
@@ -293,7 +319,8 @@ function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
 end
 
 function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n, ionization, ionization_energy,
-                                         dens_e, ppar_i, nu_ei, me_over_mi, T_wall, zgrid)
+                                         dens_e, ppar_i, nu_ei, me_over_mi, T_wall, z)
+    begin_r_z_region()
     # heat_source currently only used for testing                                         
     # @loop_r_z ir iz begin
     #     heat_source[iz,ir] = (2/3) * ppar_e[iz,ir] * dupar_dz[iz,ir]
@@ -315,19 +342,19 @@ function calculate_electron_heat_source!(heat_source, ppar_e, dupar_dz, dens_n,
     # end
     # Gaussian heat deposition profile, with a decay of 5 e-foldings at the ends of the domain in z
     @loop_r_z ir iz begin
-        heat_source[iz,ir] = 50*exp(-5*(zgrid[iz]/zgrid[end])^2)
+        heat_source[iz,ir] = 50*exp(-5*(2.0*z.grid[iz]/z.L)^2)
     end
     return nothing
 end
 
-function enforce_parallel_BC_on_electron_pressure!(ppar, dens, T_wall, ppar_i)
-    # assume T_e = T_i at boundaries in z
-    @loop_r ir begin
-        #ppar[1,ir] = 0.5 * dens[1,ir] * T_wall
-        #ppar[end,ir] = 0.5 * dens[end,ir] * T_wall
-        ppar[1,ir] = ppar_i[1,ir,1]
-        ppar[end,ir] = ppar_i[end,ir,1]
-    end
-end
+#function enforce_parallel_BC_on_electron_pressure!(ppar, dens, T_wall, ppar_i)
+#    # assume T_e = T_i at boundaries in z
+#    @loop_r ir begin
+#        #ppar[1,ir] = 0.5 * dens[1,ir] * T_wall
+#        #ppar[end,ir] = 0.5 * dens[end,ir] * T_wall
+#        ppar[1,ir] = ppar_i[1,ir,1]
+#        ppar[end,ir] = ppar_i[end,ir,1]
+#    end
+#end
 
 end
diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index b6e9ecbab..2f7baff64 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -2,12 +2,14 @@ module electron_kinetic_equation
 
 using Dates
 using LinearAlgebra
+using MPI
 
 export get_electron_critical_velocities
 
 using ..looping
 using ..derivatives: derivative_z!
 using ..calculus: derivative!, second_derivative!, integral
+using ..communication
 using ..interpolation: interpolate_to_grid_1d!
 using ..type_definitions: mk_float
 using ..array_allocation: allocate_float
@@ -104,13 +106,17 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
     num_diss_params, max_electron_pdf_iterations)
 
+    begin_r_z_region()
+
     #println("TMP FOR TESTING: SETTING UPAR_I = UPAR_E = 0!!!")
     #moments.electron.upar .= 0.0
 
     # there will be a better way of doing this
     # store the incoming ppar
-    ppar_old = allocate_float(z.n,1)
-    ppar_old .= ppar
+    #ppar_old = allocate_float(z.n,1)
+    #ppar_old .= ppar
+
+    upar = moments.electron.upar
 
     # create a (z,r) dimension dummy array for use in taking derivatives
     dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
@@ -123,7 +129,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
 
     # compute the z-derivative of the input electron parallel flow, needed for the electron kinetic equation
-    @views derivative_z!(moments.electron.dupar_dz, moments.electron.upar, buffer_r_1, buffer_r_2, buffer_r_3,
+    @views derivative_z!(moments.electron.dupar_dz, upar, buffer_r_1, buffer_r_2, buffer_r_3,
                          buffer_r_4, z_spectral, z)
 
     # compute the z-derivative of the input electron parallel pressure, needed for the electron kinetic equation
@@ -134,11 +140,14 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #@. dens = 1.0
     #@. ddens_dz = (dppar_dz / ppar)*0.9
 
-    # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
-    # and parallel pressure
-    @. dvth_dz = 0.5 * vthe * (dppar_dz / ppar - ddens_dz / dens)
-    # update the electron thermal speed itself using the updated electron parallel pressure
-    @. vthe = sqrt(abs(2.0 * ppar / (dens * composition.me_over_mi)))
+    @loop_r_z ir iz begin
+        # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
+        # and parallel pressure
+        dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
+        # update the electron thermal speed itself using the updated electron parallel pressure
+        vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
+    end
+
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
     @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
                          buffer_r_4, z_spectral, z)
@@ -169,65 +178,94 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                         num_diss_params, dt_max)
     # Divide by wpa to relax CFL condition at large wpa - only looking for steady
     # state here, so does not matter that this makes time evolution incorrect.
+    # Also increase the effective timestep for z-values far from the sheath boundary -
+    # these have a less-limited timestep so letting them evolve faster speeds up
+    # convergence to the steady state.
+    Lz = z.L
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-        residual[ivpa,ivperp,iz,ir] /= sqrt(1.0 + vpa.grid[ivpa]^2)
+        zval = z.grid[iz]
+        znorm = 2.0*zval/Lz
+        residual[ivpa,ivperp,iz,ir] *=
+            (1.0 + 20.0*(1.0 - znorm^2)) /
+            sqrt(1.0 + vpa.grid[ivpa]^2)
     end
-    # open files to write the electron heat flux and pdf to file                                
-    io_upar = open("upar.txt", "w")
-    io_qpar = open("qpar.txt", "w")
-    io_ppar = open("ppar.txt", "w")
-    io_pdf = open("pdf.txt", "w")
-    io_vth = open("vth.txt", "w")
-    if !electron_pdf_converged
-        # need to exit or handle this appropriately
-        @loop_vpa ivpa begin
+    if n_blocks[] == 1
+        text_output_suffix = ""
+    else
+        text_output_suffix = "$(iblock_index[])"
+    end
+    begin_serial_region()
+    @serial_region begin
+        # open files to write the electron heat flux and pdf to file                                
+        io_upar = open("upar$text_output_suffix.txt", "w")
+        io_qpar = open("qpar$text_output_suffix.txt", "w")
+        io_ppar = open("ppar$text_output_suffix.txt", "w")
+        io_pdf = open("pdf$text_output_suffix.txt", "w")
+        io_vth = open("vth$text_output_suffix.txt", "w")
+        if !electron_pdf_converged
+            # need to exit or handle this appropriately
+            @loop_vpa ivpa begin
+                @loop_z iz begin
+                    println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                end
+                println(io_pdf,"")
+            end
             @loop_z iz begin
-                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                println(io_upar, "z: ", z.grid[iz], " upar: ", upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
             end
-            println(io_pdf,"")
-        end
-        @loop_z iz begin
-            println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-            println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-            println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-            println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+            println(io_upar,"")
+            println(io_qpar,"")
+            println(io_ppar,"")
+            println(io_vth,"")
         end
-        println(io_upar,"")
-        println(io_qpar,"")
-        println(io_ppar,"")
-        println(io_vth,"")
+        io_pdf_stages = open("pdf_zright$text_output_suffix.txt", "w")
     end
-    #stop()
-
-    io_pdf_stages = open("pdf_zright.txt", "w")
 
     # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
     #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf), moments.electron.upar, vthe, vpa)
-    #println("TMP FOR TESTING -- enforce_boundary_condition_on_electron_pdf needs uncommenting!!!")
-    # evolve (artificially) in time until the residual is less than the tolerance
+    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
+
+    result_pdf = nothing
+    result_ppar = nothing
+    result_vth = nothing
+    result_qpar = nothing
+    result_phi = nothing
+    result_residual = nothing
     output_interval = 1000
-    result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
-    result_pdf[:,:,1] .= pdf[:,1,:,1]
-    result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-    result_ppar[:,1] .= ppar[:,1]
-    result_vth = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-    result_vth[:,1] .= vthe[:,1]
-    result_qpar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-    result_qpar[:,1] .= qpar[:,1]
-    result_phi = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-    result_phi[:,1] .= phi[:,1]
-    result_residual = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
+    begin_serial_region()
+    @serial_region begin
+        result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
+        result_pdf[:,:,1] .= pdf[:,1,:,1]
+        result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+        result_ppar[:,1] .= ppar[:,1]
+        result_vth = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+        result_vth[:,1] .= vthe[:,1]
+        result_qpar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+        result_qpar[:,1] .= qpar[:,1]
+        result_phi = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
+        result_phi[:,1] .= phi[:,1]
+        result_residual = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
+    end
+    # evolve (artificially) in time until the residual is less than the tolerance
     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+        begin_r_z_region()
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
-        @. pdf -= dt_electron * residual
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            pdf[ivpa,ivperp,iz,ir] -= dt_electron * residual[ivpa,ivperp,iz,ir]
+        end
 
         # enforce the boundary condition(s) on the electron pdf
-        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, moments.electron.upar, vpa, vpa_spectral, composition.me_over_mi)
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
+
         #pdf = max.(pdf, 0.0)
-        for ir ∈ 1:size(ppar, 2), iz ∈ 2:size(ppar,1)-1
+
+        begin_r_z_region()
+        @loop_r_z ir iz begin
             #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
             @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
         end
@@ -236,10 +274,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         #error("foo")
         
         if (mod(iteration,output_interval)==1)
-            @loop_vpa ivpa begin
-                println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " iteration: ", iteration, " flag: ", 1)
+            begin_serial_region()
+            @serial_region begin
+                @loop_vpa ivpa begin
+                    println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " iteration: ", iteration, " flag: ", 1)
+                end
+                println(io_pdf_stages,"")
             end
-            println(io_pdf_stages,"")
         end
 
         # update the time following the pdf update
@@ -256,7 +297,9 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         # Compute the upwinded z-derivative of the electron parallel pressure for the 
         # electron energy equation
-        dummy_zr .= -moments.electron.upar
+        @loop_r_z ir iz begin
+            dummy_zr[iz,ir] = -upar[iz,ir]
+        end
         @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
@@ -266,35 +309,49 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
 
+        if (mod(iteration,100) == 0)
+            begin_serial_region()
+            @serial_region begin
+                println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,[1,end]], " average_residual: ", average_residual)
+            end
+        end
         if (mod(iteration,output_interval) == 0)
-            println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,[1,end]], " average_residual: ", average_residual)
-            @loop_z iz begin
-                println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+            begin_serial_region()
+            @serial_region begin
+                @loop_z iz begin
+                    println(io_upar, "z: ", z.grid[iz], " upar: ", upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
+                    println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+                end
+                println(io_upar,"")
+                println(io_qpar,"")
+                println(io_ppar,"")
+                println(io_vth,"")
+                result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
+                result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
+                result_vth[:,iteration÷output_interval+1] .= vthe[:,1]
+                result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
+                result_phi[:,iteration÷output_interval+1] .= phi[:,1]
             end
-            println(io_upar,"")
-            println(io_qpar,"")
-            println(io_ppar,"")
-            println(io_vth,"")
-            result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
-            result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
-            result_vth[:,iteration÷output_interval+1] .= vthe[:,1]
-            result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
-            result_phi[:,iteration÷output_interval+1] .= phi[:,1]
         end
 
         #dt_energy = dt_electron * 10.0
 
         # get an updated iterate of the electron parallel pressure
+        begin_r_z_region()
         #ppar .= ppar_old
-        wpa3_moment = @. qpar / vthe^3
+        wpa3_moment = @view scratch_dummy.buffer_zrs_1[:,:,1]
+        @loop_r_z ir iz begin
+            wpa3_moment[iz,ir] = qpar[iz,ir] / vthe[iz,ir]^3
+        end
         for i in 1:n_ppar_subcycles
-            @. qpar = vthe^3 * wpa3_moment
+            @loop_r_z ir iz begin
+                qpar[iz,ir] = vthe[iz,ir]^3 * wpa3_moment[iz,ir]
+                dummy_zr[iz,ir] = -upar[iz,ir]
+            end
             # Compute the upwinded z-derivative of the electron parallel pressure for the 
             # electron energy equation
-            dummy_zr .= -moments.electron.upar
             @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
@@ -305,17 +362,23 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             end
 
             #dt_energy = dt_electron
-            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z.grid)
-            fvec.electron_ppar .= ppar
+            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                fvec.electron_ppar[iz,ir] = ppar[iz,ir]
+            end
         
-        # compute the z-derivative of the updated electron parallel pressure
-        @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
-                             buffer_r_4, z_spectral, z)
-        # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
-        # and parallel pressure
-        @. dvth_dz = 0.5 * vthe * (dppar_dz / ppar - ddens_dz / dens)
-        # update the electron thermal speed itself using the updated electron parallel pressure
-        @. vthe = sqrt(abs(2.0 * ppar / (dens * composition.me_over_mi)))
+            # compute the z-derivative of the updated electron parallel pressure
+            @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
+                                 buffer_r_4, z_spectral, z)
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
+                # and parallel pressure
+                dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
+                # update the electron thermal speed itself using the updated electron parallel pressure
+                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
+            end
         end
 
         #@views derivative_z!(dqpar_dz, qpar, 
@@ -330,13 +393,20 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                             num_diss_params, dt_max)
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, abs.(pdf), moments.electron.upar, vthe, vpa)
+        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
         if (mod(iteration,output_interval) == 0)
-            result_residual[:,:,iteration÷output_interval+1] .= residual[:,1,:,1]
+            begin_serial_region()
+            @serial_region begin
+                result_residual[:,:,iteration÷output_interval+1] .= residual[:,1,:,1]
+            end
         end
 
         # Divide by wpa to relax CFL condition at large wpa - only looking for steady
         # state here, so does not matter that this makes time evolution incorrect.
+        # Also increase the effective timestep for z-values far from the sheath boundary -
+        # these have a less-limited timestep so letting them evolve faster speeds up
+        # convergence to the steady state.
+        begin_r_z_vperp_region()
         Lz = z.L
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
             zval = z.grid[iz]
@@ -350,24 +420,27 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
         iteration += 1
     end
-    if !electron_pdf_converged
-        # need to exit or handle this appropriately
-        println("!!!max number of iterations for electron pdf update exceeded!!!")
-        println("Stopping at ", Dates.format(now(), dateformat"H:MM:SS"))
-        @loop_vpa ivpa begin
-            @loop_z iz begin
-                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+    begin_serial_region()
+    @serial_region begin
+        if !electron_pdf_converged
+            # need to exit or handle this appropriately
+            println("!!!max number of iterations for electron pdf update exceeded!!!")
+            println("Stopping at ", Dates.format(now(), dateformat"H:MM:SS"))
+            @loop_vpa ivpa begin
+                @loop_z iz begin
+                    println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                end
+                println(io_pdf,"")
             end
-            println(io_pdf,"")
         end
+        close(io_upar)
+        close(io_qpar)
+        close(io_ppar)
+        close(io_vth)
+        close(io_pdf)
+        close(io_pdf_stages)
     end
-    close(io_upar)
-    close(io_qpar)
-    close(io_ppar)
-    close(io_vth)
-    close(io_pdf)
-    close(io_pdf_stages)
-    return result_pdf, dens, moments.electron.upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual
+    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
@@ -379,6 +452,8 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
     # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
     # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
 
+    begin_r_region()
+
     # pdf_adjustment_option determines the velocity-dependent pre-factor for the
     # corrections to the pdf needed to ensure moment constraints are satisfied
     pdf_adjustment_option = "vpa4_gaussian"
@@ -392,12 +467,14 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
     reversed_pdf = vpa.scratch
 
     # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
-    ivpa_zero = (vpa.n+1)÷2
+    #ivpa_zero = (vpa.n+1)÷2
 
     @loop_r ir begin
         # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
         #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
         #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
         reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[1,ir] / vthe[1,ir]
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
@@ -599,6 +676,8 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
     @loop_r ir begin
         # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
         #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
         reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[end,ir] / vthe[end,ir]
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
@@ -875,7 +954,7 @@ function update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
     # initialise the RHS to zero
     rhs .= 0.0
     # get critical velocities beyond which electrons are lost to the wall
-    crit_speed_zmin, crit_speed_zmax = get_electron_critical_velocities(phi, vthe, composition.me_over_mi)
+    crit_speed_zmin, crit_speed_zmax = get_electron_critical_velocities(phi, vthe, composition.me_over_mi, z)
     # add the contribution to rhs from the term proportional to the pdf (rather than its derivatives)
     add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
     # add the contribution to rhs from the wpa advection term
@@ -1054,12 +1133,15 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
                                              num_diss_params, dt_electron)
 
     # initialise the residual to zero                                             
-    residual .= 0.0
+    begin_r_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        residual[ivpa,ivperp,iz,ir] = 0.0
+    end
     # calculate the contribution to the residual from the z advection term
     electron_z_advection!(residual, pdf, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
     #dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
-    single_term .= residual
-    max_term .= abs.(residual)
+    #single_term .= residual
+    #max_term .= abs.(residual)
     #println("z_adv residual = ", maximum(abs.(single_term)))
     #println("z_advection: ", sum(residual), " dqpar_dz: ", sum(abs.(dqpar_dz)))
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
@@ -1067,16 +1149,16 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
                             vpa_advect, vpa, vpa_spectral, scratch_dummy)
     #dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
-    @. single_term = residual - single_term
-    max_term .= max.(max_term, abs.(single_term))
-    @. single_term = residual
+    #@. single_term = residual - single_term
+    #max_term .= max.(max_term, abs.(single_term))
+    #@. single_term = residual
     #println("v_adv residual = ", maximum(abs.(single_term)))
     #add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
     # add in the contribution to the residual from the term proportional to the pdf
     add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
-    @. single_term = residual - single_term
-    max_term .= max.(max_term, abs.(single_term))
-    @. single_term = residual
+    #@. single_term = residual - single_term
+    #max_term .= max.(max_term, abs.(single_term))
+    #@. single_term = residual
     #println("pdf_term residual = ", maximum(abs.(single_term)))
     # @loop_vpa ivpa begin
     #     @loop_z iz begin
@@ -1087,9 +1169,9 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     # println("")
     # add in numerical dissipation terms
     add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
-    @. single_term = residual - single_term
+    #@. single_term = residual - single_term
     #println("dissipation residual = ", maximum(abs.(single_term)))
-    max_term .= max.(max_term, abs.(single_term))
+    #max_term .= max.(max_term, abs.(single_term))
     # add in particle and heat source term(s)
     #@. single_term = residual
     #add_source_term!(residual, vpa.grid, z.grid, dvth_dz)
@@ -1112,6 +1194,7 @@ end
 function simple_z_advection!(advection_term, pdf, vth, z, vpa, dt_max_in)
     dt_max = dt_max_in
     # take the z derivative of the input pdf
+    begin_r_vperp_vpa_region()
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
         speed = vth[iz,ir] * vpa[ivpa]
         dt_max = min(dt_max, 0.5*z.cell_width[iz]/(max(abs(speed),1e-3)))
@@ -1151,6 +1234,7 @@ end
 function simple_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_max_in)
     dt_max = dt_max_in
     # take the vpa derivative in the input pdf
+    begin_r_z_vperp_region()
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
         speed = ((vth[iz,ir] * dppar_dz[iz,ir] + vpa.grid[ivpa] * dqpar_dz[iz,ir]) 
                 / (2 * ppar[iz,ir]) - vpa.grid[ivpa]^2 * dvth_dz[iz,ir])
@@ -1183,6 +1267,7 @@ end
 
 function add_source_term!(source_term, vpa, z, dvth_dz)
     # add in particle and heat source term(s)
+    begin_r_z_vperp_vpa_region()
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
     #    source_term[ivpa,ivperp,iz,ir] -= 40*exp(-vpa[ivpa]^2)*exp(-z[iz]^2)
     #    source_term[ivpa,ivperp,iz,ir] -= vpa[ivpa]*exp(-vpa[ivpa]^2)*(2*vpa[ivpa]^2-3)*dvth_dz[iz,ir]
@@ -1205,6 +1290,7 @@ function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa,
     #                         buffer_r_4, z_spectral, z)
     #    @. residual[ivpa,ivperp,:,:] -= num_diss_params.z_dissipation_coefficient * dummy_zr2
     #end
+    begin_r_z_vperp_region()
     @loop_r_z_vperp ir iz ivperp begin
         #@views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
         #@views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
@@ -1383,39 +1469,52 @@ function calculate_contribution_from_z_advection!(z_advection_term, pdf, vthe, z
                   scratch_dummy.buffer_vpavperpr_2, scratch_dummy.buffer_vpavperpr_3,
                   scratch_dummy.buffer_vpavperpr_4, spectral, z)
     # contribution from the z-advection term is wpa * vthe * d(pdf)/dz
-    @loop_vperp_vpa ivperp ivpa begin
-        @. z_advection_term[ivpa, ivperp, :, :] = z_advection_term[ivpa, ivperp, :, :] #* vpa[ivpa] * vthe[:, :]
+    begin_r_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        @. z_advection_term[ivpa,ivperp,iz,ir] = z_advection_term[ivpa,ivperp,iz,ir] #* vpa[ivpa] * vthe[:, :]
     end
     return nothing
 end
 
 function add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz,
                                                     vpa, vpa_spectral)
+    begin_r_z_vperp_region()
     @loop_r_z_vperp ir iz ivperp begin
         # calculate the wpa-derivative of the pdf and store in the scratch array vpa.scratch
         @views derivative!(vpa.scratch, pdf[:, ivperp, iz, ir], vpa, vpa_spectral)
         # contribution from the wpa-advection term is ( ) * d(pdf)/dwpa
-        @. residual[:, ivperp, iz, ir] += vpa.scratch[:] * (0.5 * vth[iz, ir] / ppar[iz, ir] * dppar_dz[iz, ir]
-                + 0.5 * vpa.grid[:] / ppar[iz, ir] * dqpar_dz[iz, ir] - vpa.grid[:]^2 * dvth_dz[iz, ir])
+        @. residual[:,ivperp,iz,ir] += vpa.scratch * (0.5 * vth[iz,ir] / ppar[iz,ir] * dppar_dz[iz,ir]
+                + 0.5 * vpa.grid / ppar[iz,ir] * dqpar_dz[iz,ir] - vpa.grid^2 * dvth_dz[iz,ir])
     end
     return nothing
 end
 
 # calculate the pre-factor multiplying the modified electron pdf
 function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa)
-    @loop_vperp_vpa ivperp ivpa begin
-        @. pdf_dot_prefactor[ivpa, ivperp, :, :] = 0.5 * dqpar_dz[:, :] / ppar[:, :] - vpa[ivpa] * vth[:, :] * 
-        (ddens_dz[:, :] / dens[:, :] + dvth_dz[:, :] / vth[:, :])
+    begin_r_z_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        @. pdf_dot_prefactor[ivpa,ivperp,iz,ir] = 0.5 * dqpar_dz[iz,ir] / ppar[iz,ir] - vpa[ivpa] * vth[iz,ir] * 
+        (ddens_dz[iz,ir] / dens[iz,ir] + dvth_dz[iz,ir] / vth[iz,ir])
     end
     return nothing
 end
 
 # add contribution to the residual coming from the term proporational to the pdf
 function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z)
-    @loop_vperp_vpa ivperp ivpa begin
-        @. residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :] - vpa[ivpa] * vth[:, :] * 
-                                (ddens_dz[:, :] / dens[:, :] - dvth_dz[:, :] / vth[:, :])) * pdf[ivpa, ivperp, :, :]
-        #@. residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf[ivpa, ivperp, :, :]
+    begin_r_z_vperp_vpa_region()
+    @loop_r_z ir iz begin
+        this_dqpar_dz = dqpar_dz[iz,ir]
+        this_ppar = ppar[iz,ir]
+        this_vth = vth[iz,ir]
+        this_ddens_dz = ddens_dz[iz,ir]
+        this_dens = dens[iz,ir]
+        this_dvth_dz = dvth_dz[iz,ir]
+        this_vth = vth[iz,ir]
+        @loop_vperp_vpa ivperp ivpa begin
+            residual[ivpa,ivperp,iz,ir] -= (-0.5 * this_dqpar_dz / this_ppar - vpa[ivpa] * this_vth *
+                                    (this_ddens_dz / this_dens - this_dvth_dz / this_vth)) * pdf[ivpa,ivperp,iz,ir]
+            #residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf[ivpa, ivperp, :, :]
+        end
     end
     return nothing
 end
@@ -1432,18 +1531,36 @@ end
 #     return nothing
 # end
 
-function get_electron_critical_velocities(phi, vthe, me_over_mi)
+function get_electron_critical_velocities(phi, vthe, me_over_mi, z)
     # get the critical velocities beyond which electrons are lost to the wall
-    crit_speed_zmin = sqrt(phi[1, 1] / (me_over_mi * vthe[1, 1]^2))
-    crit_speed_zmax = -sqrt(max(phi[end, 1],0.0) / (me_over_mi * vthe[end, 1]^2))
+    if z.irank == 0 && block_rank[] == 0
+        crit_speed_zmin = sqrt(phi[1, 1] / (me_over_mi * vthe[1, 1]^2))
+    else
+        crit_speed_zmin = 0.0
+    end
+    @serial_region begin
+        crit_speed_zmin = MPI.Bcast(crit_speed_zmin, 0, z.comm)
+    end
+    crit_speed_zmin = MPI.Bcast(crit_speed_zmin, 0, comm_block[])
+
+    if z.irank == z.nrank - 1 && block_rank[] == 0
+        crit_speed_zmax = -sqrt(max(phi[end, 1],0.0) / (me_over_mi * vthe[end, 1]^2))
+    else
+        crit_speed_zmin = 0.0
+    end
+    @serial_region begin
+        crit_speed_zmax = MPI.Bcast(crit_speed_zmax, z.nrank-1, z.comm)
+    end
+    crit_speed_zmax = MPI.Bcast(crit_speed_zmax, 0, comm_block[])
+
     return crit_speed_zmin, crit_speed_zmax
 end
 
-function check_electron_pdf_convergence(residual, norm, upar, vthe, vpa)
+function check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
     #average_residual = 0.0
     # @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-    #     if norm[ivpa, ivperp, iz, ir] > eps(mk_float)
-    #         average_residual += abs(residual[ivpa, ivperp, iz, ir]) / norm[ivpa, ivperp, iz, ir]
+    #     if pdf[ivpa, ivperp, iz, ir] > eps(mk_float)
+    #         average_residual += abs(residual[ivpa, ivperp, iz, ir]) / pdf[ivpa, ivperp, iz, ir]
     #     else
     #         average_residual += abs(residual[ivpa, ivperp, iz, ir])
     #     end
@@ -1454,18 +1571,30 @@ function check_electron_pdf_convergence(residual, norm, upar, vthe, vpa)
     # those that are set by the sheath boundary condition are not used by the time
     # advance, and so might not converge to 0.
     # First, sum the contributions from the bulk of the domain
-    sum_residual = sum(abs.(@view residual[:,:,2:end-1,:]))
-    # Then add the contributions from the evolved parts of the sheath boundary points
-    @loop_r ir begin
-        vpa_unnorm_lower = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
-        iv0_lower = findfirst(x -> x>0.0, vpa_unnorm_lower)
-        vpa_unnorm_upper = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
-        iv0_upper = findlast(x -> x>0.0, vpa_unnorm_upper)
-        sum_residual += sum(abs.(@view residual[1:iv0_lower-1,:,1,ir])) +
-                        sum(abs.(@view residual[iv0_upper+1:end,:,end,ir]))
+    begin_r_z_vperp_region()
+    sum_residual = 0.0
+    sum_pdf = 0.0
+    @loop_r_z ir iz begin
+        if z.irank == 0 && iz == 1
+            vpa_unnorm_lower = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
+            iv0_end = findfirst(x -> x>0.0, vpa_unnorm_lower) - 1
+        else
+            iv0_end = vpa.n
+        end
+        if z.irank == z.nrank-1 && iz == z.n
+            vpa_unnorm_upper = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
+            iv0_start = findlast(x -> x>0.0, vpa_unnorm_upper) + 1
+        else
+            iv0_start = 1
+        end
+        @loop_vperp ivperp begin
+            sum_residual += sum(abs.(@view residual[iv0_start:iv0_end,ivperp,iz,ir]))
+            sum_pdf += sum(abs.(@view pdf[iv0_start:iv0_end,ivperp,iz,ir]))
+        end
     end
+    sum_residual, sum_pdf = MPI.Allreduce([sum_residual, sum_pdf], +, comm_world)
 
-    average_residual = sum_residual / sum(norm)
+    average_residual = sum_residual / sum_pdf
 
     electron_pdf_converged = (average_residual < 1e-3)
 
@@ -1473,7 +1602,13 @@ function check_electron_pdf_convergence(residual, norm, upar, vthe, vpa)
 end
 
 function check_electron_pdf_convergence(residual)
-    average_residual = sum(abs.(residual)) / length(residual)
+    begin_r_z_vperp_region()
+    sum_residual = 0.0
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        sum_residual += abs.(residual[ivpa,ivperp,iz,ir])
+    end
+    sum_residual, sum_length = MPI.Allreduce((sum_residual, length(residual) / block_size[]), +, comm_world)
+    average_residual = sum_residual / sum_length
     electron_pdf_converged = (average_residual < 1e-3)
     return average_residual, electron_pdf_converged
 end
diff --git a/src/electron_vpa_advection.jl b/src/electron_vpa_advection.jl
index 3d60b8945..8aa8c5de1 100644
--- a/src/electron_vpa_advection.jl
+++ b/src/electron_vpa_advection.jl
@@ -19,10 +19,12 @@ function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar
     d2pdf_dvpa2 = scratch_dummy.buffer_vpavperpzr_2
     begin_r_z_vperp_region()
     # get the updated speed along the wpa direction using the current pdf
-    @views update_electron_speed_vpa!(advect[1], ppar[:,:], vth[:,:], dppar_dz[:,:], dqpar_dz[:,:], dvth_dz[:,:], vpa.grid)
+    @views update_electron_speed_vpa!(advect[1], ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa.grid)
     # update adv_fac -- note that there is no factor of dt here because
     # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
-    @views @. advect[1].adv_fac[:,:,:,:] = -advect[1].speed[:,:,:,:]
+    @loop_r_z_vperp ir iz ivperp begin
+        @views @. advect[1].adv_fac[:,ivperp,iz,ir] = -advect[1].speed[:,ivperp,iz,ir]
+    end
     #calculate the upwind derivative of the electron pdf w.r.t. wpa
     @loop_r_z_vperp ir iz ivperp begin
         @views derivative!(dpdf_dvpa[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, advect[1].adv_fac[:,ivperp,iz,ir], spectral)
@@ -31,9 +33,9 @@ function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar
     #    @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, spectral)
     #end
     # calculate the advection term
-    @loop_vpa ivpa begin
-        @. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:]
-        #@. advection_term[ivpa,:,:,:] -= advect[1].adv_fac[ivpa,:,:,:] * dpdf_dvpa[ivpa,:,:,:] + 0.0001*d2pdf_dvpa2[ivpa,:,:,:]
+    @loop_r_z_vperp ir iz ivperp begin
+        @. advection_term[:,ivperp,iz,ir] -= advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir]
+        #@. advection_term[:,ivperp,iz,ir] -= advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir] + 0.0001*d2pdf_dvpa2[:,ivperp,iz,ir]
     end
     #@loop_vpa ivpa begin
     #    println("electron_vpa_advection: ", advection_term[ivpa,1,10,1], " vpa: ", vpa.grid[ivpa], " dpdf_dvpa: ", dpdf_dvpa[ivpa,1,10,1],
diff --git a/src/electron_z_advection.jl b/src/electron_z_advection.jl
index e9cf1e43f..f987ce028 100644
--- a/src/electron_z_advection.jl
+++ b/src/electron_z_advection.jl
@@ -20,10 +20,12 @@ function electron_z_advection!(advection_term, pdf, vth, advect, z, vpa, spectra
     d2pdf_dz2 = scratch_dummy.buffer_vpavperpzr_2
     begin_r_vperp_vpa_region()
     # get the updated speed along the z direction using the current pdf
-    @views update_electron_speed_z!(advect[1], vth[:,:], vpa)
+    @views update_electron_speed_z!(advect[1], vth, vpa)
     # update adv_fac -- note that there is no factor of dt here because
     # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
-    @views @. advect[1].adv_fac[:,:,:,:] = -advect[1].speed[:,:,:,:]
+    @loop_r_vperp_vpa ir ivperp ivpa begin
+        @views advect[1].adv_fac[:,ivpa,ivperp,ir] = -advect[1].speed[:,ivpa,ivperp,ir]
+    end
     #calculate the upwind derivative
     derivative_z!(dpdf_dz, pdf,
                   advect, scratch_dummy.buffer_vpavperpr_1,
@@ -34,9 +36,9 @@ function electron_z_advection!(advection_term, pdf, vth, advect, z, vpa, spectra
     #    @views second_derivative!(d2pdf_dz2[ivpa,ivperp,:,ir], pdf[ivpa,ivperp,:,ir], z, spectral)
     #end
     # calculate the advection term
-    @loop_z iz begin
-        @. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:]
-        #@. advection_term[:,:,iz,:] -= advect[1].adv_fac[iz,:,:,:] * dpdf_dz[:,:,iz,:] + 0.0001*d2pdf_dz2[:,:,iz,:]
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        advection_term[ivpa,ivperp,iz,ir] -= advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir]
+        #advection_term[ivpa,ivperp,iz,ir] -= advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir] + 0.0001*d2pdf_dz2[ivpa,ivperp,iz,ir]
     end
     return nothing
 end
diff --git a/src/em_fields.jl b/src/em_fields.jl
index e7cb0d63d..91912cc2c 100644
--- a/src/em_fields.jl
+++ b/src/em_fields.jl
@@ -69,13 +69,17 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
                            # for each radial position in parallel if possible 
         
     # TODO: parallelise this...
-    @serial_region begin
-        ne = @view scratch_dummy.dummy_zrs[:,:,1]
-        jpar_i = @view scratch_dummy.buffer_rs_1[:,:,1]
-        N_e = @view scratch_dummy.buffer_rs_2[:,:,1]
-        if composition.electron_physics == boltzmann_electron_response
+    ne = @view scratch_dummy.dummy_zrs[:,:,1]
+    jpar_i = @view scratch_dummy.buffer_rs_1[:,:,1]
+    N_e = @view scratch_dummy.buffer_rs_2[:,:,1]
+    if composition.electron_physics == boltzmann_electron_response
+        begin_serial_region()
+        @serial_region begin
             N_e .= 1.0
-        elseif composition.electron_physics == boltzmann_electron_response_with_simple_sheath
+        end
+    elseif composition.electron_physics == boltzmann_electron_response_with_simple_sheath
+        begin_serial_region()
+        @serial_region begin
             # calculate Sum_{i} Z_i n_i u_i = J_||i at z = 0
             jpar_i .= 0.0
             for is in 1:n_ion_species
@@ -98,17 +102,17 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
             end
             # See P.C. Stangeby, The Plasma Boundary of Magnetic Fusion Devices, IOP Publishing, Chpt 2, p75
         end
-        if composition.electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
-            @loop_r_z ir iz begin
-                fields.phi[iz,ir] = composition.T_e * log(dens_e[iz,ir] / N_e[ir])
-            end
-        elseif composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
-            calculate_Epar_from_electron_force_balance!(fields.Ez, dens_e, moments.electron.dppar_dz,
-                collisions.nu_ei, moments.electron.parallel_friction,
-                composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
-                fvec.density_neutral, fvec.uz_neutral, fvec.electron_upar)
-            calculate_phi_from_Epar!(fields.phi, fields.Ez, z.cell_width)
+    end
+    if composition.electron_physics ∈ (boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath)
+        @loop_r_z ir iz begin
+            fields.phi[iz,ir] = composition.T_e * log(dens_e[iz,ir] / N_e[ir])
         end
+    elseif composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        calculate_Epar_from_electron_force_balance!(fields.Ez, dens_e, moments.electron.dppar_dz,
+            collisions.nu_ei, moments.electron.parallel_friction,
+            composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
+            fvec.density_neutral, fvec.uz_neutral, fvec.electron_upar)
+        calculate_phi_from_Epar!(fields.phi, fields.Ez, z)
     end
     ## can calculate phi at z = L and hence phi_wall(z=L) using jpar_i at z =L if needed
     _block_synchronize()
@@ -145,12 +149,29 @@ function update_phi!(fields, fvec, z, r, composition, collisions, moments,
     return nothing
 end
 
-function calculate_phi_from_Epar!(phi, Epar, dz)
+function calculate_phi_from_Epar!(phi, Epar, z)
     # simple calculation of phi from Epar for now, with zero phi assumed at boundary
-    phi[1,:] .= 3.0
-    @loop_r_z ir iz begin
-        if iz > 1
-            phi[iz,ir] = phi[iz-1,ir] - dz[iz-1]*Epar[iz,ir]
+    begin_serial_region()
+
+    dz = z.cell_width
+    @serial_region begin
+        phi[1,:] .= 3.0
+        @loop_r_z ir iz begin
+            if iz > 1
+                phi[iz,ir] = phi[iz-1,ir] - dz[iz-1]*Epar[iz,ir]
+            end
+        end
+
+        # Add contributions to integral along z from processes at smaller z-values than
+        # this one.
+        this_delta_phi = phi[end,:] .- phi[1,:]
+        for irank ∈ 1:z.nrank-1
+            delta_phi = MPI.Bcast(this_delta_phi, irank, z.comm)
+            if z.irank > irank
+                @loop_r_z ir iz begin
+                    phi[iz,ir] += delta_phi[ir]
+                end
+            end
         end
     end
     return nothing
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index dffe8ffbe..84b568bdb 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -300,61 +300,63 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
-    @serial_region begin 
-        moments.electron.dens_updated = false
-        # initialise the electron density profile
-        init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
-        # initialise the electron parallel flow profile
-        init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
-            moments.ion.upar, moments.ion.dens, composition.electron_physics)
-        # initialise the electron thermal speed profile
-        init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
-        # calculate the electron temperature from the thermal speed
-        @. moments.electron.temp = composition.me_over_mi * moments.electron.vth^2
-        # the electron temperature has now been updated
-        moments.electron.temp_updated = true
-        # calculate the electron parallel pressure from the density and temperature
-        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
-        # the electron parallel pressure now been updated
-        moments.electron.ppar_updated = true
-
-        # calculate the zed derivative of the initial electron density
-        @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
-            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # calculate the zed derivative of the initial electron temperature
-        @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
-            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # calculate the zed derivative of the initial electron thermal speed
-        @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, 
-            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # calculate the zed derivative of the initial electron parallel pressure
-        @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
-            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # calculate the electron parallel heat flux;
-        # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
-            moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
-            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
-        # calculate the zed derivative of the initial electron parallel heat flux
-        @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
-            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # calculate the electron-ion parallel friction force
-        calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
-            moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
-            composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
-        
-        # initialize the scratch arrays containing pdfs and moments for the first RK stage
-        # the electron pdf is yet to be initialised but with the current code logic, the scratch
-        # arrays need to exist and be otherwise initialised in order to compute the initial electron pdf
-        initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
-        # get the initial electrostatic potential and parallel electric field
-        update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+    moments.electron.dens_updated = false
+    # initialise the electron density profile
+    init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
+    # initialise the electron parallel flow profile
+    init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
+        moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
+    # initialise the electron thermal speed profile
+    init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
+    # calculate the electron temperature from the thermal speed
+    @loop_r_z ir iz begin
+        moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
+    end
+    # the electron temperature has now been updated
+    moments.electron.temp_updated = true
+    # calculate the electron parallel pressure from the density and temperature
+    @loop_r_z ir iz begin
+        moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
     end
+    # the electron parallel pressure now been updated
+    moments.electron.ppar_updated = true
+
+    # calculate the zed derivative of the initial electron density
+    @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
+        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # calculate the zed derivative of the initial electron temperature
+    @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
+        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # calculate the zed derivative of the initial electron thermal speed
+    @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, 
+        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # calculate the zed derivative of the initial electron parallel pressure
+    @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
+        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # calculate the electron parallel heat flux;
+    # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
+    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+        moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
+        collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+    # calculate the zed derivative of the initial electron parallel heat flux
+    @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
+        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # calculate the electron-ion parallel friction force
+    calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+        moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+        composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
+    
+    # initialize the scratch arrays containing pdfs and moments for the first RK stage
+    # the electron pdf is yet to be initialised but with the current code logic, the scratch
+    # arrays need to exist and be otherwise initialised in order to compute the initial electron pdf
+    initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    # get the initial electrostatic potential and parallel electric field
+    update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
 
     # initialize the electron pdf that satisfies the electron kinetic equation
     return initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, z, vpa, vperp, z_spectral, vpa_spectral, 
@@ -373,8 +375,9 @@ that may be evolved separately via fluid equations
 """
 function initialize_scratch_arrays!(scratch, moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
     # populate each of the structs
-    for istage ∈ 1:n_rk_stages+1
-        @serial_region begin
+    begin_serial_region()
+    @serial_region begin
+        for istage ∈ 1:n_rk_stages+1
             # initialise the scratch arrays for the ion pdf and velocity moments
             scratch[istage].pdf .= pdf_ion_in
             scratch[istage].density .= moments.ion.dens
@@ -484,7 +487,6 @@ end
 
 function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spectral, vpa_spectral, z_advect, vpa_advect,
                                   scratch_dummy, collisions, composition, num_diss_params, dt)
-    @serial_region begin
         @loop_r ir begin
             # this is the initial guess for the electron pdf
             # it will be iteratively updated to satisfy the time-independent
@@ -492,33 +494,36 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
             @views init_electron_pdf_over_density!(pdf.electron.norm[:,:,:,ir], moments.electron.dens[:,ir],
                 moments.electron.upar[:,ir], moments.electron.vth[:,ir], phi[:,ir], z, vpa, vperp, composition.me_over_mi)
         end
-        # update the electron pdf in the fvec struct
-        fvec.pdf_electron .= pdf.electron.norm
-        # now that the initial electron pdf is given, the electron parallel heat flux should be updated
-        # if using kinetic electrons
-        if composition.electron_physics == kinetic_electrons
-            moments.electron.qpar_updated = false
-            calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
-                moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
-                moments.electron.dT_dz, moments.ion.upar, 
-                collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
-            # update dqpar/dz for electrons
-            # calculate the zed derivative of the initial electron parallel heat flux
-            @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
-                scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-                scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-            # now that we have our initial guess for the electron pdf, we iterate
-            # using the time-independent electron kinetic equation to find a self-consistent
-            # solution for the electron pdf
-            max_electron_pdf_iterations = 500000
-            #max_electron_pdf_iterations = 10000
-            return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
-                                        moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
-                                        phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
-                                        moments.electron.dqpar_dz, moments.electron.dvth_dz, z, vpa, z_spectral, 
-                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
-                                        num_diss_params, max_electron_pdf_iterations)
+    # now that the initial electron pdf is given, the electron parallel heat flux should be updated
+    # if using kinetic electrons
+    if composition.electron_physics == kinetic_electrons
+        begin_serial_region()
+        @serial_region begin
+            # update the electron pdf in the fvec struct
+            fvec.pdf_electron .= pdf.electron.norm
         end
+
+        moments.electron.qpar_updated = false
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+            moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
+            moments.electron.dT_dz, moments.ion.upar, 
+            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+        # update dqpar/dz for electrons
+        # calculate the zed derivative of the initial electron parallel heat flux
+        @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
+            scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+            scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+        # now that we have our initial guess for the electron pdf, we iterate
+        # using the time-independent electron kinetic equation to find a self-consistent
+        # solution for the electron pdf
+        max_electron_pdf_iterations = 500000
+        #max_electron_pdf_iterations = 10000
+        return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
+                                    moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
+                                    phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
+                                    moments.electron.dqpar_dz, moments.electron.dvth_dz, z, vpa, z_spectral, 
+                                    vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
+                                    num_diss_params, max_electron_pdf_iterations)
     end
 end
 
@@ -765,8 +770,8 @@ end
 """
 initialise the electron parallel flow density
 """
-function init_electron_upar!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
-    calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model)
+function init_electron_upar!(upar_e, updated, dens_e, upar_i, dens_i, electron_model, r, z)
+    calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model, r, z)
     return nothing
 end
 
@@ -1177,7 +1182,7 @@ this 'initital' value for the electron will just be the first guess in an iterat
 function init_electron_pdf_over_density!(pdf, density, upar, vth, phi, z, vpa, vperp, me_over_mi)
     if z.bc == "wall"
         # get critical velocities beyond which electrons are lost to the wall
-        vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi)
+        vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
         println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
         # loop over all z values on this rank, initialising a shifted Maxwellian velocity distribution
         sharp_fac = 10
diff --git a/src/time_advance.jl b/src/time_advance.jl
index d7f079a03..64ec5dae4 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -2155,7 +2155,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
     if advance.electron_energy
         electron_energy_equation!(fvec_out.electron_ppar, fvec_out.density, fvec_in, moments, collisions, dt,
-                                  composition, num_diss_params, z.grid)
+                                  composition, num_diss_params, z)
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not

From 390b32c2a5cf554d084738fd0ba17d9e8c425087 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 10 Feb 2024 17:30:44 +0000
Subject: [PATCH 080/428] Don't use hard_force_moment_constraints() at sheath
 entrance points

---
 src/electron_kinetic_equation.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 2f7baff64..1bbd8eee1 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -266,6 +266,9 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
         begin_r_z_region()
         @loop_r_z ir iz begin
+            if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
+                continue
+            end
             #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
             @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
         end

From a8c5c2ce86c77e2e8c10ccb725323fc3bae54ec0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 10 Feb 2024 18:15:39 +0000
Subject: [PATCH 081/428] Fix calculation of vmax and vmin

---
 src/electron_kinetic_equation.jl | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 1bbd8eee1..db23908c3 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -530,8 +530,16 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         integral_excess = upar_integral - upar0
         fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
         #println("fraction_of_pdf=", fraction_of_pdf)
-        vmax = 0.5*(vpa_unnorm[ivpa+1] + vpa_unnorm[ivpa]) +
-               fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa+1])
+
+        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+        # integral from the point at ivpa is required to make upar_integral<upar0) the
+        # cut-off velocity is half way between ivpa and ivpa+1, while when
+        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+        # ivpa is required to make upar_integral<upar0) the cut-off is half way between
+        # ivpa-1 and ivpa.
+        vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
+               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
+
         #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
         wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
         #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
@@ -742,8 +750,16 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         integral_excess = upar_integral - upar_end
         fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
         #println("B fraction_of_pdf=", fraction_of_pdf)
-        vmin = 0.5*(vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
-               fraction_of_pdf*(vpa_unnorm[ivpa] - vpa_unnorm[ivpa-1])
+
+        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+        # integral from the point at ivpa is required to make upar_integral>upar_end) the
+        # cut-off velocity is half way between ivpa-1 and ivpa, while when
+        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+        # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
+        # ivpa and ivpa+1.
+        vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
+               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
+
         #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
         wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
         @loop_vpa ivpa begin

From fffd9f1ca42a465d5e93c141723f36ea4a98d364 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 11 Feb 2024 23:47:26 +0000
Subject: [PATCH 082/428] Fix use of *_updated fields in
 moments_electron_substruct

These need to be of type Ref{Bool} so that they can be passed into
functions and have their value changed. Also means the `struct
moments_electron_substruct` does not need to be mutable.
---
 src/electron_fluid_equations.jl | 12 ++++++------
 src/initial_conditions.jl       |  8 ++++----
 src/velocity_moments.jl         | 22 +++++++++++-----------
 3 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/src/electron_fluid_equations.jl b/src/electron_fluid_equations.jl
index 9d3064dda..8071e773f 100644
--- a/src/electron_fluid_equations.jl
+++ b/src/electron_fluid_equations.jl
@@ -29,7 +29,7 @@ output:
 """
 function calculate_electron_density!(dens_e, updated, dens_i)
     # only update the electron density if it has not already been updated
-    if !updated
+    if !updated[]
         begin_r_z_region()
         # enforce quasineutrality
         @loop_r_z ir iz begin
@@ -40,7 +40,7 @@ function calculate_electron_density!(dens_e, updated, dens_i)
         end
     end
     # set flag indicating that the electron density has been updated
-    updated = true
+    updated[] = true
     return nothing
 end
 
@@ -59,7 +59,7 @@ output:
 """
 function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens_e, upar_i, dens_i, electron_model, r, z)
     # only calculate the electron parallel flow if it is not already updated
-    if !updated
+    if !updated[]
         begin_r_region()
         # get the number of zed grid points, nz
         nz = size(upar_e,1)
@@ -105,8 +105,8 @@ function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens
                 upar_e[iz,ir] /= dens_e[iz,ir]
             end
         end
+        updated[] = true
     end
-    updated = true
     return nothing
 end
 
@@ -280,7 +280,7 @@ output:
 function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth_e, dTe_dz, upar_i, 
                                   nu_ei, me_over_mi, electron_model, vpa)
     # only calculate qpar_e if needs updating
-    if qpar_updated == false
+    if !qpar_updated[]
         if electron_model == braginskii_fluid
             begin_r_z_region()
             # use the classical Braginskii expression for the electron heat flux
@@ -301,7 +301,7 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth
         end
     end
     # qpar has been updated
-    qpar_updated = true
+    qpar_updated[] = true
     return nothing
 end
 
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 84b568bdb..110281155 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -300,7 +300,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
-    moments.electron.dens_updated = false
+    moments.electron.dens_updated[] = false
     # initialise the electron density profile
     init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
     # initialise the electron parallel flow profile
@@ -313,13 +313,13 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
         moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
     end
     # the electron temperature has now been updated
-    moments.electron.temp_updated = true
+    moments.electron.temp_updated[] = true
     # calculate the electron parallel pressure from the density and temperature
     @loop_r_z ir iz begin
         moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
     end
     # the electron parallel pressure now been updated
-    moments.electron.ppar_updated = true
+    moments.electron.ppar_updated[] = true
 
     # calculate the zed derivative of the initial electron density
     @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
@@ -503,7 +503,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
             fvec.pdf_electron .= pdf.electron.norm
         end
 
-        moments.electron.qpar_updated = false
+        moments.electron.qpar_updated[] = false
         calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
             moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
             moments.electron.dT_dz, moments.ion.upar, 
diff --git a/src/velocity_moments.jl b/src/velocity_moments.jl
index 8de992fda..968966a31 100644
--- a/src/velocity_moments.jl
+++ b/src/velocity_moments.jl
@@ -132,27 +132,27 @@ end
 """
 moments_electron_substruct is a struct that contains moment information for electrons
 """
-mutable struct moments_electron_substruct
+struct moments_electron_substruct
     # this is the particle density
     dens::MPISharedArray{mk_float,2}
     # flag that keeps track of if the density needs updating before use
-    dens_updated::Bool
+    dens_updated::Ref{Bool}
     # this is the parallel flow
     upar::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not upar needs updating before use
-    upar_updated::Bool
+    upar_updated::Ref{Bool}
     # this is the parallel pressure
     ppar::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not ppar needs updating before use
-    ppar_updated::Bool
+    ppar_updated::Ref{Bool}
     # this is the temperature
     temp::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not temp needs updating before use
-    temp_updated::Bool
+    temp_updated::Ref{Bool}
     # this is the parallel heat flux
     qpar::MPISharedArray{mk_float,2}
     # flag that keeps track of whether or not qpar needs updating before use
-    qpar_updated::Bool
+    qpar_updated::Ref{Bool}
     # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
     vth::MPISharedArray{mk_float,2}
     # this is the parallel friction force between ions and electrons
@@ -411,23 +411,23 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     # allocate array used for the particle density
     density = allocate_shared_float(nz, nr)
     # initialise Bool variable that indicates if the density is updated for each species
-    density_updated = false
+    density_updated = Ref(false)
     # allocate array used for the parallel flow
     parallel_flow = allocate_shared_float(nz, nr)
     # allocate Bool variable that indicates if the parallel flow is updated for each species
-    parallel_flow_updated = false
+    parallel_flow_updated = Ref(false)
     # allocate array used for the parallel pressure
     parallel_pressure = allocate_shared_float(nz, nr)
     # allocate Bool variable that indicates if the parallel pressure is updated for each species
-    parallel_pressure_updated = false
+    parallel_pressure_updated = Ref(false)
     # allocate array used for the temperature
     temperature = allocate_shared_float(nz, nr)
     # allocate Bool variable that indicates if the temperature is updated for each species
-    temperature_updated = false
+    temperature_updated = Ref(false)
     # allocate array used for the parallel flow
     parallel_heat_flux = allocate_shared_float(nz, nr)
     # allocate Bool variables that indicates if the parallel flow is updated for each species
-    parallel_heat_flux_updated = false
+    parallel_heat_flux_updated = Ref(false)
     # allocate array used for the election-ion parallel friction force
     parallel_friction_force = allocate_shared_float(nz, nr)
     # allocate array used for electron heat source

From 3d75f1e56b0f27fd44f587f47a0033626f7e508a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 11 Feb 2024 23:51:22 +0000
Subject: [PATCH 083/428] Recalculate dqpar_dz when subcycling ppar

---
 src/electron_kinetic_equation.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index db23908c3..10ea88528 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -353,6 +353,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 qpar[iz,ir] = vthe[iz,ir]^3 * wpa3_moment[iz,ir]
                 dummy_zr[iz,ir] = -upar[iz,ir]
             end
+            @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
+                                 buffer_r_4, z_spectral, z)
             # Compute the upwinded z-derivative of the electron parallel pressure for the 
             # electron energy equation
             @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,

From 1306dc0fe6f01f5cde06adf29fdef22d06faea81 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 11 Feb 2024 23:53:33 +0000
Subject: [PATCH 084/428] Fix initialisation of electron pdf

Use upar_e=upar_i and the usual electron boundary condition function to
apply a boundary condition to a Maxwellian pdf at the sheath entrance
boundaries. Then blend the boundary values into a Maxwellian in the bulk
of the domain.
---
 src/initial_conditions.jl | 90 ++++++++++++++++++++++++---------------
 1 file changed, 55 insertions(+), 35 deletions(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 110281155..cb6c51571 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -38,7 +38,7 @@ using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
-using ..electron_kinetic_equation: update_electron_pdf!, get_electron_critical_velocities
+using ..electron_kinetic_equation: update_electron_pdf!, enforce_boundary_condition_on_electron_pdf!
 using ..input_structs: boltzmann_electron_response_with_simple_sheath, kinetic_electrons
 using ..derivatives: derivative_z!
 
@@ -337,6 +337,10 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
         scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
         scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # Initialise the array for the electron pd
+    init_electron_pdf_over_density_and_boundary_phi!(
+        pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
+        moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
     # calculate the electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
     calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
@@ -487,13 +491,6 @@ end
 
 function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spectral, vpa_spectral, z_advect, vpa_advect,
                                   scratch_dummy, collisions, composition, num_diss_params, dt)
-        @loop_r ir begin
-            # this is the initial guess for the electron pdf
-            # it will be iteratively updated to satisfy the time-independent
-            # electron kinetic equation
-            @views init_electron_pdf_over_density!(pdf.electron.norm[:,:,:,ir], moments.electron.dens[:,ir],
-                moments.electron.upar[:,ir], moments.electron.vth[:,ir], phi[:,ir], z, vpa, vperp, composition.me_over_mi)
-        end
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
@@ -1174,38 +1171,61 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
 end
 
 """
-init_electron_pdf_over_density initialises the normalised electron pdf = pdf_e * vth_e / dens_e;
+init_electron_pdf_over_density_and_boundary_phi initialises the normalised electron pdf = pdf_e *
+vth_e / dens_e and the boundary values of the electrostatic potential phi;
 care is taken to ensure that the parallel boundary condition is satisfied;
 NB: as the electron pdf is obtained via a time-independent equation,
 this 'initital' value for the electron will just be the first guess in an iterative solution
 """
-function init_electron_pdf_over_density!(pdf, density, upar, vth, phi, z, vpa, vperp, me_over_mi)
+function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
+        vpa, vperp, vpa_spectral, me_over_mi)
+
     if z.bc == "wall"
-        # get critical velocities beyond which electrons are lost to the wall
-        vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
-        println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
-        # loop over all z values on this rank, initialising a shifted Maxwellian velocity distribution
-        sharp_fac = 10
-        blend_fac = 100
-        @loop_z_vperp iz ivperp begin
-            #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
-            #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
-            #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
-            @. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2) * (
-                                  (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
-                                  tanh(sharp_fac*(vpa.grid[:]-vpa_crit_zmin))) *
-                                  (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
-                                  tanh(-sharp_fac*(vpa.grid[:]-vpa_crit_zmax)))) #/
-                                  #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
-                                  #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
-#                                   exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
-                                   #                                      exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
-
-            # ensure that the normalised electron pdf integrates to unity
-            norm_factor = integrate_over_vspace(pdf[:,ivperp,iz], vpa.wgts)
-            @. pdf[:,ivperp,iz] /= norm_factor
-            #println("TMP FOR TESTING -- init electron pdf")
-            #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
+        @loop_r ir begin
+            # Initialise an unshifted Maxwellian as a first step
+            @loop_z iz begin
+                vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
+                @loop_vperp ivperp begin
+                    @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
+                end
+            end
+            # Apply the sheath boundary condition to get cut-off boundary distribution
+            # functions and boundary values of phi
+            enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vpa, vpa_spectral, me_over_mi)
+            # get critical velocities beyond which electrons are lost to the wall
+            #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
+            #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
+            # Blend boundary distribution function into bulk of domain to avoid
+            # discontinuities (as much as possible)
+            blend_fac = 100
+            if z.nrank > 1
+                error("Distributed MPI not supported in this init yet")
+            end
+            @loop_z_vperp iz ivperp begin
+                #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
+                #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
+                #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
+                blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
+                blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
+                @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
+                                         blend_fac_lower * pdf[:,ivperp,1,ir] +
+                                         blend_fac_upper * pdf[:,ivperp,end,ir]
+                #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
+                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
+                #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *
+                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
+                #                          tanh(-sharp_fac*(vpa.grid-vpa_crit_zmax)))) #/
+                                         #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
+                                         #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
+                                         #exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
+                                         #exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
+
+                # ensure that the normalised electron pdf integrates to unity
+                norm_factor = integrate_over_vspace(pdf[:,ivperp,iz,ir], vpa.wgts)
+                @. pdf[:,ivperp,iz,ir] /= norm_factor
+                #println("TMP FOR TESTING -- init electron pdf")
+                #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
+            end
         end
     else
         println("!!! currently, only the wall BC is supported for kinetic electrons !!!")

From dc3aa05bce67ce3afb519b0c540a021d4869c2d0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 14:58:12 +0000
Subject: [PATCH 085/428] Use z-diffusion of electron ppar to help reach
 initial steady state

---
 examples/kinetic-electrons/wall+sheath-bc_kinetic.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index e5a51d1f2..faad6f68c 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -81,5 +81,6 @@ ascii_output = true
 
 [numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
+moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
 vpa_dissipation_coefficient = 0.2

From 89cde9b48e887ec9a1324da13b37cceb5563af56 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 14:59:48 +0000
Subject: [PATCH 086/428] Adjust electron pseudo-timestep size to get stable
 initialization run

---
 src/electron_kinetic_equation.jl | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 10ea88528..359087839 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -153,10 +153,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                          buffer_r_4, z_spectral, z)
 
     #dt_electron = dt * sqrt(composition.me_over_mi)
-    dt_max = 3.0e-8 #1.0
-    #dt_max = 2.5e-9 #1.0
-    #dt_energy = 0.001
-    dt_energy = 1.0e-7
+    #dt_max = 3.0e-8
+    dt_max = 1.0e-8
+    #dt_max = 2.5e-9
+
+    #dt_energy = 1.0e-7
+    dt_energy = 1.0e-8
+    #dt_energy = 2.5e-9
     #n_ppar_subcycles = 1000
     #n_ppar_subcycles = 200
     n_ppar_subcycles = 1

From b7f31c27d1172436944aa01376b326c5f732ff14 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 15:00:58 +0000
Subject: [PATCH 087/428] Apply z-speedup to pressure equation as well as g_e
 equation

---
 src/electron_kinetic_equation.jl | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 359087839..8d2e18fd1 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -165,6 +165,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     n_ppar_subcycles = 1
     time = 0.0
 
+    #z_speedup_fac = 20.0
+    #z_speedup_fac = 5.0
+    z_speedup_fac = 1.0
+
     # define residual to point to a dummy array;
     # to be filled with the contributions to the electron kinetic equation; i.e., d(pdf)/dt = -residual
     residual = scratch_dummy.buffer_vpavperpzr_6
@@ -189,7 +193,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         zval = z.grid[iz]
         znorm = 2.0*zval/Lz
         residual[ivpa,ivperp,iz,ir] *=
-            (1.0 + 20.0*(1.0 - znorm^2)) /
+            (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
             sqrt(1.0 + vpa.grid[ivpa]^2)
     end
     if n_blocks[] == 1
@@ -371,6 +375,16 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
             #dt_energy = dt_electron
             electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
+
+            # Apply same 'speed up' hack to ppar that we do to the distribution function,
+            # but without the wpa dependence.
+            @loop_r_z ir iz begin
+                zval = z.grid[iz]
+                znorm = 2.0*zval/Lz
+                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
+                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
+                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
+            end
             begin_r_z_region()
             @loop_r_z ir iz begin
                 fvec.electron_ppar[iz,ir] = ppar[iz,ir]
@@ -420,7 +434,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             zval = z.grid[iz]
             znorm = 2.0*zval/Lz
             residual[ivpa,ivperp,iz,ir] *=
-                (1.0 + 20.0*(1.0 - znorm^2)) /
+                (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
                 sqrt(1.0 + vpa.grid[ivpa]^2)
         end
         if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))

From 00915c45b214d4070e8baa06dbdfcceaa9b54cbd Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 15:01:44 +0000
Subject: [PATCH 088/428] Remove negative values from electron pdf

...by just resetting them to zero.
---
 src/electron_kinetic_equation.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 8d2e18fd1..96310b2a2 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -264,13 +264,15 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             pdf[ivpa,ivperp,iz,ir] -= dt_electron * residual[ivpa,ivperp,iz,ir]
         end
 
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            pdf[ivpa,ivperp,iz,ir] = max(pdf[ivpa,ivperp,iz,ir], 0.0)
+        end
+
         # enforce the boundary condition(s) on the electron pdf
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
 
-        #pdf = max.(pdf, 0.0)
-
         begin_r_z_region()
         @loop_r_z ir iz begin
             if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)

From 611e848e1ffc21047e5fb5c5dfdf697d06e59cf0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 15:04:14 +0000
Subject: [PATCH 089/428] DEBUGGING! Improve temporary debugging output

---
 src/electron_kinetic_equation.jl | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 96310b2a2..e7e9ee2ba 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -324,7 +324,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         if (mod(iteration,100) == 0)
             begin_serial_region()
             @serial_region begin
-                println("time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[1,[1,end]], " average_residual: ", average_residual)
+                println("iteration: ", iteration, " time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[[1,end],1], " average_residual: ", average_residual)
             end
         end
         if (mod(iteration,output_interval) == 0)
@@ -346,6 +346,17 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
                 result_phi[:,iteration÷output_interval+1] .= phi[:,1]
             end
+        else
+            begin_serial_region()
+            @serial_region begin
+                if iteration÷output_interval+1+iteration%output_interval < size(result_pdf, 3)
+                    result_pdf[:,:,iteration÷output_interval+1+iteration%output_interval] .= pdf[:,1,:,1]
+                    result_ppar[:,iteration÷output_interval+1+iteration%output_interval] .= ppar[:,1]
+                    result_vth[:,iteration÷output_interval+1+iteration%output_interval] .= vthe[:,1]
+                    result_qpar[:,iteration÷output_interval+1+iteration%output_interval] .= qpar[:,1]
+                    result_phi[:,iteration÷output_interval+1+iteration%output_interval] .= phi[:,1]
+                end
+            end
         end
 
         #dt_energy = dt_electron * 10.0
@@ -463,6 +474,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         close(io_vth)
         close(io_pdf)
         close(io_pdf_stages)
+
+        result_pdf = result_pdf[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
+        result_ppar = result_ppar[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
+        result_vth = result_vth[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
+        result_qpar = result_qpar[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
+        result_phi = result_phi[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
+        result_residual = result_residual[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
     end
     return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual
 end

From 631750e3926c997b3b40367a709fecf9fbca25ae Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 15:04:45 +0000
Subject: [PATCH 090/428] DEBUGGING! Still return output arrays if error
 occured during electron initialization

---
 src/electron_kinetic_equation.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index e7e9ee2ba..d64940872 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -257,6 +257,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         result_residual = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
     end
     # evolve (artificially) in time until the residual is less than the tolerance
+    try
     while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
         begin_r_z_region()
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
@@ -455,6 +456,9 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
         iteration += 1
     end
+    catch e
+        println("Error: $e")
+    end
     begin_serial_region()
     @serial_region begin
         if !electron_pdf_converged

From 9ca9475bfe4b89e178f0509062f9de221d74eef2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 14 Feb 2024 15:05:36 +0000
Subject: [PATCH 091/428] When using ppar z-diffusion for electrons, zero out
 end points

If z-diffusion is allowed to update the values of electron ppar at the
end points in z (the sheath entrance points), then there are not enough
boundary conditions and so the simulation becomes numerically unstable.
This instability can be avoided (without having to impose unwanted
boundary conditions) by setting the d2ppar_dz2 array to 0.0 at the end
points. The diffusion is just a numerical hack anyway, which should be
removed once a smooth enough initial condition has been reached, so this
extra hack should be OK (?).
---
 src/electron_kinetic_equation.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index d64940872..830de4492 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -385,6 +385,15 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             if num_diss_params.moment_dissipation_coefficient > 0.0
                 @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+                begin_serial_region()
+                @serial_region begin
+                    if z.irank == 0
+                        moments.electron.d2ppar_dz2[1,:] .= 0.0
+                    end
+                    if z.irank == z.nrank - 1
+                        moments.electron.d2ppar_dz2[end,:] .= 0.0
+                    end
+                end
             end
 
             #dt_energy = dt_electron

From 73488e1b5958ee69749eaacb5a726b9ab8762f0b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 16 Feb 2024 11:53:33 +0000
Subject: [PATCH 092/428] Increase maximum number of iterations for electron
 initialisation

Now that electron kinetic equation solve is parallelised, can afford to
do more iterations to get closer to steady state. 2,000,000 is still not
enough to make the solution really steady...

[skip ci]
---
 src/initial_conditions.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index cb6c51571..0d9766d55 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -513,7 +513,8 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
         # now that we have our initial guess for the electron pdf, we iterate
         # using the time-independent electron kinetic equation to find a self-consistent
         # solution for the electron pdf
-        max_electron_pdf_iterations = 500000
+        max_electron_pdf_iterations = 2000000
+        #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
         return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
                                     moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,

From f75d8152b93839948fe1cf258be67f9bd7ebe00a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 16 Feb 2024 16:21:05 +0000
Subject: [PATCH 093/428] Split I/O for ions, electrons and neutrals into
 separate functions

The separate functions are also called by a single output function so
the original call-site does not get more complicated, but the individual
functions can be called on their own if necessary.
---
 src/file_io.jl         | 890 ++++++++++++++++++++++++++---------------
 src/moment_kinetics.jl |  13 +-
 src/time_advance.jl    |  18 +-
 3 files changed, 591 insertions(+), 330 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index 85de97a15..fb433c464 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -603,263 +603,36 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_time = create_dynamic_variable!(dynamic, "time", mk_float; parallel_io=parallel_io,
                                            description="simulation time")
 
-        # io_phi is the handle referring to the electrostatic potential phi
-        io_phi = create_dynamic_variable!(dynamic, "phi", mk_float, z, r;
-                                          parallel_io=parallel_io,
-                                          description="electrostatic potential",
-                                          units="T_ref/e")
-        # io_Er is the handle for the radial component of the electric field
-        io_Er = create_dynamic_variable!(dynamic, "Er", mk_float, z, r;
-                                         parallel_io=parallel_io,
-                                         description="radial electric field",
-                                         units="T_ref/e L_ref")
-        # io_Ez is the handle for the zed component of the electric field
-        io_Ez = create_dynamic_variable!(dynamic, "Ez", mk_float, z, r;
-                                         parallel_io=parallel_io,
-                                         description="vertical electric field",
-                                         units="T_ref/e L_ref")
-
-        # io_density is the handle for the ion particle density
-        io_density = create_dynamic_variable!(dynamic, "density", mk_float, z, r;
-                                              n_ion_species=n_ion_species,
-                                              parallel_io=parallel_io,
-                                              description="ion species density",
-                                              units="n_ref")
-
-        # io_upar is the handle for the ion parallel flow density
-        io_upar = create_dynamic_variable!(dynamic, "parallel_flow", mk_float, z, r;
-                                           n_ion_species=n_ion_species,
-                                           parallel_io=parallel_io,
-                                           description="ion species parallel flow",
-                                           units="c_ref = sqrt(2*T_ref/mi)")
-
-        # io_ppar is the handle for the ion parallel pressure
-        io_ppar = create_dynamic_variable!(dynamic, "parallel_pressure", mk_float, z, r;
-                                           n_ion_species=n_ion_species,
-                                           parallel_io=parallel_io,
-                                           description="ion species parallel pressure",
-                                           units="n_ref*T_ref")
-        
-        # io_pperp is the handle for the ion parallel pressure
-        io_pperp = create_dynamic_variable!(dynamic, "perpendicular_pressure", mk_float, z, r;
-                                           n_ion_species=n_ion_species,
-                                           parallel_io=parallel_io,
-                                           description="ion species perpendicular pressure",
-                                           units="n_ref*T_ref")
-
-        # io_qpar is the handle for the ion parallel heat flux
-        io_qpar = create_dynamic_variable!(dynamic, "parallel_heat_flux", mk_float, z, r;
-                                           n_ion_species=n_ion_species,
-                                           parallel_io=parallel_io,
-                                           description="ion species parallel heat flux",
-                                           units="n_ref*T_ref*c_ref")
-
-        # io_vth is the handle for the ion thermal speed
-        io_vth = create_dynamic_variable!(dynamic, "thermal_speed", mk_float, z, r;
-                                          n_ion_species=n_ion_species,
-                                          parallel_io=parallel_io,
-                                          description="ion species thermal speed",
-                                          units="c_ref")
-
-        # io_density is the handle for the ion particle density
-        io_electron_density = create_dynamic_variable!(dynamic, "electron_density", mk_float, z, r;
-                                              parallel_io=parallel_io,
-                                              description="electron species density",
-                                              units="n_ref")
-
-        # io_electron_upar is the handle for the electron parallel flow density
-        io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
-                                           parallel_io=parallel_io,
-                                           description="electron species parallel flow",
-                                           units="c_ref = sqrt(2*T_ref/mi)")
-
-        # io_electron_ppar is the handle for the electron parallel pressure
-        io_electron_ppar = create_dynamic_variable!(dynamic, "electron_parallel_pressure", mk_float, z, r;
-                                           parallel_io=parallel_io,
-                                           description="electron species parallel pressure",
-                                           units="n_ref*T_ref")
-
-        # io_electron_qpar is the handle for the electron parallel heat flux
-        io_electron_qpar = create_dynamic_variable!(dynamic, "electron_parallel_heat_flux", mk_float, z, r;
-                                           parallel_io=parallel_io,
-                                           description="electron species parallel heat flux",
-                                           units="n_ref*T_ref*c_ref")
-
-        # io_electron_vth is the handle for the electron thermal speed
-        io_electron_vth = create_dynamic_variable!(dynamic, "electron_thermal_speed", mk_float, z, r;
-                                          parallel_io=parallel_io,
-                                          description="electron species thermal speed",
-                                          units="c_ref")
-
-        # io_dSdt is the handle for the entropy production (due to collisions)
-        io_dSdt = create_dynamic_variable!(dynamic, "entropy_production", mk_float, z, r;
-                                          n_ion_species=n_ion_species,
-                                          parallel_io=parallel_io,
-                                          description="ion species entropy production",
-                                          units="")
-
-        if parallel_io || z.irank == 0
-            # io_chodura_lower is the handle for the ion thermal speed
-            io_chodura_lower = create_dynamic_variable!(dynamic, "chodura_integral_lower", mk_float, r;
-                                              n_ion_species=n_ion_species,
-                                              parallel_io=parallel_io,
-                                              description="Generalised Chodura integral lower sheath entrance",
-                                              units="c_ref")
-        else
-            io_chodura_lower = nothing
-        end
-        if parallel_io || z.irank == z.nrank - 1
-            # io_chodura_upper is the handle for the ion thermal speed
-            io_chodura_upper = create_dynamic_variable!(dynamic, "chodura_integral_upper", mk_float, r;
-                                              n_ion_species=n_ion_species,
-                                              parallel_io=parallel_io,
-                                              description="Generalised Chodura integral upper sheath entrance",
-                                              units="c_ref")
-        else
-            io_chodura_upper = nothing
-        end
-        # io_density_neutral is the handle for the neutral particle density
-        io_density_neutral = create_dynamic_variable!(dynamic, "density_neutral", mk_float, z, r;
-                                                      n_neutral_species=n_neutral_species,
-                                                      parallel_io=parallel_io,
-                                                      description="neutral species density",
-                                                      units="n_ref")
-
-        # io_uz_neutral is the handle for the neutral z momentum density
-        io_uz_neutral = create_dynamic_variable!(dynamic, "uz_neutral", mk_float, z, r;
-                                                 n_neutral_species=n_neutral_species,
-                                                 parallel_io=parallel_io,
-                                                 description="neutral species mean z velocity",
-                                                 units="c_ref = sqrt(2*T_ref/mi)")
-
-        # io_pz_neutral is the handle for the neutral species zz pressure
-        io_pz_neutral = create_dynamic_variable!(dynamic, "pz_neutral", mk_float, z, r;
-                                                 n_neutral_species=n_neutral_species,
-                                                 parallel_io=parallel_io,
-                                                 description="neutral species mean zz pressure",
-                                                 units="n_ref*T_ref")
-
-        # io_qz_neutral is the handle for the neutral z heat flux
-        io_qz_neutral = create_dynamic_variable!(dynamic, "qz_neutral", mk_float, z, r;
-                                                 n_neutral_species=n_neutral_species,
-                                                 parallel_io=parallel_io,
-                                                 description="neutral species z heat flux",
-                                                 units="n_ref*T_ref*c_ref")
-
-        # io_thermal_speed_neutral is the handle for the neutral thermal speed
-        io_thermal_speed_neutral = create_dynamic_variable!(
-            dynamic, "thermal_speed_neutral", mk_float, z, r;
-            n_neutral_species=n_neutral_species,
-            parallel_io=parallel_io, description="neutral species thermal speed",
-            units="c_ref")
-
-        ion_source_settings = external_source_settings.ion
-        if ion_source_settings.active
-            external_source_amplitude = create_dynamic_variable!(
-                dynamic, "external_source_amplitude", mk_float, z, r;
-                parallel_io=parallel_io, description="Amplitude of the external source for ions",
-                units="n_ref/c_ref^3*c_ref/L_ref")
-            if evolve_density
-                external_source_density_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_density_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external density source for ions",
-                    units="n_ref*c_ref/L_ref")
-            else
-                external_source_density_amplitude = nothing
-            end
-            if evolve_upar
-                external_source_momentum_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_momentum_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external momentum source for ions",
-                    units="m_ref*n_ref*c_ref*c_ref/L_ref")
-            else
-                external_source_momentum_amplitude = nothing
-            end
-            if evolve_ppar
-                external_source_pressure_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_pressure_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external pressure source for ions",
-                    units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
-            else
-                external_source_pressure_amplitude = nothing
-            end
-            if ion_source_settings.PI_density_controller_I != 0.0 &&
-                    ion_source_settings.source_type ∈ ("density_profile_control", "density_midpoint_control")
-                if ion_source_settings.source_type == "density_profile_control"
-                    external_source_controller_integral = create_dynamic_variable!(
-                        dynamic, "external_source_controller_integral", mk_float, z, r;
-                        parallel_io=parallel_io,
-                        description="Integral term for the PID controller of the external source for ions")
-                else
-                    external_source_controller_integral = create_dynamic_variable!(
-                        dynamic, "external_source_controller_integral", mk_float;
-                        parallel_io=parallel_io,
-                        description="Integral term for the PID controller of the external source for ions")
-                end
-            else
-                external_source_controller_integral = nothing
-            end
-        else
-            external_source_amplitude = nothing
-            external_source_density_amplitude = nothing
-            external_source_momentum_amplitude = nothing
-            external_source_pressure_amplitude = nothing
-            external_source_controller_integral = nothing
-        end
+        io_phi, io_Er, io_Ez =
+            define_dynamic_em_field_variables!(fid, r, z, parallel_io,
+                                               external_source_settings)
+
+        io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
+        external_source_amplitude, external_source_density_amplitude,
+        external_source_momentum_amplitude, external_source_pressure_amplitude,
+        external_source_controller_integral, io_chodura_lower, io_chodura_upper =
+            define_dynamic_ion_moment_variables!(fid, n_ion_species, r, z, parallel_io,
+                                                 external_source_settings, evolve_density,
+                                                 evolve_upar, evolve_ppar)
+
+        io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
+        io_electron_vth =
+            define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
+                                                      external_source_settings,
+                                                      evolve_density, evolve_upar,
+                                                      evolve_ppar)
 
-        neutral_source_settings = external_source_settings.neutral
-        if n_neutral_species > 0 && neutral_source_settings.active
-            external_source_neutral_amplitude = create_dynamic_variable!(
-                dynamic, "external_source_neutral_amplitude", mk_float, z, r;
-                parallel_io=parallel_io, description="Amplitude of the external source for neutrals",
-                units="n_ref/c_ref^3*c_ref/L_ref")
-            if evolve_density
-                external_source_neutral_density_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_neutral_density_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external density source for neutrals",
-                    units="n_ref*c_ref/L_ref")
-            else
-                external_source_neutral_density_amplitude = nothing
-            end
-            if evolve_upar
-                external_source_neutral_momentum_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_neutral_momentum_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external momentum source for neutrals",
-                    units="m_ref*n_ref*c_ref*c_ref/L_ref")
-            else
-                external_source_neutral_momentum_amplitude = nothing
-            end
-            if evolve_ppar
-                external_source_neutral_pressure_amplitude = create_dynamic_variable!(
-                    dynamic, "external_source_neutral_pressure_amplitude", mk_float, z, r;
-                    parallel_io=parallel_io, description="Amplitude of the external pressure source for neutrals",
-                    units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
-            else
-                external_source_neutral_pressure_amplitude = nothing
-            end
-            if neutral_source_settings.PI_density_controller_I != 0.0 &&
-                    neutral_source_settings.source_type ∈ ("density_profile_control", "density_midpoint_control")
-                if neutral_source_settings.source_type == "density_profile_control"
-                    external_source_neutral_controller_integral = create_dynamic_variable!(
-                        dynamic, "external_source_neutral_controller_integral", mk_float, z, r;
-                        parallel_io=parallel_io,
-                        description="Integral term for the PID controller of the external source for neutrals")
-                else
-                    external_source_neutral_controller_integral = create_dynamic_variable!(
-                        dynamic, "external_source_neutral_controller_integral", mk_float;
-                        parallel_io=parallel_io,
-                        description="Integral term for the PID controller of the external source for neutrals")
-                end
-            else
-                external_source_neutral_controller_integral = nothing
-            end
-        else
-            external_source_neutral_amplitude = nothing
-            external_source_neutral_density_amplitude = nothing
-            external_source_neutral_momentum_amplitude = nothing
-            external_source_neutral_pressure_amplitude = nothing
-            external_source_neutral_controller_integral = nothing
-        end
+        io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
+        io_thermal_speed_neutral, external_source_neutral_amplitude,
+        external_source_neutral_density_amplitude,
+        external_source_neutral_momentum_amplitude,
+        external_source_neutral_pressure_amplitude,
+        external_source_neutral_controller_integral =
+            define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r, z,
+                                                     parallel_io,
+                                                     external_source_settings,
+                                                     evolve_density, evolve_upar,
+                                                     evolve_ppar)
 
         io_time_for_run = create_dynamic_variable!(
             dynamic, "time_for_run", mk_float; parallel_io=parallel_io,
@@ -888,6 +661,321 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
     return nothing
 end
 
+"""
+define dynamic (time-evolving) electromagnetic field variables for writing to the hdf5
+file
+"""
+function define_dynamic_em_field_variables!(fid, r::coordinate, z::coordinate,
+                                            parallel_io)
+
+    dynamic = get_group(fid, "dynamic_data")
+
+    # io_phi is the handle referring to the electrostatic potential phi
+    io_phi = create_dynamic_variable!(dynamic, "phi", mk_float, z, r;
+                                      parallel_io=parallel_io,
+                                      description="electrostatic potential",
+                                      units="T_ref/e")
+    # io_Er is the handle for the radial component of the electric field
+    io_Er = create_dynamic_variable!(dynamic, "Er", mk_float, z, r;
+                                     parallel_io=parallel_io,
+                                     description="radial electric field",
+                                     units="T_ref/e L_ref")
+    # io_Ez is the handle for the zed component of the electric field
+    io_Ez = create_dynamic_variable!(dynamic, "Ez", mk_float, z, r;
+                                     parallel_io=parallel_io,
+                                     description="vertical electric field",
+                                     units="T_ref/e L_ref")
+
+    return io_phi, io_Er, io_Ez
+end
+
+"""
+define dynamic (time-evolving) ion moment variables for writing to the hdf5 file
+"""
+function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
+        z::coordinate, parallel_io, external_source_settings, evolve_density, evolve_upar,
+        evolve_ppar)
+
+    dynamic = get_group(fid, "dynamic_data")
+
+    # io_density is the handle for the ion particle density
+    io_density = create_dynamic_variable!(dynamic, "density", mk_float, z, r;
+                                          n_ion_species=n_ion_species,
+                                          parallel_io=parallel_io,
+                                          description="ion species density",
+                                          units="n_ref")
+
+    # io_upar is the handle for the ion parallel flow density
+    io_upar = create_dynamic_variable!(dynamic, "parallel_flow", mk_float, z, r;
+                                       n_ion_species=n_ion_species,
+                                       parallel_io=parallel_io,
+                                       description="ion species parallel flow",
+                                       units="c_ref = sqrt(2*T_ref/mi)")
+
+    # io_ppar is the handle for the ion parallel pressure
+    io_ppar = create_dynamic_variable!(dynamic, "parallel_pressure", mk_float, z, r;
+                                       n_ion_species=n_ion_species,
+                                       parallel_io=parallel_io,
+                                       description="ion species parallel pressure",
+                                       units="n_ref*T_ref")
+
+    # io_pperp is the handle for the ion parallel pressure
+    io_pperp = create_dynamic_variable!(dynamic, "perpendicular_pressure", mk_float, z, r;
+                                        n_ion_species=n_ion_species,
+                                        parallel_io=parallel_io,
+                                        description="ion species perpendicular pressure",
+                                        units="n_ref*T_ref")
+
+    # io_qpar is the handle for the ion parallel heat flux
+    io_qpar = create_dynamic_variable!(dynamic, "parallel_heat_flux", mk_float, z, r;
+                                       n_ion_species=n_ion_species,
+                                       parallel_io=parallel_io,
+                                       description="ion species parallel heat flux",
+                                       units="n_ref*T_ref*c_ref")
+
+    # io_vth is the handle for the ion thermal speed
+    io_vth = create_dynamic_variable!(dynamic, "thermal_speed", mk_float, z, r;
+                                      n_ion_species=n_ion_species,
+                                      parallel_io=parallel_io,
+                                      description="ion species thermal speed",
+                                      units="c_ref")
+
+    # io_dSdt is the handle for the entropy production (due to collisions)
+    io_dSdt = create_dynamic_variable!(dynamic, "entropy_production", mk_float, z, r;
+                                      n_ion_species=n_ion_species,
+                                      parallel_io=parallel_io,
+                                      description="ion species entropy production",
+                                      units="")
+
+    ion_source_settings = external_source_settings.ion
+    if ion_source_settings.active
+        external_source_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external source for ions",
+            units="n_ref/c_ref^3*c_ref/L_ref")
+        if evolve_density
+            external_source_density_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_density_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external density source for ions",
+                units="n_ref*c_ref/L_ref")
+        else
+            external_source_density_amplitude = nothing
+        end
+        if evolve_upar
+            external_source_momentum_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_momentum_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external momentum source for ions",
+                units="m_ref*n_ref*c_ref*c_ref/L_ref")
+        else
+            external_source_momentum_amplitude = nothing
+        end
+        if evolve_ppar
+            external_source_pressure_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_pressure_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external pressure source for ions",
+                units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
+        else
+            external_source_pressure_amplitude = nothing
+        end
+        if ion_source_settings.PI_density_controller_I != 0.0 &&
+                ion_source_settings.source_type ∈ ("density_profile_control", "density_midpoint_control")
+            if ion_source_settings.source_type == "density_profile_control"
+                external_source_controller_integral = create_dynamic_variable!(
+                    dynamic, "external_source_controller_integral", mk_float, z, r;
+                    parallel_io=parallel_io,
+                    description="Integral term for the PID controller of the external source for ions")
+            else
+                external_source_controller_integral = create_dynamic_variable!(
+                    dynamic, "external_source_controller_integral", mk_float;
+                    parallel_io=parallel_io,
+                    description="Integral term for the PID controller of the external source for ions")
+            end
+        else
+            external_source_controller_integral = nothing
+        end
+    else
+        external_source_amplitude = nothing
+        external_source_density_amplitude = nothing
+        external_source_momentum_amplitude = nothing
+        external_source_pressure_amplitude = nothing
+        external_source_controller_integral = nothing
+    end
+
+    if parallel_io || z.irank == 0
+        # io_chodura_lower is the handle for the ion thermal speed
+        io_chodura_lower = create_dynamic_variable!(dynamic, "chodura_integral_lower", mk_float, r;
+                                          n_ion_species=n_ion_species,
+                                          parallel_io=parallel_io,
+                                          description="Generalised Chodura integral lower sheath entrance",
+                                          units="c_ref")
+    else
+        io_chodura_lower = nothing
+    end
+    if parallel_io || z.irank == z.nrank - 1
+        # io_chodura_upper is the handle for the ion thermal speed
+        io_chodura_upper = create_dynamic_variable!(dynamic, "chodura_integral_upper", mk_float, r;
+                                          n_ion_species=n_ion_species,
+                                          parallel_io=parallel_io,
+                                          description="Generalised Chodura integral upper sheath entrance",
+                                          units="c_ref")
+    else
+        io_chodura_upper = nothing
+    end
+
+    return io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
+           external_source_amplitude, external_source_density_amplitude,
+           external_source_momentum_amplitude, external_source_pressure_amplitude,
+           external_source_controller_integral, io_chodura_lower, io_chodura_upper
+end
+
+"""
+define dynamic (time-evolving) electron moment variables for writing to the hdf5 file
+"""
+function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordinate,
+        parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar)
+
+    dynamic = get_group(fid, "dynamic_data")
+
+    # io_density is the handle for the ion particle density
+    io_electron_density = create_dynamic_variable!(dynamic, "electron_density", mk_float, z, r;
+                                          parallel_io=parallel_io,
+                                          description="electron species density",
+                                          units="n_ref")
+
+    # io_electron_upar is the handle for the electron parallel flow density
+    io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
+                                       parallel_io=parallel_io,
+                                       description="electron species parallel flow",
+                                       units="c_ref = sqrt(2*T_ref/mi)")
+
+    # io_electron_ppar is the handle for the electron parallel pressure
+    io_electron_ppar = create_dynamic_variable!(dynamic, "electron_parallel_pressure", mk_float, z, r;
+                                       parallel_io=parallel_io,
+                                       description="electron species parallel pressure",
+                                       units="n_ref*T_ref")
+
+    # io_electron_qpar is the handle for the electron parallel heat flux
+    io_electron_qpar = create_dynamic_variable!(dynamic, "electron_parallel_heat_flux", mk_float, z, r;
+                                       parallel_io=parallel_io,
+                                       description="electron species parallel heat flux",
+                                       units="n_ref*T_ref*c_ref")
+
+    # io_electron_vth is the handle for the electron thermal speed
+    io_electron_vth = create_dynamic_variable!(dynamic, "electron_thermal_speed", mk_float, z, r;
+                                      parallel_io=parallel_io,
+                                      description="electron species thermal speed",
+                                      units="c_ref")
+
+    return io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
+           io_electron_vth
+end
+
+"""
+define dynamic (time-evolving) neutral moment variables for writing to the hdf5 file
+"""
+function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coordinate,
+        z::coordinate, parallel_io, external_source_settings, evolve_density, evolve_upar,
+        evolve_ppar)
+
+    dynamic = get_group(fid, "dynamic_data")
+
+    # io_density_neutral is the handle for the neutral particle density
+    io_density_neutral = create_dynamic_variable!(dynamic, "density_neutral", mk_float, z, r;
+                                                  n_neutral_species=n_neutral_species,
+                                                  parallel_io=parallel_io,
+                                                  description="neutral species density",
+                                                  units="n_ref")
+
+    # io_uz_neutral is the handle for the neutral z momentum density
+    io_uz_neutral = create_dynamic_variable!(dynamic, "uz_neutral", mk_float, z, r;
+                                             n_neutral_species=n_neutral_species,
+                                             parallel_io=parallel_io,
+                                             description="neutral species mean z velocity",
+                                             units="c_ref = sqrt(2*T_ref/mi)")
+
+    # io_pz_neutral is the handle for the neutral species zz pressure
+    io_pz_neutral = create_dynamic_variable!(dynamic, "pz_neutral", mk_float, z, r;
+                                             n_neutral_species=n_neutral_species,
+                                             parallel_io=parallel_io,
+                                             description="neutral species mean zz pressure",
+                                             units="n_ref*T_ref")
+
+    # io_qz_neutral is the handle for the neutral z heat flux
+    io_qz_neutral = create_dynamic_variable!(dynamic, "qz_neutral", mk_float, z, r;
+                                             n_neutral_species=n_neutral_species,
+                                             parallel_io=parallel_io,
+                                             description="neutral species z heat flux",
+                                             units="n_ref*T_ref*c_ref")
+
+    # io_thermal_speed_neutral is the handle for the neutral thermal speed
+    io_thermal_speed_neutral = create_dynamic_variable!(
+        dynamic, "thermal_speed_neutral", mk_float, z, r;
+        n_neutral_species=n_neutral_species,
+        parallel_io=parallel_io, description="neutral species thermal speed",
+        units="c_ref")
+
+    neutral_source_settings = external_source_settings.neutral
+    if n_neutral_species > 0 && neutral_source_settings.active
+        external_source_neutral_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_neutral_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external source for neutrals",
+            units="n_ref/c_ref^3*c_ref/L_ref")
+        if evolve_density
+            external_source_neutral_density_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_neutral_density_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external density source for neutrals",
+                units="n_ref*c_ref/L_ref")
+        else
+            external_source_neutral_density_amplitude = nothing
+        end
+        if evolve_upar
+            external_source_neutral_momentum_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_neutral_momentum_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external momentum source for neutrals",
+                units="m_ref*n_ref*c_ref*c_ref/L_ref")
+        else
+            external_source_neutral_momentum_amplitude = nothing
+        end
+        if evolve_ppar
+            external_source_neutral_pressure_amplitude = create_dynamic_variable!(
+                dynamic, "external_source_neutral_pressure_amplitude", mk_float, z, r;
+                parallel_io=parallel_io, description="Amplitude of the external pressure source for neutrals",
+                units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
+        else
+            external_source_neutral_pressure_amplitude = nothing
+        end
+        if neutral_source_settings.PI_density_controller_I != 0.0 &&
+                neutral_source_settings.source_type ∈ ("density_profile_control", "density_midpoint_control")
+            if neutral_source_settings.source_type == "density_profile_control"
+                external_source_neutral_controller_integral = create_dynamic_variable!(
+                    dynamic, "external_source_neutral_controller_integral", mk_float, z, r;
+                    parallel_io=parallel_io,
+                    description="Integral term for the PID controller of the external source for neutrals")
+            else
+                external_source_neutral_controller_integral = create_dynamic_variable!(
+                    dynamic, "external_source_neutral_controller_integral", mk_float;
+                    parallel_io=parallel_io,
+                    description="Integral term for the PID controller of the external source for neutrals")
+            end
+        else
+            external_source_neutral_controller_integral = nothing
+        end
+    else
+        external_source_neutral_amplitude = nothing
+        external_source_neutral_density_amplitude = nothing
+        external_source_neutral_momentum_amplitude = nothing
+        external_source_neutral_pressure_amplitude = nothing
+        external_source_neutral_controller_integral = nothing
+    end
+
+    return io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
+           io_thermal_speed_neutral, external_source_neutral_amplitude,
+           external_source_neutral_density_amplitude,
+           external_source_neutral_momentum_amplitude,
+           external_source_neutral_pressure_amplitude,
+           external_source_neutral_controller_integral
+end
+
 """
 define dynamic (time-evolving) distribution function variables for writing to the output
 file
@@ -1170,11 +1258,12 @@ function append_to_dynamic_var() end
 end
 
 """
-write time-dependent moments data to the binary output file
+write time-dependent moments data for ions, electrons and neutrals to the binary output
+file
 """
-function write_moments_data_to_binary(moments, fields, t, n_ion_species,
-                                      n_neutral_species, io_or_file_info_moments, t_idx,
-                                      time_for_run, r, z)
+function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
+                                          n_neutral_species, io_or_file_info_moments,
+                                          t_idx, time_for_run, r, z)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
@@ -1191,11 +1280,71 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
         # add the time for this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.time, t, t_idx, parallel_io)
 
+        write_em_fields_data_to_binary(fields, io_or_file_info_moments, t_idx, r, z)
+
+        write_ion_moments_data_to_binary(moments, n_ion_species, io_or_file_info_moments,
+                                         t_idx, r, z)
+
+        write_electron_moments_data_to_binary(moments, io_or_file_info_moments, t_idx, r,
+                                              z)
+
+        write_neutral_moments_data_to_binary(moments, n_neutral_species,
+                                             io_or_file_info_moments, t_idx, r, z)
+
+        append_to_dynamic_var(io_moments.time_for_run, time_for_run, t_idx, parallel_io)
+
+        closefile && close(io_moments.fid)
+    end
+
+    return nothing
+end
+
+"""
+write time-dependent EM fields data to the binary output file
+"""
+function write_em_fields_data_to_binary(fields, io_or_file_info_moments, t_idx, r, z)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_moments, io_moments_info)
+            io_moments = io_or_file_info_moments
+            closefile = false
+        else
+            io_moments = reopen_moments_io(io_or_file_info_moments)
+            closefile = true
+        end
+
+        parallel_io = io_moments.parallel_io
+
         # add the electrostatic potential and electric field components at this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.phi, fields.phi, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.Er, fields.Er, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.Ez, fields.Ez, t_idx, parallel_io, z, r)
 
+        closefile && close(io_moments.fid)
+    end
+
+    return nothing
+end
+
+"""
+write time-dependent moments data for ions to the binary output file
+"""
+function write_ion_moments_data_to_binary(moments, n_ion_species, io_or_file_info_moments,
+                                          t_idx, r, z)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_moments, io_moments_info)
+            io_moments = io_or_file_info_moments
+            closefile = false
+        else
+            io_moments = reopen_moments_io(io_or_file_info_moments)
+            closefile = true
+        end
+
+        parallel_io = io_moments.parallel_io
+
         # add the density data at this time slice to the output file
         append_to_dynamic_var(io_moments.density, moments.ion.dens, t_idx,
                               parallel_io, z, r, n_ion_species)
@@ -1260,70 +1409,126 @@ function write_moments_data_to_binary(moments, fields, t, n_ion_species,
                                       t_idx, parallel_io, z, r)
             end
         end
-        # add the electron velocity-moments data at this time slice to the output file
-        append_to_dynamic_var(io_moments.electron_density, moments.electron.dens.data, t_idx, z, r)
-        append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar.data, t_idx, z, r)
-        append_to_dynamic_var(io_moments.electron_parallel_pressure, moments.electron.ppar.data, t_idx, z, r)
-        append_to_dynamic_var(io_moments.electron_parallel_heat_flux, moments.electron.qpar.data, t_idx, z, r)
-        append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth.data, t_idx, z, r)
-        if n_neutral_species > 0
-            append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
-                                  parallel_io, z, r, n_neutral_species)
-            append_to_dynamic_var(io_moments.uz_neutral, moments.neutral.uz, t_idx,
-                                  parallel_io, z, r, n_neutral_species)
-            append_to_dynamic_var(io_moments.pz_neutral, moments.neutral.pz, t_idx,
-                                  parallel_io, z, r, n_neutral_species)
-            append_to_dynamic_var(io_moments.qz_neutral, moments.neutral.qz, t_idx,
-                                  parallel_io, z, r, n_neutral_species)
-            append_to_dynamic_var(io_moments.thermal_speed_neutral, moments.neutral.vth,
-                                  t_idx, parallel_io, z, r, n_neutral_species)
-
-            if io_moments.external_source_neutral_amplitude !== nothing
-                append_to_dynamic_var(io_moments.external_source_neutral_amplitude,
-                                      moments.neutral.external_source_amplitude, t_idx,
-                                      parallel_io, z, r)
-                if moments.evolve_density
-                    append_to_dynamic_var(io_moments.external_source_neutral_density_amplitude,
-                                          moments.neutral.external_source_density_amplitude,
-                                          t_idx, parallel_io, z, r)
-                end
-                if moments.evolve_upar
-                    append_to_dynamic_var(io_moments.external_source_neutral_momentum_amplitude,
-                                          moments.neutral.external_source_momentum_amplitude,
-                                          t_idx, parallel_io, z, r)
-                end
-                if moments.evolve_ppar
-                    append_to_dynamic_var(io_moments.external_source_neutral_pressure_amplitude,
-                                          moments.neutral.external_source_pressure_amplitude,
-                                          t_idx, parallel_io, z, r)
-                end
+
+        closefile && close(io_moments.fid)
+    end
+
+    return nothing
+end
+
+"""
+write time-dependent moments data for electrons to the binary output file
+"""
+function write_electron_moments_data_to_binary(moments, io_or_file_info_moments, t_idx, r,
+                                               z)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_moments, io_moments_info)
+            io_moments = io_or_file_info_moments
+            closefile = false
+        else
+            io_moments = reopen_moments_io(io_or_file_info_moments)
+            closefile = true
+        end
+
+        parallel_io = io_moments.parallel_io
+
+        append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx,
+                              parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar,
+                              t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_pressure,
+                              moments.electron.ppar, t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_heat_flux,
+                              moments.electron.qpar, t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
+                              t_idx, parallel_io, z, r)
+
+        closefile && close(io_moments.fid)
+    end
+
+    return nothing
+end
+
+"""
+write time-dependent moments data for neutrals to the binary output file
+"""
+function write_neutral_moments_data_to_binary(moments, n_neutral_species,
+                                              io_or_file_info_moments, t_idx, r, z)
+    if n_neutral_species ≤ 0
+        return nothing
+    end
+
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_moments, io_moments_info)
+            io_moments = io_or_file_info_moments
+            closefile = false
+        else
+            io_moments = reopen_moments_io(io_or_file_info_moments)
+            closefile = true
+        end
+
+        parallel_io = io_moments.parallel_io
+
+        append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
+                              parallel_io, z, r, n_neutral_species)
+        append_to_dynamic_var(io_moments.uz_neutral, moments.neutral.uz, t_idx,
+                              parallel_io, z, r, n_neutral_species)
+        append_to_dynamic_var(io_moments.pz_neutral, moments.neutral.pz, t_idx,
+                              parallel_io, z, r, n_neutral_species)
+        append_to_dynamic_var(io_moments.qz_neutral, moments.neutral.qz, t_idx,
+                              parallel_io, z, r, n_neutral_species)
+        append_to_dynamic_var(io_moments.thermal_speed_neutral, moments.neutral.vth,
+                              t_idx, parallel_io, z, r, n_neutral_species)
+
+        if io_moments.external_source_neutral_amplitude !== nothing
+            append_to_dynamic_var(io_moments.external_source_neutral_amplitude,
+                                  moments.neutral.external_source_amplitude, t_idx,
+                                  parallel_io, z, r)
+            if moments.evolve_density
+                append_to_dynamic_var(io_moments.external_source_neutral_density_amplitude,
+                                      moments.neutral.external_source_density_amplitude,
+                                      t_idx, parallel_io, z, r)
             end
-            if io_moments.external_source_neutral_controller_integral !== nothing
-                if size(moments.neutral.external_source_neutral_controller_integral) == (1,1)
-                    append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
-                                          moments.neutral.external_source_controller_integral[1,1],
-                                          t_idx, parallel_io)
-                else
-                    append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
-                                          moments.neutral.external_source_controller_integral,
-                                          t_idx, parallel_io, z, r)
-                end
+            if moments.evolve_upar
+                append_to_dynamic_var(io_moments.external_source_neutral_momentum_amplitude,
+                                      moments.neutral.external_source_momentum_amplitude,
+                                      t_idx, parallel_io, z, r)
+            end
+            if moments.evolve_ppar
+                append_to_dynamic_var(io_moments.external_source_neutral_pressure_amplitude,
+                                      moments.neutral.external_source_pressure_amplitude,
+                                      t_idx, parallel_io, z, r)
+            end
+        end
+        if io_moments.external_source_neutral_controller_integral !== nothing
+            if size(moments.neutral.external_source_neutral_controller_integral) == (1,1)
+                append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
+                                      moments.neutral.external_source_controller_integral[1,1],
+                                      t_idx, parallel_io)
+            else
+                append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
+                                      moments.neutral.external_source_controller_integral,
+                                      t_idx, parallel_io, z, r)
             end
         end
-
-        append_to_dynamic_var(io_moments.time_for_run, time_for_run, t_idx, parallel_io)
 
         closefile && close(io_moments.fid)
     end
+
     return nothing
 end
 
 """
-write time-dependent distribution function data to the binary output file
+write time-dependent distribution function data for ions and neutrals to the
+binary output file
 """
-function write_dfns_data_to_binary(ff, ff_neutral, moments, fields, t, n_ion_species,
-                                   n_neutral_species, io_or_file_info_dfns, t_idx,
-                                   time_for_run, r, z, vperp, vpa, vzeta, vr, vz)
+function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
+                                       n_neutral_species, io_or_file_info_dfns, t_idx,
+                                       time_for_run, r, z, vperp, vpa, vzeta, vr, vz)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
@@ -1337,14 +1542,67 @@ function write_dfns_data_to_binary(ff, ff_neutral, moments, fields, t, n_ion_spe
 
         # Write the moments for this time slice to the output file.
         # This also updates the time.
-        write_moments_data_to_binary(moments, fields, t, n_ion_species, n_neutral_species,
-                                     io_dfns.io_moments, t_idx, time_for_run, r, z)
+        write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
+                                         n_neutral_species, io_dfns.io_moments, t_idx,
+                                         time_for_run, r, z)
+
+        # add the distribution function data at this time slice to the output file
+        write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_or_file_info_dfns,
+                                      t_idx, r, z, vperp, vpa)
+        write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species,
+                                          io_or_file_info_dfns, t_idx, r, z, vzeta, vr,
+                                          vz)
+
+        closefile && close(io_dfns.fid)
+    end
+    return nothing
+end
+
+"""
+write time-dependent distribution function data for ions to the binary output file
+"""
+function write_ion_dfns_data_to_binary(ff, n_ion_species, io_or_file_info_dfns, t_idx, r,
+                                       z, vperp, vpa)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_dfns, io_dfns_info)
+            io_dfns = io_or_file_info_dfns
+            closefile = false
+        else
+            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
+            closefile = true
+        end
 
         parallel_io = io_dfns.parallel_io
 
-        # add the distribution function data at this time slice to the output file
         append_to_dynamic_var(io_dfns.f, ff, t_idx, parallel_io, vpa, vperp, z, r,
                               n_ion_species)
+
+        closefile && close(io_dfns.fid)
+    end
+    return nothing
+end
+
+"""
+write time-dependent distribution function data for neutrals to the binary output file
+"""
+function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
+                                           io_or_file_info_dfns, t_idx, r, z, vzeta, vr,
+                                           vz)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_dfns, io_dfns_info)
+            io_dfns = io_or_file_info_dfns
+            closefile = false
+        else
+            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
+            closefile = true
+        end
+
+        parallel_io = io_dfns.parallel_io
+
         if n_neutral_species > 0
             append_to_dynamic_var(io_dfns.f_neutral, ff_neutral, t_idx, parallel_io, vz,
                                   vr, vzeta, z, r, n_neutral_species)
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index abc0a5ae4..5483d87b5 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -80,7 +80,7 @@ using Glob
 
 using .file_io: setup_file_io, finish_file_io
 using .file_io: write_data_to_ascii
-using .file_io: write_moments_data_to_binary, write_dfns_data_to_binary
+using .file_io: write_all_moments_data_to_binary, write_all_dfns_data_to_binary
 using .command_line_options: get_options
 using .communication
 using .communication: _block_synchronize
@@ -438,11 +438,12 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
-    write_moments_data_to_binary(moments, fields, code_time, composition.n_ion_species,
-        composition.n_neutral_species, io_moments, 1, 0.0, r, z)
-    write_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
-         code_time, composition.n_ion_species, composition.n_neutral_species, io_dfns, 1,
-         0.0, r, z, vperp, vpa, vzeta, vr, vz)
+    write_all_moments_data_to_binary(moments, fields, code_time,
+        composition.n_ion_species, composition.n_neutral_species, io_moments, 1, 0.0, r,
+        z)
+    write_all_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
+        code_time, composition.n_ion_species, composition.n_neutral_species, io_dfns, 1,
+        0.0, r, z, vperp, vpa, vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 
diff --git a/src/time_advance.jl b/src/time_advance.jl
index 64ec5dae4..6409fb045 100644
--- a/src/time_advance.jl
+++ b/src/time_advance.jl
@@ -15,7 +15,7 @@ using ..array_allocation: allocate_float, allocate_shared_float
 using ..communication
 using ..communication: _block_synchronize
 using ..debugging
-using ..file_io: write_data_to_ascii, write_moments_data_to_binary, write_dfns_data_to_binary, debug_dump
+using ..file_io: write_data_to_ascii, write_all_moments_data_to_binary, write_all_dfns_data_to_binary, debug_dump
 using ..looping
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: update_moments!, update_moments_neutral!, reset_moments_status!
@@ -1060,9 +1060,10 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
             write_data_to_ascii(moments, fields, vpa, vperp, z, r, t,
                                 composition.n_ion_species, composition.n_neutral_species,
                                 ascii_io)
-            write_moments_data_to_binary(moments, fields, t, composition.n_ion_species,
-                                         composition.n_neutral_species, io_moments,
-                                         iwrite_moments, time_for_run, r, z)
+            write_all_moments_data_to_binary(moments, fields, t,
+                                             composition.n_ion_species,
+                                             composition.n_neutral_species, io_moments,
+                                             iwrite_moments, time_for_run, r, z)
 
             if t_input.steady_state_residual
                 # Calculate some residuals to see how close simulation is to steady state
@@ -1235,10 +1236,11 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                     flush(stdout)
                 end
             end
-            write_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
-                                      t, composition.n_ion_species,
-                                      composition.n_neutral_species, io_dfns, iwrite_dfns,
-                                      time_for_run, r, z, vperp, vpa, vzeta, vr, vz)
+            write_all_dfns_data_to_binary(pdf, moments, fields, t,
+                                          composition.n_ion_species,
+                                          composition.n_neutral_species, io_dfns,
+                                          iwrite_dfns, time_for_run, r, z, vperp, vpa,
+                                          vzeta, vr, vz)
             iwrite_dfns += 1
             begin_s_r_z_vperp_region()
             @debug_detect_redundant_block_synchronize begin

From 350e090dcc37b4717d8256434a15ed9eb6160139 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 16 Feb 2024 16:35:10 +0000
Subject: [PATCH 094/428] Add output for electron distribution function

---
 src/file_io.jl | 72 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 56 insertions(+), 16 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index fb433c464..e925673bd 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -117,11 +117,13 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn, Tchodura_lower,
 structure containing the data/metadata needed for binary file i/o
 distribution function data only
 """
-struct io_dfns_info{Tfile, Tfi, Tfn, Tmoments}
+struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the ion species distribution function variable
     f::Tfi
+    # handle for the electron distribution function variable
+    f_electron::Tfe
     # handle for the neutral species distribution function variable
     f_neutral::Tfn
 
@@ -980,14 +982,13 @@ end
 define dynamic (time-evolving) distribution function variables for writing to the output
 file
 """
-function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz,
-                                       n_ion_species, n_neutral_species, parallel_io,
-                                       external_source_settings, evolve_density,
-                                       evolve_upar, evolve_ppar)
+function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, composition,
+                                       parallel_io, external_source_settings,
+                                       evolve_density, evolve_upar, evolve_ppar)
 
     @serial_region begin
-        io_moments = define_dynamic_moment_variables!(fid, n_ion_species,
-                                                      n_neutral_species, r, z,
+        io_moments = define_dynamic_moment_variables!(fid, composition.n_ion_species,
+                                                      composition.n_neutral_species, r, z,
                                                       parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -997,17 +998,27 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz,
 
         # io_f is the handle for the ion pdf
         io_f = create_dynamic_variable!(dynamic, "f", mk_float, vpa, vperp, z, r;
-                                        n_ion_species=n_ion_species,
+                                        n_ion_species=composition.n_ion_species,
                                         parallel_io=parallel_io,
                                         description="ion species distribution function")
 
+        if composition.electron_physics == kinetic_electrons
+            # io_f_electron is the handle for the electron pdf
+            io_f_electron = create_dynamic_variable!(dynamic, "f_electron", mk_float, vpa,
+                                                     vperp, z, r;
+                                                     parallel_io=parallel_io,
+                                                     description="electron distribution function")
+        else
+            io_f_electron = nothing
+        end
+
         # io_f_neutral is the handle for the neutral pdf
         io_f_neutral = create_dynamic_variable!(dynamic, "f_neutral", mk_float, vz, vr, vzeta, z, r;
-                                                n_neutral_species=n_neutral_species,
+                                                n_neutral_species=composition.n_neutral_species,
                                                 parallel_io=parallel_io,
                                                 description="neutral species distribution function")
 
-        return io_dfns_info(fid, io_f, io_f_neutral, parallel_io, io_moments)
+        return io_dfns_info(fid, io_f, io_f_electron, io_f_neutral, parallel_io, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1179,9 +1190,8 @@ function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vper
         ### create variables for time-dependent quantities and store them ###
         ### in a struct for later access ###
         io_dfns = define_dynamic_dfn_variables!(
-            fid, r, z, vperp, vpa, vzeta, vr, vz, composition.n_ion_species,
-            composition.n_neutral_species, parallel_io, external_source_settings,
-            evolve_density, evolve_upar, evolve_ppar)
+            fid, r, z, vperp, vpa, vzeta, vr, vz, composition, parallel_io,
+            external_source_settings, evolve_density, evolve_upar, evolve_ppar)
 
         close(fid)
 
@@ -1232,8 +1242,8 @@ function reopen_dfns_io(file_info)
                                      getvar("time_for_run"),
                                      parallel_io)
 
-        return io_dfns_info(fid, getvar("f"), getvar("f_neutral"), parallel_io,
-                            io_moments)
+        return io_dfns_info(fid, getvar("f"), getvar("f_electron"), getvar("f_neutral"),
+                            parallel_io, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1523,7 +1533,7 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
 end
 
 """
-write time-dependent distribution function data for ions and neutrals to the
+write time-dependent distribution function data for ions, electrons and neutrals to the
 binary output file
 """
 function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
@@ -1549,6 +1559,8 @@ function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
         # add the distribution function data at this time slice to the output file
         write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_or_file_info_dfns,
                                       t_idx, r, z, vperp, vpa)
+        write_electron_dfns_data_to_binary(pdf.electron.norm, io_or_file_info_dfns, t_idx,
+                                           r, z, vperp, vpa)
         write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species,
                                           io_or_file_info_dfns, t_idx, r, z, vzeta, vr,
                                           vz)
@@ -1584,6 +1596,34 @@ function write_ion_dfns_data_to_binary(ff, n_ion_species, io_or_file_info_dfns,
     return nothing
 end
 
+"""
+write time-dependent distribution function data for electrons to the binary output file
+"""
+function write_electron_dfns_data_to_binary(ff_electron, io_or_file_info_dfns, t_idx, r,
+                                            z, vperp, vpa)
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_dfns, io_dfns_info)
+            io_dfns = io_or_file_info_dfns
+            closefile = false
+        else
+            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
+            closefile = true
+        end
+
+        parallel_io = io_dfns.parallel_io
+
+        if io_dfns.f_electron !== nothing
+            append_to_dynamic_var(io_dfns.f_electron, ff_electron, t_idx, parallel_io,
+                                  vpa, vperp, z, r)
+        end
+
+        closefile && close(io_dfns.fid)
+    end
+    return nothing
+end
+
 """
 write time-dependent distribution function data for neutrals to the binary output file
 """

From 647d53e390f4a3868e6faec03f50226c1f4a1fc7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 17 Feb 2024 13:40:26 +0000
Subject: [PATCH 095/428] Save run_id in io_input

This makes it easier to add the run_id to multiple output files, which
might not be created by the same function.
---
 src/file_io.jl               | 3 +--
 src/input_structs.jl         | 1 +
 src/moment_kinetics_input.jl | 2 ++
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index e925673bd..f78499bbb 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -20,7 +20,6 @@ using ..type_definitions: mk_float, mk_int
 using LibGit2
 using MPI
 using Pkg
-using UUIDs
 using TOML
 
 @debug_shared_array using ..communication: DebugMPISharedArray
@@ -170,7 +169,7 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
             ascii = ascii_ios(nothing, nothing, nothing, nothing, nothing)
         end
 
-        run_id = string(uuid4())
+        run_id = io_input.run_id
 
         io_moments = setup_moments_io(out_prefix, io_input.binary_format, r, z,
                                       composition, collisions, evolve_density,
diff --git a/src/input_structs.jl b/src/input_structs.jl
index 373246e67..7aac6fc6e 100644
--- a/src/input_structs.jl
+++ b/src/input_structs.jl
@@ -371,6 +371,7 @@ Base.@kwdef struct io_input
     ascii_output::Bool
     binary_format::binary_format_type
     parallel_io::Bool
+    run_id::String
 end
 
 """
diff --git a/src/moment_kinetics_input.jl b/src/moment_kinetics_input.jl
index 615e02694..b57923de3 100644
--- a/src/moment_kinetics_input.jl
+++ b/src/moment_kinetics_input.jl
@@ -21,6 +21,7 @@ using ..reference_parameters
 
 using MPI
 using TOML
+using UUIDs
 
 """
 Read input from a TOML file
@@ -522,6 +523,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     io_settings["binary_format"] = get(io_settings, "binary_format", hdf5)
     io_settings["parallel_io"] = get(io_settings, "parallel_io",
                                      io_has_parallel(Val(io_settings["binary_format"])))
+    io_settings["run_id"] = string(uuid4())
     io_immutable = io_input(; output_dir=output_dir, run_name=run_name,
                               Dict(Symbol(k)=>v for (k,v) in io_settings)...)
 

From 1e4104c35bc6f10bf5eab1f938be224789fea5d5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 17 Feb 2024 18:15:09 +0000
Subject: [PATCH 096/428] Finish implementing electron I/O for moments/dfn

Need to add the variables also in the reopen_*_io() functions.
---
 src/file_io.jl | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/file_io.jl b/src/file_io.jl
index f78499bbb..23be48b6c 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -1126,9 +1126,15 @@ function reopen_moments_io(file_info)
                                getvar("parallel_heat_flux"),
                                getvar("thermal_speed"), getvar("entropy_production"),
                                getvar("chodura_integral_lower"),
-                               getvar("chodura_integral_upper"), getvar("density_neutral"),
-                               getvar("uz_neutral"), getvar("pz_neutral"),
-                               getvar("qz_neutral"), getvar("thermal_speed_neutral"),
+                               getvar("chodura_integral_upper"),
+                               getvar("electron_density"),
+                               getvar("electron_parallel_flow"),
+                               getvar("electron_parallel_pressure"),
+                               getvar("electron_parallel_heat_flux"),
+                               getvar("electron_thermal_speed"),
+                               getvar("density_neutral"), getvar("uz_neutral"),
+                               getvar("pz_neutral"), getvar("qz_neutral"),
+                               getvar("thermal_speed_neutral"),
                                getvar("external_source_amplitude"),
                                getvar("external_source_density_amplitude"),
                                getvar("external_source_momentum_amplitude"),
@@ -1224,9 +1230,14 @@ function reopen_dfns_io(file_info)
                                      getvar("perpendicular_pressure"),
                                      getvar("parallel_heat_flux"), getvar("thermal_speed"),
                                      getvar("entropy_production"), getvar("chodura_integral_lower"),
-                                     getvar("chodura_integral_upper"), getvar("density_neutral"),
-                                     getvar("uz_neutral"), getvar("pz_neutral"),
-                                     getvar("qz_neutral"),
+                                     getvar("chodura_integral_upper"),
+                                     getvar("electron_density"),
+                                     getvar("electron_parallel_flow"),
+                                     getvar("electron_parallel_pressure"),
+                                     getvar("electron_parallel_heat_flux"),
+                                     getvar("electron_thermal_speed"),
+                                     getvar("density_neutral"), getvar("uz_neutral"),
+                                     getvar("pz_neutral"), getvar("qz_neutral"),
                                      getvar("thermal_speed_neutral"),
                                      getvar("external_source_amplitude"),
                                      getvar("external_source_density_amplitude"),

From 340cbc44616aab63b2291c0232918706547b3508 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 11:47:53 +0000
Subject: [PATCH 097/428] Update scratch arrays within
 update_electron_pdf_with_time_advance()

Removes the need to call initialize_scratch_arrays!() a second time.
---
 src/electron_kinetic_equation.jl | 5 +++++
 src/initial_conditions.jl        | 6 +++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index 830de4492..ec9853334 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -468,6 +468,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     catch e
         println("Error: $e")
     end
+    # Update the 'scratch' arrays with the final result
+    begin_r_z_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        fvec.pdf_electron[ivpa,ivperp,iz,ir] = pdf[ivpa,ivperp,iz,ir]
+    end
     begin_serial_region()
     @serial_region begin
         if !electron_pdf_converged
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 0d9766d55..0fccd41e7 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -357,7 +357,9 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     
     # initialize the scratch arrays containing pdfs and moments for the first RK stage
     # the electron pdf is yet to be initialised but with the current code logic, the scratch
-    # arrays need to exist and be otherwise initialised in order to compute the initial electron pdf
+    # arrays need to exist and be otherwise initialised in order to compute the initial
+    # electron pdf. The electron arrays will be updated as necessary by
+    # initialize_electron_pdf!().
     initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
     # get the initial electrostatic potential and parallel electric field
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
@@ -367,8 +369,6 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                              advection_structs.electron_z_advect, advection_structs.electron_vpa_advect,
                              scratch_dummy, collisions, composition, 
                              num_diss_params, t_input.dt)
-    # re-initialize the scratch arrays now that the electron pdf has been initialised
-    initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
 
     return nothing
 end

From 376d15f20d93c4eda64d0d651c615171c81ca196 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 11:50:03 +0000
Subject: [PATCH 098/428] Write the initial electron state to a binary file

Write the converged initial state of the electrons to a file. Also (for
debugging) write the state at several intermediate steps.
---
 src/electron_kinetic_equation.jl |  36 +++++++----
 src/file_io.jl                   | 106 +++++++++++++++++++++++++++++--
 src/initial_conditions.jl        |  58 +++++++++++++----
 src/moment_kinetics.jl           |   2 +-
 4 files changed, 173 insertions(+), 29 deletions(-)

diff --git a/src/electron_kinetic_equation.jl b/src/electron_kinetic_equation.jl
index ec9853334..811ca1107 100644
--- a/src/electron_kinetic_equation.jl
+++ b/src/electron_kinetic_equation.jl
@@ -17,6 +17,7 @@ using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!
 using ..electron_vpa_advection: electron_vpa_advection!
+using ..file_io: write_electron_dfns_data_to_binary, write_electron_moments_data_to_binary
 using ..moment_constraints: hard_force_moment_constraints!
 using ..velocity_moments: integrate_over_vspace
 
@@ -46,12 +47,10 @@ The electron kinetic equation is:
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
-function update_electron_pdf!(fvec, pdf, moments, dens, vthe, 
-    ppar, qpar, qpar_updated, 
-    phi, ddens_dz, dppar_dz, 
-    dqpar_dz, dvth_dz, z, vpa, z_spectral, 
-    vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition, 
-    num_diss_params, max_electron_pdf_iterations)
+function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
+        phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
+        vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
+        num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -61,8 +60,8 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe,
     if solution_method == "artificial_time_derivative"
         return update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
             moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
-            z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
-            num_diss_params, max_electron_pdf_iterations)
+            r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
+            num_diss_params, max_electron_pdf_iterations; io_initial_electron=io_initial_electron)
     elseif solution_method == "shooting_method"
         update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar, qpar_updated, 
             phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy, composition)
@@ -103,8 +102,8 @@ OUTPUT:
 """
 function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
     moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition,
-    z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, 
-    num_diss_params, max_electron_pdf_iterations)
+    r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
+    num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing)
 
     begin_r_z_region()
 
@@ -242,8 +241,16 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     result_phi = nothing
     result_residual = nothing
     output_interval = 1000
+    output_counter = 0
     begin_serial_region()
     @serial_region begin
+        output_counter += 1
+        if io_initial_electron !== nothing
+            write_electron_dfns_data_to_binary(pdf, io_initial_electron,
+                                               output_counter, r, z, vperp, vpa)
+            write_electron_moments_data_to_binary(moments, io_initial_electron,
+                                                  output_counter, r, z)
+        end
         result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
         result_pdf[:,:,1] .= pdf[:,1,:,1]
         result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
@@ -341,6 +348,13 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 println(io_qpar,"")
                 println(io_ppar,"")
                 println(io_vth,"")
+                output_counter += 1
+                if io_initial_electron !== nothing
+                    write_electron_dfns_data_to_binary(pdf, io_initial_electron,
+                                                       output_counter, r, z, vperp, vpa)
+                    write_electron_moments_data_to_binary(moments, io_initial_electron,
+                                                          output_counter, r, z)
+                end
                 result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
                 result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
                 result_vth[:,iteration÷output_interval+1] .= vthe[:,1]
@@ -500,7 +514,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         result_phi = result_phi[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
         result_residual = result_residual[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
     end
-    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual
+    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual, output_counter
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
diff --git a/src/file_io.jl b/src/file_io.jl
index 23be48b6c..7dab01ce8 100644
--- a/src/file_io.jl
+++ b/src/file_io.jl
@@ -44,7 +44,7 @@ end
 structure containing the data/metadata needed for binary file i/o
 moments & fields only
 """
-struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn, Tchodura_lower,
+struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
                        Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
                        Texti5, Textn1, Textn2, Textn3, Textn4, Textn5}
     # file identifier for the binary file to which data is written
@@ -76,15 +76,15 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmomn, Tchodura_lower,
     # handle for chodura diagnostic (upper)
     chodura_integral_upper::Tchodura_upper
     # handle for the electron species density
-    electron_density::Tphi
+    electron_density::Tmome
     # handle for the electron species parallel flow
-    electron_parallel_flow::Tphi
+    electron_parallel_flow::Tmome
     # handle for the electron species parallel pressure
-    electron_parallel_pressure::Tphi
+    electron_parallel_pressure::Tmome
     # handle for the electron species parallel heat flux
-    electron_parallel_heat_flux::Tphi
+    electron_parallel_heat_flux::Tmome
     # handle for the electron species thermal speed
-    electron_thermal_speed::Tphi
+    electron_thermal_speed::Tmome
 
     # handle for the neutral species density
     density_neutral::Tmomn
@@ -133,6 +133,28 @@ struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments}
     io_moments::Tmoments
 end
 
+"""
+structure containing the data/metadata needed for binary file i/o
+for electron initialization
+"""
+struct io_electron_initialization_info{Tfile, Tfe, Tmom}
+    # file identifier for the binary file to which data is written
+    fid::Tfile
+    # handle for the electron distribution function variable
+    f_electron::Tfe
+    # handle for the electron density variable
+    electron_density::Tmom
+    # handle for the electron parallel flow variable
+    electron_parallel_flow::Tmom
+    # handle for the electron parallel pressure variable
+    electron_parallel_pressure::Tmom
+    # handle for the electron parallel heat flux variable
+    electron_parallel_heat_flux::Tmom
+
+    # Use parallel I/O?
+    parallel_io::Bool
+end
+
 """
     io_has_parallel(Val(binary_format))
 
@@ -190,6 +212,78 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
     return nothing, nothing, nothing
 end
 
+"""
+open output file to save the initial electron pressure and distribution function
+"""
+function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition,
+                                   collisions, evolve_density, evolve_upar, evolve_ppar,
+                                   external_source_settings, input_dict,
+                                   previous_runs_info)
+    begin_serial_region()
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        # check to see if output_dir exists in the current directory
+        # if not, create it
+        isdir(io_input.output_dir) || mkdir(io_input.output_dir)
+        out_prefix = joinpath(io_input.output_dir, io_input.run_name)
+
+        run_id = io_input.run_id
+        parallel_io = io_input.parallel_io
+        io_comm = comm_inter_block[]
+
+        # dummy value for restart_time_index, which is not used for the initial electron
+        # state
+        restart_time_index = -1
+
+        electrons_prefix = string(out_prefix, ".initial_electrons")
+        if !parallel_io
+            electrons_prefix *= ".$(iblock_index[])"
+        end
+        fid, file_info = open_output_file(electrons_prefix, io_input.binary_format,
+                                          parallel_io, io_comm)
+
+        # write a header to the output file
+        add_attribute!(fid, "file_info",
+                       "Output initial electron state from the moment_kinetics code")
+
+        # write some overview information to the output file
+        write_overview!(fid, composition, collisions, parallel_io, evolve_density,
+                        evolve_upar, evolve_ppar, -1.0)
+
+        # write provenance tracking information to the output file
+        write_provenance_tracking_info!(fid, parallel_io, run_id, restart_time_index,
+                                        input_dict, previous_runs_info)
+
+        # write the input settings
+        write_input!(fid, input_dict, parallel_io)
+
+        ### define coordinate dimensions ###
+        coords_group = define_spatial_coordinates!(fid, z, r, parallel_io)
+        add_vspace_coordinates!(coords_group, vz, vr, vzeta, vpa, vperp, parallel_io)
+
+        ### create variables for time-dependent quantities ###
+        dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
+        io_f_electron = create_dynamic_variable!(dynamic, "f_electron", mk_float, vpa,
+                                                 vperp, z, r;
+                                                 parallel_io=parallel_io,
+                                                 description="electron distribution function")
+
+        io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
+        io_electron_vth =
+            define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
+                                                      external_source_settings,
+                                                      evolve_density, evolve_upar,
+                                                      evolve_ppar)
+
+        close(fid)
+
+        return file_info
+    end
+    # For other processes in the block, return (nothing, nothing, nothing)
+    return nothing, nothing, nothing
+end
+
 """
 Get a (sub-)group from a file or group
 """
diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl
index 0fccd41e7..23e6e58d4 100644
--- a/src/initial_conditions.jl
+++ b/src/initial_conditions.jl
@@ -24,6 +24,8 @@ using ..external_sources
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
 using ..em_fields: update_phi!
+using ..file_io: setup_initial_electron_io, write_electron_dfns_data_to_binary,
+                 write_electron_moments_data_to_binary
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
@@ -212,7 +214,8 @@ with a self-consistent initial condition
 function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry, composition, r, z,
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, vz_spectral,
                                species, collisions, external_source_settings, manufactured_solns_input,
-                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs)
+                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
+                               io_input, input_dict)
     if manufactured_solns_input.use_for_init
         init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z,
                                              r, composition.n_ion_species,
@@ -365,10 +368,12 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
 
     # initialize the electron pdf that satisfies the electron kinetic equation
-    return initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, z, vpa, vperp, z_spectral, vpa_spectral, 
-                             advection_structs.electron_z_advect, advection_structs.electron_vpa_advect,
-                             scratch_dummy, collisions, composition, 
-                             num_diss_params, t_input.dt)
+    return initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa,
+                                    vperp, vzeta, vr, vz, z_spectral, vpa_spectral,
+                                    advection_structs.electron_z_advect,
+                                    advection_structs.electron_vpa_advect, scratch_dummy,
+                                    collisions, composition, external_source_settings,
+                                    num_diss_params, t_input.dt, io_input, input_dict)
 
     return nothing
 end
@@ -489,13 +494,31 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
     return nothing
 end
 
-function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spectral, vpa_spectral, z_advect, vpa_advect,
-                                  scratch_dummy, collisions, composition, num_diss_params, dt)
+function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
+                                  vz, z_spectral, vpa_spectral, z_advect, vpa_advect,
+                                  scratch_dummy, collisions, composition,
+                                  external_source_settings, num_diss_params, dt, io_input,
+                                  input_dict)
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
         begin_serial_region()
         @serial_region begin
+            if false && isfile(initial_electron_output_filename)
+            else
+                previous_runs_info = nothing
+            end
+            # Setup I/O for initial electron state
+            io_initial_electron = setup_initial_electron_io(io_input, vz, vr, vzeta, vpa,
+                                                            vperp, z, r, composition,
+                                                            collisions,
+                                                            moments.evolve_density,
+                                                            moments.evolve_upar,
+                                                            moments.evolve_ppar,
+                                                            external_source_settings,
+                                                            input_dict,
+                                                            previous_runs_info)
+
             # update the electron pdf in the fvec struct
             fvec.pdf_electron .= pdf.electron.norm
         end
@@ -516,12 +539,25 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, z, vpa, vperp, z_spec
         max_electron_pdf_iterations = 2000000
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
-        return @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
+        result = @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
                                     moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
                                     phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
-                                    moments.electron.dqpar_dz, moments.electron.dvth_dz, z, vpa, z_spectral, 
-                                    vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
-                                    num_diss_params, max_electron_pdf_iterations)
+                                    moments.electron.dqpar_dz, moments.electron.dvth_dz,
+                                    r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect,
+                                    vpa_advect, scratch_dummy, dt, collisions,
+                                    composition, num_diss_params,
+                                    max_electron_pdf_iterations;
+                                    io_initial_electron=io_initial_electron)
+
+        # Write the converged initial state for the electrons to a file so that it can be
+        # re-used if the simulation is re-run.
+        n_debug_outputs = output_counter = result[end]
+        t_idx = n_debug_outputs+1
+        write_electron_dfns_data_to_binary(pdf.electron.norm, io_initial_electron, t_idx,
+                                           r, z, vperp, vpa)
+        write_electron_moments_data_to_binary(moments, io_initial_electron, t_idx, r, z)
+
+        return result
     end
 end
 
diff --git a/src/moment_kinetics.jl b/src/moment_kinetics.jl
index 5483d87b5..8a64acdbc 100644
--- a/src/moment_kinetics.jl
+++ b/src/moment_kinetics.jl
@@ -343,7 +343,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
                               z_spectral, r_spectral, vpa_spectral, vz_spectral, species,
                               collisions, external_source_settings,
                               manufactured_solns_input, scratch_dummy, scratch, t_input,
-                              num_diss_params, advection_structs)
+                              num_diss_params, advection_structs, io_input, input_dict)
         # initialize time variable
         code_time = 0.
         previous_runs_info = nothing

From 9baafe8144e68730c37e827dcbc07f496b64dac8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 12:20:40 +0000
Subject: [PATCH 099/428] Correct docstring for `get_run_info_no_setup()`

The `do_setup` and `setup_input_file` arguments do not exist.
---
 moment_kinetics/src/load_data.jl | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index cb4f29e06..fb13e6f50 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -2280,10 +2280,9 @@ function ilocal_func(iglobal,irank,nlocal)
 end
 
 """
-    get_run_info_no_setup(run_dir...; itime_min=1, itime_max=0, itime_skip=1, dfns=false,
-                          do_setup=true, setup_input_file=nothing)
+    get_run_info_no_setup(run_dir...; itime_min=1, itime_max=0, itime_skip=1, dfns=false)
     get_run_info_no_setup((run_dir, restart_index)...; itime_min=1, itime_max=0,
-                          itime_skip=1, dfns=false, do_setup=true, setup_input_file=nothing)
+                          itime_skip=1, dfns=false)
 
 Get file handles and other info for a single run
 

From b8b433ec73e678e7003fc6896741f7ee5b403500 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 12:33:38 +0000
Subject: [PATCH 100/428] Fix I/O of initial electron state for parallel runs

---
 moment_kinetics/src/file_io.jl            | 16 ++++++++++++++++
 moment_kinetics/src/initial_conditions.jl |  4 +++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 9f16153cc..ac19dda1d 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1833,6 +1833,22 @@ function finish_file_io(ascii_io::Union{ascii_ios,Nothing},
     return nothing
 end
 
+"""
+close output files for electron initialization
+"""
+function finish_initial_electron_io(
+        binary_initial_electron::Union{io_initial_electron_info,Tuple,Nothing})
+
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if binary_initial_electron !== nothing && !isa(binary_initial_electron, Tuple)
+            close(binary_initial_electron.fid)
+        end
+    end
+    return nothing
+end
+
 # Include the non-optional implementations of binary I/O functions
 include("file_io_hdf5.jl")
 
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 23e6e58d4..431277000 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -25,7 +25,7 @@ using ..interpolation: interpolate_to_grid_1d!
 using ..looping
 using ..em_fields: update_phi!
 using ..file_io: setup_initial_electron_io, write_electron_dfns_data_to_binary,
-                 write_electron_moments_data_to_binary
+                 write_electron_moments_data_to_binary, finish_initial_electron_io
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
@@ -503,6 +503,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
         begin_serial_region()
+        io_initial_electron = nothing
         @serial_region begin
             if false && isfile(initial_electron_output_filename)
             else
@@ -556,6 +557,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         write_electron_dfns_data_to_binary(pdf.electron.norm, io_initial_electron, t_idx,
                                            r, z, vperp, vpa)
         write_electron_moments_data_to_binary(moments, io_initial_electron, t_idx, r, z)
+        finish_initial_electron_io(io_initial_electron)
 
         return result
     end

From 2ace23f17cf241e71cad075ccd62b322269e3374 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 12:35:28 +0000
Subject: [PATCH 101/428] Remove 's' from filename for output of initial
 electron state

---
 moment_kinetics/src/file_io.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index ac19dda1d..8117a84f1 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -274,7 +274,7 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
         # state
         restart_time_index = -1
 
-        electrons_prefix = string(out_prefix, ".initial_electrons")
+        electrons_prefix = string(out_prefix, ".initial_electron")
         if !parallel_io
             electrons_prefix *= ".$(iblock_index[])"
         end

From a5f4fd306d4aa500ead8554360cfe0c072bde57b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 15:12:24 +0000
Subject: [PATCH 102/428] Allow coordinate to be created when not using MPI

When not using MPI (i.e. in post-processing) allocate `scratch_shared`
and `scratch_shared_2` as regular arrays, not shared memory ones, as
shared-memory arrays cannot be allocated using a null communicator.
---
 moment_kinetics/src/coordinates.jl            | 11 +++-
 moment_kinetics/src/moment_kinetics_input.jl  | 25 +++++---
 moment_kinetics/test/calculus_tests.jl        | 60 ++++++++++++++-----
 moment_kinetics/test/interpolation_tests.jl   |  4 +-
 .../test/velocity_integral_tests.jl           | 10 ++--
 5 files changed, 79 insertions(+), 31 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index a5224b14e..d2ef23b94 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -113,7 +113,7 @@ create arrays associated with a given coordinate,
 setup the coordinate grid, and populate the coordinate structure
 containing all of this information
 """
-function define_coordinate(input, parallel_io::Bool=false; init_YY::Bool=true)
+function define_coordinate(input, parallel_io::Bool=false; ignore_MPI=false, init_YY::Bool=true)
     # total number of grid points is ngrid for the first element
     # plus ngrid-1 unique points for each additional element due
     # to the repetition of a point at the element boundary
@@ -143,8 +143,13 @@ function define_coordinate(input, parallel_io::Bool=false; init_YY::Bool=true)
     duniform_dgrid = allocate_float(input.ngrid, input.nelement_local)
     # scratch is an array used for intermediate calculations requiring n entries
     scratch = allocate_float(n_local)
-    scratch_shared = allocate_shared_float(n_local)
-    scratch_shared2 = allocate_shared_float(n_local)
+    if ignore_MPI
+        scratch_shared = allocate_float(n_local)
+        scratch_shared2 = allocate_float(n_local)
+    else
+        scratch_shared = allocate_shared_float(n_local)
+        scratch_shared2 = allocate_shared_float(n_local)
+    end
     # scratch_2d is an array used for intermediate calculations requiring ngrid x nelement entries
     scratch_2d = allocate_float(input.ngrid, input.nelement_local)
     # struct containing the advection speed options/inputs for this coordinate
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index e8a4c13c7..660409d53 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -536,21 +536,30 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
                               Dict(Symbol(k)=>v for (k,v) in io_settings)...)
 
     # initialize z grid and write grid point locations to file
-    z, z_spectral = define_coordinate(z_immutable, io_immutable.parallel_io)
+    z, z_spectral = define_coordinate(z_immutable, io_immutable.parallel_io;
+                                      ignore_MPI=ignore_MPI)
     # initialize r grid and write grid point locations to file
-    r, r_spectral = define_coordinate(r_immutable, io_immutable.parallel_io)
+    r, r_spectral = define_coordinate(r_immutable, io_immutable.parallel_io;
+                                      ignore_MPI=ignore_MPI)
     # initialize vpa grid and write grid point locations to file
-    vpa, vpa_spectral = define_coordinate(vpa_immutable, io_immutable.parallel_io)
+    vpa, vpa_spectral = define_coordinate(vpa_immutable, io_immutable.parallel_io;
+                                          ignore_MPI=ignore_MPI)
     # initialize vperp grid and write grid point locations to file
-    vperp, vperp_spectral = define_coordinate(vperp_immutable, io_immutable.parallel_io)
+    vperp, vperp_spectral = define_coordinate(vperp_immutable, io_immutable.parallel_io;
+                                              ignore_MPI=ignore_MPI)
     # initialize gyrophase grid and write grid point locations to file
-    gyrophase, gyrophase_spectral = define_coordinate(gyrophase_immutable, io_immutable.parallel_io)
+    gyrophase, gyrophase_spectral = define_coordinate(gyrophase_immutable,
+                                                      io_immutable.parallel_io;
+                                                      ignore_MPI=ignore_MPI)
     # initialize vz grid and write grid point locations to file
-    vz, vz_spectral = define_coordinate(vz_immutable, io_immutable.parallel_io)
+    vz, vz_spectral = define_coordinate(vz_immutable, io_immutable.parallel_io;
+                                        ignore_MPI=ignore_MPI)
     # initialize vr grid and write grid point locations to file
-    vr, vr_spectral = define_coordinate(vr_immutable, io_immutable.parallel_io)
+    vr, vr_spectral = define_coordinate(vr_immutable, io_immutable.parallel_io;
+                                        ignore_MPI=ignore_MPI)
     # initialize vr grid and write grid point locations to file
-    vzeta, vzeta_spectral = define_coordinate(vzeta_immutable, io_immutable.parallel_io)
+    vzeta, vzeta_spectral = define_coordinate(vzeta_immutable, io_immutable.parallel_io;
+                                              ignore_MPI=ignore_MPI)
 
     external_source_settings = setup_external_sources!(scan_input, r, z)
 
diff --git a/moment_kinetics/test/calculus_tests.jl b/moment_kinetics/test/calculus_tests.jl
index be8ac68ac..9be3ca7c1 100644
--- a/moment_kinetics/test/calculus_tests.jl
+++ b/moment_kinetics/test/calculus_tests.jl
@@ -43,7 +43,9 @@ function runtests()
                     discretization, fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x'
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
                 # create array for the function f(x) to be differentiated/integrated
                 f = Array{Float64,1}(undef, x.n)
                 # create array for the derivative df/dx
@@ -93,7 +95,9 @@ function runtests()
                     "finite_difference", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x'
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 # create array for the derivative df/dx and the expected result
                 df = Array{Float64,1}(undef, x.n)
@@ -143,7 +147,9 @@ function runtests()
                     "finite_difference", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x'
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 # create array for the derivative df/dx and the expected result
                 df = Array{Float64,1}(undef, x.n)
@@ -189,7 +195,9 @@ function runtests()
                     "finite_difference", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x'
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 # create array for the derivative df/dx and the expected result
                 df = Array{Float64,1}(undef, x.n)
@@ -243,7 +251,9 @@ function runtests()
                     "finite_difference", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                # create the coordinate struct 'x'
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 # create array for the derivative df/dx and the expected result
                 df = Array{Float64,1}(undef, x.n)
@@ -459,7 +469,9 @@ function runtests()
                     "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
@@ -655,7 +667,9 @@ function runtests()
                     "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
@@ -698,7 +712,9 @@ function runtests()
                     "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
                 # test polynomials up to order ngrid-1
                 for n ∈ 0:ngrid-1
                     # create array for the function f(x) to be differentiated/integrated
@@ -748,7 +764,9 @@ function runtests()
                     "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
                 # test polynomials up to order ngrid-1
                 for n ∈ 0:ngrid-1
                     # create array for the function f(x) to be differentiated/integrated
@@ -884,7 +902,9 @@ function runtests()
                     "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input,init_YY=false)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; init_YY=false, ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
@@ -1001,7 +1021,9 @@ function runtests()
                     "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input,init_YY=false)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; init_YY=false, ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
@@ -1045,7 +1067,9 @@ function runtests()
                     "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input,init_YY=false)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; init_YY=false, ignore_MPI=true)
                 # test polynomials up to order ngrid-1
                 for n ∈ 0:ngrid-1
                     # create array for the function f(x) to be differentiated/integrated
@@ -1096,7 +1120,9 @@ function runtests()
                     "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input,init_YY=false)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; init_YY=false, ignore_MPI=true)
                 # test polynomials up to order ngrid-1
                 for n ∈ 0:ngrid-1
                     # create array for the function f(x) to be differentiated/integrated
@@ -1310,7 +1336,9 @@ function runtests()
                     "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
@@ -1418,7 +1446,9 @@ function runtests()
                     "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
                     element_spacing_option)
                 # create the coordinate struct 'x' and info for derivatives, etc.
-                x, spectral = define_coordinate(input,init_YY=false)
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; init_YY=false, ignore_MPI=true)
 
                 offset = randn(rng)
                 f = @. sinpi(2.0 * x.grid / L) + offset
diff --git a/moment_kinetics/test/interpolation_tests.jl b/moment_kinetics/test/interpolation_tests.jl
index 1951183a1..9b88deda7 100644
--- a/moment_kinetics/test/interpolation_tests.jl
+++ b/moment_kinetics/test/interpolation_tests.jl
@@ -44,7 +44,9 @@ function runtests()
                                 discretization, fd_option, cheb_option, bc, adv_input, comm,
                                 element_spacing_option)
             # create the coordinate struct 'z'
-            z, spectral = define_coordinate(input)
+            # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+            # errors due to communicators not being fully set up.
+            z, spectral = define_coordinate(input; ignore_MPI=true)
 
             test_grid = [z for z in range(-zlim, zlim, length=ntest)]
 
diff --git a/moment_kinetics/test/velocity_integral_tests.jl b/moment_kinetics/test/velocity_integral_tests.jl
index 2ceaa7e86..365f92905 100644
--- a/moment_kinetics/test/velocity_integral_tests.jl
+++ b/moment_kinetics/test/velocity_integral_tests.jl
@@ -45,10 +45,12 @@ function runtests()
                                  irank, Lvperp, discretization, fd_option, cheb_option, bc, adv_input,
                                  comm, "uniform")
         # create the coordinate struct 'x'
-        vpa, vpa_spectral = define_coordinate(vpa_input)
-        vperp, vperp_spectral = define_coordinate(vperp_input)
-        vz, vz_spectral = define_coordinate(vz_input)
-        vr, vr_spectral = define_coordinate(vr_input)
+        # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+        # errors due to communicators not being fully set up.
+        vpa, vpa_spectral = define_coordinate(vpa_input; ignore_MPI=true)
+        vperp, vperp_spectral = define_coordinate(vperp_input; ignore_MPI=true)
+        vz, vz_spectral = define_coordinate(vz_input; ignore_MPI=true)
+        vr, vr_spectral = define_coordinate(vr_input; ignore_MPI=true)
 
         dfn = allocate_float(vpa.n,vperp.n)
         dfn1D = allocate_float(vz.n, vr.n)

From c14a936a4d914c249ec87444e4e3d54f6be262b3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 12:38:09 +0000
Subject: [PATCH 103/428] Support loading from *.initial_electron.* files in
 get_run_info()

---
 .../src/makie_post_processing.jl              |   9 +-
 moment_kinetics/src/load_data.jl              | 247 +++++++++++++++++-
 2 files changed, 246 insertions(+), 10 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index c7b4fb850..b76710d7d 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -721,9 +721,11 @@ end
 
 """
     get_run_info(run_dir...; itime_min=1, itime_max=0,
-                 itime_skip=1, dfns=false, do_setup=true, setup_input_file=nothing)
+                 itime_skip=1, dfns=false, initial_electron=false, do_setup=true,
+                 setup_input_file=nothing)
     get_run_info((run_dir, restart_index)...; itime_min=1, itime_max=0,
-                 itime_skip=1, dfns=false, do_setup=true, setup_input_file=nothing)
+                 itime_skip=1, dfns=false, initial_electron=false, do_setup=true,
+                 setup_input_file=nothing)
 
 Get file handles and other info for a single run
 
@@ -741,7 +743,8 @@ argument can be a String `run_dir` giving a directory to read output from or a T
 mix Strings and Tuples in a call).
 
 By default load data from moments files, pass `dfns=true` to load from distribution
-functions files.
+functions files, or `initial_electron=true` and `dfns=true` to load from initial electron
+state files.
 
 The `itime_min`, `itime_max` and `itime_skip` options can be used to select only a slice
 of time points when loading data. In `makie_post_process` these options are read from the
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index fb13e6f50..ae7066fda 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -30,15 +30,21 @@ const em_variables = ("phi", "Er", "Ez")
 const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "thermal_speed", "temperature", "parallel_heat_flux",
                               "collision_frequency", "sound_speed", "mach_number")
+const electron_moment_variables = ("electron_density", "electron_parallel_flow",
+                                   "electron_parallel_pressure", "electron_thermal_speed",
+                                   "electron_temperature", "electron_parallel_heat_flux")
 const neutral_moment_variables = ("density_neutral", "uz_neutral", "pz_neutral",
                                   "thermal_speed_neutral", "temperature_neutral",
                                   "qz_neutral")
 const all_moment_variables = tuple(em_variables..., ion_moment_variables...,
+                                   electron_moment_variables...,
                                    neutral_moment_variables...)
 
 const ion_dfn_variables = ("f",)
+const electron_dfn_variables = ("f_electron",)
 const neutral_dfn_variables = ("f_neutral",)
-const all_dfn_variables = tuple(ion_dfn_variables..., neutral_dfn_variables...)
+const all_dfn_variables = tuple(ion_dfn_variables..., electron_dfn_variables...,
+                                neutral_dfn_variables...)
 
 const ion_variables = tuple(ion_moment_variables..., ion_dfn_variables)
 const neutral_variables = tuple(neutral_moment_variables..., neutral_dfn_variables)
@@ -2045,6 +2051,202 @@ function load_distributed_ion_pdf_slice(run_names::Tuple, nblocks::Tuple, t_rang
     return f_global
 end
 
+"""
+Read a slice of an electron distribution function
+
+run_names is a tuple. If it has more than one entry, this means that there are multiple
+restarts (which are sequential in time), so concatenate the data from each entry together.
+
+The slice to take is specified by the keyword arguments.
+"""
+function load_distributed_electron_pdf_slice(run_names::Tuple, nblocks::Tuple, t_range,
+                                             n_species::mk_int, r::coordinate,
+                                             z::coordinate, vperp::coordinate,
+                                             vpa::coordinate; ir=nothing, iz=nothing,
+                                             ivperp=nothing, ivpa=nothing)
+    # dimension of pdf is [vpa,vperp,z,r,t]
+
+    result_dims = mk_int[]
+    if ivpa === nothing
+        ivpa = 1:vpa.n_global
+        push!(result_dims, vpa.n_global)
+    elseif !isa(ivpa, mk_int)
+        push!(result_dims, length(ivpa))
+    end
+    if ivperp === nothing
+        ivperp = 1:vperp.n_global
+        push!(result_dims, vperp.n_global)
+    elseif !isa(ivperp, mk_int)
+        push!(result_dims, length(ivperp))
+    end
+    if iz === nothing
+        iz = 1:z.n_global
+        push!(result_dims, z.n_global)
+    elseif isa(iz, mk_int)
+        push!(result_dims, 1)
+    else
+        push!(result_dims, length(iz))
+    end
+    if ir === nothing
+        ir = 1:r.n_global
+        push!(result_dims, r.n_global)
+    elseif isa(ir, mk_int)
+        push!(result_dims, 1)
+    else
+        push!(result_dims, length(ir))
+    end
+    push!(result_dims, length(t_range))
+
+    f_global = allocate_float(result_dims...)
+
+    local_tind_start = 1
+    local_tind_end = -1
+    global_tind_start = 1
+    global_tind_end = -1
+    for (run_name, nb) in zip(run_names, nblocks)
+        for iblock in 0:nb-1
+            fid = open_readonly_output_file(run_name, "dfns", iblock=iblock, printout=false)
+
+            z_irank, z_nrank, r_irank, r_nrank = load_rank_data(fid)
+
+            # max index set to avoid double assignment of repeated points
+            # nr/nz if irank = nrank-1, (nr-1)/(nz-1) otherwise
+            imax_r = (r_irank == r.nrank - 1 ? r.n : r.n - 1)
+            imax_z = (z_irank == z.nrank - 1 ? z.n : z.n - 1)
+            local_r_range = 1:imax_r
+            local_z_range = 1:imax_z
+            global_r_range = iglobal_func(1, r_irank, r.n):iglobal_func(imax_r, r_irank, r.n)
+            global_z_range = iglobal_func(1, z_irank, z.n):iglobal_func(imax_z, z_irank, z.n)
+
+            if ir !== nothing && !any(i ∈ global_r_range for i in ir)
+                # No data for the slice on this rank
+                continue
+            elseif isa(ir, StepRange)
+                # Note that `findfirst(test, array)` returns the index `i` of the first
+                # element of `array` for which `test(array[i])` is `true`.
+                # `findlast()` similarly finds the index of the last element...
+                start_ind = findfirst(i -> i>=ir.start, global_r_range)
+                start = global_r_range[start_ind]
+                stop_ind = findlast(i -> i<=ir.stop, global_r_range)
+                stop = global_r_range[stop_ind]
+                local_r_range = (local_r_range.start + start - global_r_range.start):ir.step:(local_r_range.stop + stop - global_r_range.stop)
+                global_r_range = findfirst(i->i ∈ global_r_range, ir):findlast(i->i ∈ global_r_range, ir)
+            elseif isa(ir, UnitRange)
+                start_ind = findfirst(i -> i>=ir.start, global_r_range)
+                start = global_r_range[start_ind]
+                stop_ind = findlast(i -> i<=ir.stop, global_r_range)
+                stop = global_r_range[stop_ind]
+                local_r_range = (local_r_range.start + start - global_r_range.start):(local_r_range.stop + stop - global_r_range.stop)
+                global_r_range = findfirst(i->i ∈ global_r_range, ir):findlast(i->i ∈ global_r_range, ir)
+            elseif isa(ir, mk_int)
+                local_r_range = ir - (global_r_range.start - 1)
+                global_r_range = ir
+            end
+            if iz !== nothing && !any(i ∈ global_z_range for i in iz)
+                # No data for the slice on this rank
+                continue
+            elseif isa(iz, StepRange)
+                # Note that `findfirst(test, array)` returns the index `i` of the first
+                # element of `array` for which `test(array[i])` is `true`.
+                # `findlast()` similarly finds the index of the last element...
+                start_ind = findfirst(i -> i>=iz.start, global_z_range)
+                start = global_z_range[start_ind]
+                stop_ind = findlast(i -> i<=iz.stop, global_z_range)
+                stop = global_z_range[stop_ind]
+                local_z_range = (local_z_range.start + start - global_z_range.start):iz.step:(local_z_range.stop + stop - global_z_range.stop)
+                global_z_range = findfirst(i->i ∈ global_z_range, iz):findlast(i->i ∈ global_z_range, iz)
+            elseif isa(iz, UnitRange)
+                start_ind = findfirst(i -> i>=iz.start, global_z_range)
+                start = global_z_range[start_ind]
+                stop_ind = findlast(i -> i<=iz.stop, global_z_range)
+                stop = global_z_range[stop_ind]
+                local_z_range = (local_z_range.start + start - global_z_range.start):(local_z_range.stop + stop - global_z_range.stop)
+                global_z_range = findfirst(i->i ∈ global_z_range, iz):findlast(i->i ∈ global_z_range, iz)
+            elseif isa(iz, mk_int)
+                local_z_range = iz - (global_z_range.start - 1)
+                global_z_range = iz
+            end
+
+            f_local_slice = load_pdf_data(fid)
+
+            if local_tind_start > 1
+                # The run being loaded is a restart (as local_tind_start=1 for the first
+                # run), so skip the first point, as this is a duplicate of the last point
+                # of the previous restart
+                skip_first = 1
+            else
+                skip_first = 0
+            end
+            ntime_local = size(f_local_slice, ndims(f_local_slice)) - skip_first
+            local_tind_end = local_tind_start + ntime_local - 1
+            local_t_range = collect(it - local_tind_start + 1 + skip_first
+                                    for it ∈ t_range
+                                    if local_tind_start <= it <= local_tind_end)
+            global_tind_end = global_tind_start + length(local_t_range) - 1
+
+            f_global_slice = selectdim(f_global, ndims(f_global),
+                                       global_tind_start:global_tind_end)
+
+            # Note: use selectdim() and get the dimension from thisdim because the actual
+            # number of dimensions in f_global_slice, f_local_slice is different depending
+            # on which combination of ivpa, ivperp, iz, ir, and is was passed.
+            thisdim = ndims(f_local_slice) - 5
+            f_local_slice = selectdim(f_local_slice, thisdim, ivpa)
+
+            thisdim = ndims(f_local_slice) - 4
+            f_local_slice = selectdim(f_local_slice, thisdim, ivperp)
+
+            thisdim = ndims(f_local_slice) - 3
+            if isa(iz, mk_int)
+                f_global_slice = selectdim(f_global_slice, thisdim, 1)
+                f_local_slice = selectdim(f_local_slice, thisdim,
+                                          ilocal_func(iz, z_irank, z.n))
+            else
+                f_global_slice = selectdim(f_global_slice, thisdim, global_z_range)
+                f_local_slice = selectdim(f_local_slice, thisdim, local_z_range)
+            end
+
+            thisdim = ndims(f_local_slice) - 2
+            if isa(ir, mk_int)
+                f_global_slice = selectdim(f_global_slice, thisdim, 1)
+                f_local_slice = selectdim(f_local_slice, thisdim,
+                                          ilocal_func(ir, r_irank, r.n))
+            else
+                f_global_slice = selectdim(f_global_slice, thisdim, global_r_range)
+                f_local_slice = selectdim(f_local_slice, thisdim, local_r_range)
+            end
+
+            thisdim = ndims(f_local_slice) - 1
+            f_global_slice = selectdim(f_global_slice, thisdim)
+            f_local_slice = selectdim(f_local_slice, thisdim)
+
+            # Select time slice
+            thisdim = ndims(f_local_slice)
+            f_local_slice = selectdim(f_local_slice, thisdim, local_t_range)
+
+            f_global_slice .= f_local_slice
+            close(fid)
+        end
+        local_tind_start = local_tind_end + 1
+        global_tind_start = global_tind_end + 1
+    end
+
+    if isa(iz, mk_int)
+        thisdim = ndims(f_global) - 3
+        f_global = selectdim(f_global, thisdim, 1)
+    end
+    if isa(ir, mk_int)
+        thisdim = ndims(f_global) - 2
+        f_global = selectdim(f_global, thisdim, 1)
+    end
+    if isa(t_range, mk_int)
+        thisdim = ndims(f_global)
+        f_global = selectdim(f_global, thisdim, 1)
+    end
+
+    return f_global
+end
+
 """
 Read a slice of a neutral distribution function
 
@@ -2280,9 +2482,10 @@ function ilocal_func(iglobal,irank,nlocal)
 end
 
 """
-    get_run_info_no_setup(run_dir...; itime_min=1, itime_max=0, itime_skip=1, dfns=false)
+    get_run_info_no_setup(run_dir...; itime_min=1, itime_max=0, itime_skip=1, dfns=false,
+                          initial_electron=false)
     get_run_info_no_setup((run_dir, restart_index)...; itime_min=1, itime_max=0,
-                          itime_skip=1, dfns=false)
+                          itime_skip=1, dfns=false, initial_electron=false)
 
 Get file handles and other info for a single run
 
@@ -2300,7 +2503,8 @@ argument can be a String `run_dir` giving a directory to read output from or a T
 mix Strings and Tuples in a call).
 
 By default load data from moments files, pass `dfns=true` to load from distribution
-functions files.
+functions files, or `initial_electron=true` and `dfns=true` to load from initial electron
+state files.
 
 The `itime_min`, `itime_max` and `itime_skip` options can be used to select only a slice
 of time points when loading data. In `makie_post_process` these options are read from the
@@ -2311,14 +2515,19 @@ defaults for the remaining options. If either `itime_min` or `itime_max` are ≤
 values are used as offsets from the final time index of the run.
 """
 function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractString,Union{Int,Nothing}}}...;
-                               itime_min=1, itime_max=0, itime_skip=1, dfns=false)
+                               itime_min=1, itime_max=0, itime_skip=1, dfns=false,
+                               initial_electron=false)
     if length(run_dir) == 0
         error("No run_dir passed")
     end
+    if initial_electron && !dfns
+        error("When `initial_electron=true` is passed, `dfns=true` must also be passed")
+    end
     if length(run_dir) > 1
         run_info = Tuple(get_run_info_no_setup(r; itime_min=itime_min,
                                                itime_max=itime_max, itime_skip=itime_skip,
-                                               dfns=dfns)
+                                               dfns=dfns,
+                                               initial_electron=initial_electron)
                          for r ∈ run_dir)
         return run_info
     end
@@ -2377,7 +2586,9 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
         error("Invalid restart_index=$restart_index")
     end
 
-    if dfns
+    if initial_electron
+        ext = "initial_electron"
+    elseif dfns
         ext = "dfns"
     else
         ext = "moments"
@@ -2688,6 +2899,15 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
             end
             !isa(it, mk_int) && push!(dims, nt)
             result = allocate_float(dims...)
+        elseif nd == 5
+            # electron distribution function variable with dimensions (vpa,vperp,z,r,t)
+            dims = Vector{mk_int}()
+            !isa(ivpa, mk_int) && push!(dims, nvpa)
+            !isa(ivperp, mk_int) && push!(dims, nvperp)
+            !isa(iz, mk_int) && push!(dims, nz)
+            !isa(ir, mk_int) && push!(dims, nr)
+            !isa(it, mk_int) && push!(dims, nt)
+            result = allocate_float(dims...)
         elseif nd == 6
             # ion distribution function variable with dimensions (vpa,vperp,z,r,s,t)
             nspecies = size(variable[1], 5)
@@ -2745,6 +2965,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                         result .= v[iz,ir,tind]
                     elseif nd == 4
                         result .= v[iz,ir,is,tind]
+                    elseif nd == 5
+                        result .= v[ivpa,ivperp,iz,ir,tind]
                     elseif nd == 6
                         result .= v[ivpa,ivperp,iz,ir,is,tind]
                     elseif nd == 7
@@ -2775,6 +2997,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                         selectdim(result, ndims(result), global_it_start:global_it_end) .= v[iz,ir,tinds]
                     elseif nd == 4
                         selectdim(result, ndims(result), global_it_start:global_it_end) .= v[iz,ir,is,tinds]
+                    elseif nd == 5
+                        selectdim(result, ndims(result), global_it_start:global_it_end) .= v[ivpa,ivperp,iz,ir,tinds]
                     elseif nd == 6
                         selectdim(result, ndims(result), global_it_start:global_it_end) .= v[ivpa,ivperp,iz,ir,is,tinds]
                     elseif nd == 7
@@ -2794,6 +3018,8 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
         # Use existing distributed I/O loading functions
         if variable_name ∈ em_variables
             nd = 3
+        elseif variable_name ∈ electron_dfn_variables
+            nd = 5
         elseif variable_name ∈ ion_dfn_variables
             nd = 6
         elseif variable_name ∈ neutral_dfn_variables
@@ -2820,6 +3046,13 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                                       run_info.ext, run_info.nblocks, run_info.z_local.n,
                                       run_info.r_local.n, run_info.itime_skip)
             result = result[iz,ir,is,it]
+        elseif nd === 5
+            result = load_distributed_electron_pdf_slice(run_info.files, run_info.nblocks,
+                                                         it, run_info.n_ion_species,
+                                                         run_info.r_local,
+                                                         run_info.z_local, run_info.vperp,
+                                                         run_info.vpa; ir=ir, iz=iz,
+                                                         ivperp=ivperp, ivpa=ivpa)
         elseif nd === 6
             result = load_distributed_ion_pdf_slice(run_info.files, run_info.nblocks, it,
                                                     run_info.n_ion_species,

From 61ae463ac2d058fca865643e75ba9a4e8e4a5d42 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 14:09:47 +0000
Subject: [PATCH 104/428] Tidy up split I/O functions

---
 moment_kinetics/src/file_io.jl         | 168 +++++++++++--------------
 moment_kinetics/src/moment_kinetics.jl |   6 +-
 2 files changed, 74 insertions(+), 100 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 8117a84f1..3be192ade 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -146,7 +146,7 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_electron_initialization_info{Tfile, Tfe, Tmom}
+struct io_initial_electron_info{Tfile, Tfe, Tmom}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the electron distribution function variable
@@ -159,6 +159,8 @@ struct io_electron_initialization_info{Tfile, Tfe, Tmom}
     electron_parallel_pressure::Tmom
     # handle for the electron parallel heat flux variable
     electron_parallel_heat_flux::Tmom
+    # handle for the electron thermal speed variable
+    electron_thermal_speed::Tmom
 
     # Use parallel I/O?
     parallel_io::Bool
@@ -322,6 +324,36 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
     return nothing, nothing, nothing
 end
 
+"""
+Reopen an existing initial electron output file to append more data
+"""
+function reopen_initial_electron_io(file_info)
+    @serial_region begin
+        filename, parallel_io, io_comm = file_info
+        fid = reopen_output_file(filename, parallel_io, io_comm)
+        dyn = get_group(fid, "dynamic_data")
+
+        variable_list = get_variable_keys(dyn)
+        function getvar(name)
+            if name ∈ variable_list
+                return dyn[name]
+            else
+                return nothing
+            end
+        end
+        return io_initial_electron_info(fid, getvar("pseudotime"), getvar("f_electron"),
+                                        getvar("electron_density"),
+                                        getvar("electron_parallel_flow"),
+                                        getvar("electron_parallel_pressure"),
+                                        getvar("electron_parallel_heat_flux"),
+                                        getvar("electron_thermal_speed"), parallel_io)
+    end
+
+    # For processes other than the root process of each shared-memory group...
+    return nothing
+end
+
+
 """
 Get a (sub-)group from a file or group
 """
@@ -737,8 +769,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                            description="simulation time")
 
         io_phi, io_Er, io_Ez =
-            define_dynamic_em_field_variables!(fid, r, z, parallel_io,
-                                               external_source_settings)
+            define_dynamic_em_field_variables!(fid, r, z, parallel_io)
 
         io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
         external_source_amplitude, external_source_density_amplitude,
@@ -1441,16 +1472,14 @@ function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
         # add the time for this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.time, t, t_idx, parallel_io)
 
-        write_em_fields_data_to_binary(fields, io_or_file_info_moments, t_idx, r, z)
+        write_em_fields_data_to_binary(fields, io_moments, t_idx, r, z)
 
-        write_ion_moments_data_to_binary(moments, n_ion_species, io_or_file_info_moments,
-                                         t_idx, r, z)
+        write_ion_moments_data_to_binary(moments, n_ion_species, io_moments, t_idx, r, z)
 
-        write_electron_moments_data_to_binary(moments, io_or_file_info_moments, t_idx, r,
-                                              z)
+        write_electron_moments_data_to_binary(moments, io_moments, t_idx, r, z)
 
-        write_neutral_moments_data_to_binary(moments, n_neutral_species,
-                                             io_or_file_info_moments, t_idx, r, z)
+        write_neutral_moments_data_to_binary(moments, n_neutral_species, io_moments,
+                                             t_idx, r, z)
 
         append_to_dynamic_var(io_moments.time_for_run, time_for_run, t_idx, parallel_io)
 
@@ -1462,27 +1491,19 @@ end
 
 """
 write time-dependent EM fields data to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_em_fields_data_to_binary(fields, io_or_file_info_moments, t_idx, r, z)
+function write_em_fields_data_to_binary(fields, io_moments::io_moments_info, t_idx, r, z)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_moments, io_moments_info)
-            io_moments = io_or_file_info_moments
-            closefile = false
-        else
-            io_moments = reopen_moments_io(io_or_file_info_moments)
-            closefile = true
-        end
-
         parallel_io = io_moments.parallel_io
 
         # add the electrostatic potential and electric field components at this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.phi, fields.phi, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.Er, fields.Er, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.Ez, fields.Ez, t_idx, parallel_io, z, r)
-
-        closefile && close(io_moments.fid)
     end
 
     return nothing
@@ -1490,20 +1511,14 @@ end
 
 """
 write time-dependent moments data for ions to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_ion_moments_data_to_binary(moments, n_ion_species, io_or_file_info_moments,
-                                          t_idx, r, z)
+function write_ion_moments_data_to_binary(moments, n_ion_species,
+                                          io_moments::io_moments_info, t_idx, r, z)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_moments, io_moments_info)
-            io_moments = io_or_file_info_moments
-            closefile = false
-        else
-            io_moments = reopen_moments_io(io_or_file_info_moments)
-            closefile = true
-        end
-
         parallel_io = io_moments.parallel_io
 
         # add the density data at this time slice to the output file
@@ -1570,8 +1585,6 @@ function write_ion_moments_data_to_binary(moments, n_ion_species, io_or_file_inf
                                       t_idx, parallel_io, z, r)
             end
         end
-
-        closefile && close(io_moments.fid)
     end
 
     return nothing
@@ -1579,20 +1592,15 @@ end
 
 """
 write time-dependent moments data for electrons to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_electron_moments_data_to_binary(moments, io_or_file_info_moments, t_idx, r,
-                                               z)
+function write_electron_moments_data_to_binary(moments,
+                                               io_moments::Union{io_moments_info,io_initial_electron_info},
+                                               t_idx, r, z)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_moments, io_moments_info)
-            io_moments = io_or_file_info_moments
-            closefile = false
-        else
-            io_moments = reopen_moments_io(io_or_file_info_moments)
-            closefile = true
-        end
-
         parallel_io = io_moments.parallel_io
 
         append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx,
@@ -1605,8 +1613,6 @@ function write_electron_moments_data_to_binary(moments, io_or_file_info_moments,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
                               t_idx, parallel_io, z, r)
-
-        closefile && close(io_moments.fid)
     end
 
     return nothing
@@ -1614,9 +1620,11 @@ end
 
 """
 write time-dependent moments data for neutrals to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
 function write_neutral_moments_data_to_binary(moments, n_neutral_species,
-                                              io_or_file_info_moments, t_idx, r, z)
+                                              io_moments::io_moments_info, t_idx, r, z)
     if n_neutral_species ≤ 0
         return nothing
     end
@@ -1624,14 +1632,6 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_moments, io_moments_info)
-            io_moments = io_or_file_info_moments
-            closefile = false
-        else
-            io_moments = reopen_moments_io(io_or_file_info_moments)
-            closefile = true
-        end
-
         parallel_io = io_moments.parallel_io
 
         append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
@@ -1676,8 +1676,6 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
                                       t_idx, parallel_io, z, r)
             end
         end
-
-        closefile && close(io_moments.fid)
     end
 
     return nothing
@@ -1708,13 +1706,12 @@ function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
                                          time_for_run, r, z)
 
         # add the distribution function data at this time slice to the output file
-        write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_or_file_info_dfns,
-                                      t_idx, r, z, vperp, vpa)
-        write_electron_dfns_data_to_binary(pdf.electron.norm, io_or_file_info_dfns, t_idx,
-                                           r, z, vperp, vpa)
-        write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species,
-                                          io_or_file_info_dfns, t_idx, r, z, vzeta, vr,
-                                          vz)
+        write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_dfns, t_idx, r, z,
+                                      vperp, vpa)
+        write_electron_dfns_data_to_binary(pdf.electron.norm, io_dfns, t_idx, r, z, vperp,
+                                           vpa)
+        write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species, io_dfns,
+                                          t_idx, r, z, vzeta, vr, vz)
 
         closefile && close(io_dfns.fid)
     end
@@ -1723,83 +1720,60 @@ end
 
 """
 write time-dependent distribution function data for ions to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_ion_dfns_data_to_binary(ff, n_ion_species, io_or_file_info_dfns, t_idx, r,
+function write_ion_dfns_data_to_binary(ff, n_ion_species, io_dfns::io_dfns_info, t_idx, r,
                                        z, vperp, vpa)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_dfns, io_dfns_info)
-            io_dfns = io_or_file_info_dfns
-            closefile = false
-        else
-            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
-            closefile = true
-        end
-
         parallel_io = io_dfns.parallel_io
 
         append_to_dynamic_var(io_dfns.f, ff, t_idx, parallel_io, vpa, vperp, z, r,
                               n_ion_species)
-
-        closefile && close(io_dfns.fid)
     end
     return nothing
 end
 
 """
 write time-dependent distribution function data for electrons to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_electron_dfns_data_to_binary(ff_electron, io_or_file_info_dfns, t_idx, r,
-                                            z, vperp, vpa)
+function write_electron_dfns_data_to_binary(ff_electron,
+                                            io_dfns::Union{io_dfns_info,io_initial_electron_info},
+                                            t_idx, r, z, vperp, vpa)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_dfns, io_dfns_info)
-            io_dfns = io_or_file_info_dfns
-            closefile = false
-        else
-            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
-            closefile = true
-        end
-
         parallel_io = io_dfns.parallel_io
 
         if io_dfns.f_electron !== nothing
             append_to_dynamic_var(io_dfns.f_electron, ff_electron, t_idx, parallel_io,
                                   vpa, vperp, z, r)
         end
-
-        closefile && close(io_dfns.fid)
     end
     return nothing
 end
 
 """
 write time-dependent distribution function data for neutrals to the binary output file
+
+Note: should only be called from within a function that (re-)opens the output file.
 """
 function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
-                                           io_or_file_info_dfns, t_idx, r, z, vzeta, vr,
+                                           io_dfns::io_dfns_info, t_idx, r, z, vzeta, vr,
                                            vz)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if isa(io_or_file_info_dfns, io_dfns_info)
-            io_dfns = io_or_file_info_dfns
-            closefile = false
-        else
-            io_dfns = reopen_dfns_io(io_or_file_info_dfns)
-            closefile = true
-        end
-
         parallel_io = io_dfns.parallel_io
 
         if n_neutral_species > 0
             append_to_dynamic_var(io_dfns.f_neutral, ff_neutral, t_idx, parallel_io, vz,
                                   vr, vzeta, z, r, n_neutral_species)
         end
-
-        closefile && close(io_dfns.fid)
     end
     return nothing
 end
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index e0baacdb8..7a570be21 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -448,9 +448,9 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     write_all_moments_data_to_binary(moments, fields, code_time,
         composition.n_ion_species, composition.n_neutral_species, io_moments, 1, 0.0, r,
         z)
-    write_all_dfns_data_to_binary(pdf.ion.norm, pdf.neutral.norm, moments, fields,
-        code_time, composition.n_ion_species, composition.n_neutral_species, io_dfns, 1,
-        0.0, r, z, vperp, vpa, vzeta, vr, vz)
+    write_all_dfns_data_to_binary(pdf, moments, fields, code_time,
+        composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, 0.0, r, z,
+        vperp, vpa, vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 

From c2eaf7706ba82fbf0aea8e8dbd6f4dc62194e28b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 18 Feb 2024 13:19:15 +0000
Subject: [PATCH 105/428] Write the pseudo-time to initial electron output file

---
 .../src/electron_kinetic_equation.jl          | 16 +++----
 moment_kinetics/src/file_io.jl                | 44 ++++++++++++++++++-
 moment_kinetics/src/initial_conditions.jl     | 10 ++---
 3 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 811ca1107..f6ad04277 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -17,7 +17,7 @@ using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!
 using ..electron_vpa_advection: electron_vpa_advection!
-using ..file_io: write_electron_dfns_data_to_binary, write_electron_moments_data_to_binary
+using ..file_io: write_initial_electron_state
 using ..moment_constraints: hard_force_moment_constraints!
 using ..velocity_moments: integrate_over_vspace
 
@@ -246,10 +246,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     @serial_region begin
         output_counter += 1
         if io_initial_electron !== nothing
-            write_electron_dfns_data_to_binary(pdf, io_initial_electron,
-                                               output_counter, r, z, vperp, vpa)
-            write_electron_moments_data_to_binary(moments, io_initial_electron,
-                                                  output_counter, r, z)
+            write_initial_electron_state(pdf, moments, time, io_initial_electron,
+                                         output_counter, r, z, vperp, vpa)
         end
         result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
         result_pdf[:,:,1] .= pdf[:,1,:,1]
@@ -350,10 +348,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 println(io_vth,"")
                 output_counter += 1
                 if io_initial_electron !== nothing
-                    write_electron_dfns_data_to_binary(pdf, io_initial_electron,
-                                                       output_counter, r, z, vperp, vpa)
-                    write_electron_moments_data_to_binary(moments, io_initial_electron,
-                                                          output_counter, r, z)
+                    write_initial_electron_state(pdf, moments, time, io_initial_electron,
+                                                 output_counter, r, z, vperp, vpa)
                 end
                 result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
                 result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
@@ -514,7 +510,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         result_phi = result_phi[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
         result_residual = result_residual[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
     end
-    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual, output_counter
+    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual, time, output_counter
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 3be192ade..011b1948f 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -146,9 +146,11 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_initial_electron_info{Tfile, Tfe, Tmom}
+struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom}
     # file identifier for the binary file to which data is written
     fid::Tfile
+    # handle for the pseudotime variable
+    pseudotime::Ttime
     # handle for the electron distribution function variable
     f_electron::Tfe
     # handle for the electron density variable
@@ -304,6 +306,8 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
 
         ### create variables for time-dependent quantities ###
         dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
+        io_pseudotime = create_dynamic_variable!(dynamic, "time", mk_float; parallel_io=parallel_io,
+                                                 description="pseudotime used for electron initialization")
         io_f_electron = create_dynamic_variable!(dynamic, "f_electron", mk_float, vpa,
                                                  vperp, z, r;
                                                  parallel_io=parallel_io,
@@ -341,7 +345,7 @@ function reopen_initial_electron_io(file_info)
                 return nothing
             end
         end
-        return io_initial_electron_info(fid, getvar("pseudotime"), getvar("f_electron"),
+        return io_initial_electron_info(fid, getvar("time"), getvar("f_electron"),
                                         getvar("electron_density"),
                                         getvar("electron_parallel_flow"),
                                         getvar("electron_parallel_pressure"),
@@ -1778,6 +1782,42 @@ function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
     return nothing
 end
 
+"""
+    write_initial_electron_state(pdf, moments, io_initial_electron, t_idx, r, z,
+                                 vperp, vpa)
+
+Write the electron state to an output file.
+"""
+function write_initial_electron_state(pdf, moments, t, io_or_file_info_initial_electron,
+                                      t_idx, r, z, vperp, vpa)
+
+    @serial_region begin
+        # Only read/write from first process in each 'block'
+
+        if isa(io_or_file_info_initial_electron, io_dfns_info)
+            io_initial_electron = io_or_file_info_initial_electron
+            closefile = false
+        else
+            io_initial_electron = reopen_initial_electron_io(io_or_file_info_initial_electron)
+            closefile = true
+        end
+
+        parallel_io = io_initial_electron.parallel_io
+
+        # add the pseudo-time for this time slice to the hdf5 file
+        append_to_dynamic_var(io_initial_electron.pseudotime, t, t_idx, parallel_io)
+
+        write_electron_dfns_data_to_binary(pdf, io_initial_electron, t_idx, r, z, vperp,
+                                           vpa)
+
+        write_electron_moments_data_to_binary(moments, io_initial_electron, t_idx, r, z)
+
+        closefile && close(io_initial_electron.fid)
+    end
+
+    return nothing
+end
+
 """
 close all opened output files
 """
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 431277000..7709f8a0d 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -24,8 +24,8 @@ using ..external_sources
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
 using ..em_fields: update_phi!
-using ..file_io: setup_initial_electron_io, write_electron_dfns_data_to_binary,
-                 write_electron_moments_data_to_binary, finish_initial_electron_io
+using ..file_io: setup_initial_electron_io, write_initial_electron_state,
+                 finish_initial_electron_io
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
@@ -553,10 +553,10 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
         n_debug_outputs = output_counter = result[end]
+        t = result[end-1]
         t_idx = n_debug_outputs+1
-        write_electron_dfns_data_to_binary(pdf.electron.norm, io_initial_electron, t_idx,
-                                           r, z, vperp, vpa)
-        write_electron_moments_data_to_binary(moments, io_initial_electron, t_idx, r, z)
+        write_initial_electron_state(pdf.electron.norm, moments, t, io_initial_electron,
+                                     t_idx, r, z, vperp, vpa)
         finish_initial_electron_io(io_initial_electron)
 
         return result

From 829ffa5fe45f204ea87529018441859c6a039828 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 3 Jan 2024 16:55:06 +0000
Subject: [PATCH 106/428] Get complete printing of stacktrace when exception
 occurs

In order to call MPI.Abort(), we have to do our own handling of errors
in run_moment_kinetics(). This commit changes the way the backtrace is
printed to avoid it being truncated, from the solution in
https://discourse.julialang.org/t/how-to-print-a-backtrace/74164/4
---
 moment_kinetics/src/moment_kinetics.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 7a570be21..04f1aa7c5 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -144,8 +144,7 @@ function run_moment_kinetics(to::Union{TimerOutput,Nothing}, input_dict=Dict();
         # throws an error
         if global_size[] > 1
             println("$(typeof(e)) on process $(global_rank[]):")
-            showerror(stdout, e)
-            display(stacktrace(catch_backtrace()))
+            showerror(stdout, e, catch_backtrace())
             flush(stdout)
             flush(stderr)
             MPI.Abort(comm_world, 1)

From 6acdfbf444c75aa3e7717febadb1740e5f4d144f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Feb 2024 11:39:22 +0000
Subject: [PATCH 107/428] DEBUGGING! Remove return of arrays for debugging of
 electrons

---
 .../src/electron_kinetic_equation.jl          | 60 +++----------------
 moment_kinetics/src/initial_conditions.jl     | 42 ++++++-------
 moment_kinetics/src/moment_kinetics.jl        |  6 +-
 3 files changed, 32 insertions(+), 76 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index f6ad04277..326ba800c 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -63,13 +63,15 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_u
             r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
             num_diss_params, max_electron_pdf_iterations; io_initial_electron=io_initial_electron)
     elseif solution_method == "shooting_method"
-        update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar, qpar_updated, 
-            phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy, composition)
+        return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
+            qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
+            vpa_spectral, scratch_dummy, composition)
     elseif solution_method == "picard_iteration"
-        update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz,
-            z, vpa, vpa_spectral, scratch_dummy, max_electron_pdf_iterations)
+        return update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens_dz,
+            dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy,
+            max_electron_pdf_iterations)
     else
-        println("!!! invalid solution method specified !!!")
+        error("!!! invalid solution method specified !!!")
     end
     return nothing    
 end
@@ -234,12 +236,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
     average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
 
-    result_pdf = nothing
-    result_ppar = nothing
-    result_vth = nothing
-    result_qpar = nothing
-    result_phi = nothing
-    result_residual = nothing
     output_interval = 1000
     output_counter = 0
     begin_serial_region()
@@ -249,17 +245,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             write_initial_electron_state(pdf, moments, time, io_initial_electron,
                                          output_counter, r, z, vperp, vpa)
         end
-        result_pdf = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
-        result_pdf[:,:,1] .= pdf[:,1,:,1]
-        result_ppar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-        result_ppar[:,1] .= ppar[:,1]
-        result_vth = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-        result_vth[:,1] .= vthe[:,1]
-        result_qpar = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-        result_qpar[:,1] .= qpar[:,1]
-        result_phi = zeros(z.n, max_electron_pdf_iterations ÷ output_interval)
-        result_phi[:,1] .= phi[:,1]
-        result_residual = zeros(vpa.n, z.n, max_electron_pdf_iterations ÷ output_interval)
     end
     # evolve (artificially) in time until the residual is less than the tolerance
     try
@@ -351,22 +336,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                     write_initial_electron_state(pdf, moments, time, io_initial_electron,
                                                  output_counter, r, z, vperp, vpa)
                 end
-                result_pdf[:,:,iteration÷output_interval+1] .= pdf[:,1,:,1]
-                result_ppar[:,iteration÷output_interval+1] .= ppar[:,1]
-                result_vth[:,iteration÷output_interval+1] .= vthe[:,1]
-                result_qpar[:,iteration÷output_interval+1] .= qpar[:,1]
-                result_phi[:,iteration÷output_interval+1] .= phi[:,1]
-            end
-        else
-            begin_serial_region()
-            @serial_region begin
-                if iteration÷output_interval+1+iteration%output_interval < size(result_pdf, 3)
-                    result_pdf[:,:,iteration÷output_interval+1+iteration%output_interval] .= pdf[:,1,:,1]
-                    result_ppar[:,iteration÷output_interval+1+iteration%output_interval] .= ppar[:,1]
-                    result_vth[:,iteration÷output_interval+1+iteration%output_interval] .= vthe[:,1]
-                    result_qpar[:,iteration÷output_interval+1+iteration%output_interval] .= qpar[:,1]
-                    result_phi[:,iteration÷output_interval+1+iteration%output_interval] .= phi[:,1]
-                end
             end
         end
 
@@ -449,12 +418,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
-        if (mod(iteration,output_interval) == 0)
-            begin_serial_region()
-            @serial_region begin
-                result_residual[:,:,iteration÷output_interval+1] .= residual[:,1,:,1]
-            end
-        end
 
         # Divide by wpa to relax CFL condition at large wpa - only looking for steady
         # state here, so does not matter that this makes time evolution incorrect.
@@ -502,15 +465,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         close(io_vth)
         close(io_pdf)
         close(io_pdf_stages)
-
-        result_pdf = result_pdf[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
-        result_ppar = result_ppar[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
-        result_vth = result_vth[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
-        result_qpar = result_qpar[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
-        result_phi = result_phi[:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
-        result_residual = result_residual[:,:,1:min(iteration÷output_interval+output_interval,max_electron_pdf_iterations÷output_interval)]
     end
-    return result_pdf, dens, upar, result_ppar, result_vth, result_qpar, result_phi, z, vpa, result_residual, time, output_counter
+    return time, output_counter
 end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 7709f8a0d..dbced58a5 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -368,12 +368,12 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
 
     # initialize the electron pdf that satisfies the electron kinetic equation
-    return initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa,
-                                    vperp, vzeta, vr, vz, z_spectral, vpa_spectral,
-                                    advection_structs.electron_z_advect,
-                                    advection_structs.electron_vpa_advect, scratch_dummy,
-                                    collisions, composition, external_source_settings,
-                                    num_diss_params, t_input.dt, io_input, input_dict)
+    initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
+                             vzeta, vr, vz, z_spectral, vpa_spectral,
+                             advection_structs.electron_z_advect,
+                             advection_structs.electron_vpa_advect, scratch_dummy,
+                             collisions, composition, external_source_settings,
+                             num_diss_params, t_input.dt, io_input, input_dict)
 
     return nothing
 end
@@ -540,27 +540,29 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         max_electron_pdf_iterations = 2000000
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
-        result = @views update_electron_pdf!(fvec, pdf.electron.norm, moments, moments.electron.dens, moments.electron.vth, 
-                                    moments.electron.ppar, moments.electron.qpar, moments.electron.qpar_updated,
-                                    phi, moments.electron.ddens_dz, moments.electron.dppar_dz, 
-                                    moments.electron.dqpar_dz, moments.electron.dvth_dz,
-                                    r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect,
-                                    vpa_advect, scratch_dummy, dt, collisions,
-                                    composition, num_diss_params,
-                                    max_electron_pdf_iterations;
-                                    io_initial_electron=io_initial_electron)
+        electron_pseudotime, n_debug_outputs =
+            @views update_electron_pdf!(fvec, pdf.electron.norm, moments,
+                                        moments.electron.dens, moments.electron.vth,
+                                        moments.electron.ppar, moments.electron.qpar,
+                                        moments.electron.qpar_updated, phi,
+                                        moments.electron.ddens_dz,
+                                        moments.electron.dppar_dz,
+                                        moments.electron.dqpar_dz,
+                                        moments.electron.dvth_dz, r, z, vperp, vpa,
+                                        z_spectral, vpa_spectral, z_advect, vpa_advect,
+                                        scratch_dummy, dt, collisions, composition,
+                                        num_diss_params, max_electron_pdf_iterations;
+                                        io_initial_electron=io_initial_electron)
 
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
-        n_debug_outputs = output_counter = result[end]
-        t = result[end-1]
         t_idx = n_debug_outputs+1
-        write_initial_electron_state(pdf.electron.norm, moments, t, io_initial_electron,
-                                     t_idx, r, z, vperp, vpa)
+        write_initial_electron_state(pdf.electron.norm, moments, electron_pseudotime,
+                                     io_initial_electron, t_idx, r, z, vperp, vpa)
         finish_initial_electron_io(io_initial_electron)
 
-        return result
     end
+    return nothing
 end
 
 """
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 04f1aa7c5..d39dc1360 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -118,9 +118,7 @@ function run_moment_kinetics(to::Union{TimerOutput,Nothing}, input_dict=Dict();
     mk_state = nothing
     try
         # set up all the structs, etc. needed for a run
-        #mk_state = setup_moment_kinetics(input_dict; restart=restart,
-        #                                 restart_time_index=restart_time_index)
-        return setup_moment_kinetics(input_dict; restart=restart,
+        mk_state = setup_moment_kinetics(input_dict; restart=restart,
                                          restart_time_index=restart_time_index)
 
         # solve the 1+1D kinetic equation to advance f in time by nstep time steps
@@ -344,7 +342,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
         # each of which may be evolved separately depending on input choices.
-        return init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry,
+        init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry,
                               composition, r, z, vperp, vpa, vzeta, vr, vz,
                               z_spectral, r_spectral, vpa_spectral, vz_spectral, species,
                               collisions, external_source_settings,

From 393a8c045be9cd71f862d441b349a5666cafd612 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Feb 2024 11:51:11 +0000
Subject: [PATCH 108/428] Fix merge of electrons and master

---
 moment_kinetics/src/initial_conditions.jl     |  22 ++-
 moment_kinetics/src/moment_kinetics.jl        |  26 +--
 moment_kinetics/src/time_advance.jl           | 185 +++---------------
 moment_kinetics/src/velocity_moments.jl       |   8 +-
 .../test/Krook_collisions_tests.jl            |   4 +-
 .../fokker_planck_time_evolution_tests.jl     |   4 +-
 .../test/nonlinear_sound_wave_tests.jl        |   6 +-
 .../test/restart_interpolation_tests.jl       |   2 +-
 8 files changed, 75 insertions(+), 182 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index dbced58a5..a124c2014 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -276,6 +276,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                      pdf.ion.norm, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
 
+        initialize_external_source_amplitude!(moments, external_source_settings, vperp,
+                                              vzeta, vr, n_neutral_species)
+        initialize_external_source_controller_integral!(moments, external_source_settings,
+                                                        n_neutral_species)
+
         if n_neutral_species > 0
             update_neutral_qz!(moments.neutral.qz, moments.neutral.qz_updated,
                                moments.neutral.dens, moments.neutral.uz,
@@ -1003,6 +1008,11 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
         vz, vr, vzeta, z, vz_spectral, density, uz, pz, vth, v_norm_fac, evolve_density,
         evolve_upar, evolve_ppar, wall_flux_0, wall_flux_L)
 
+    # Reduce the ion flux by `recycling_fraction` to account for ions absorbed by the
+    # wall.
+    wall_flux_0 *= composition.recycling_fraction
+    wall_flux_L *= composition.recycling_fraction
+
     #if spec.vz_IC.initialization_option == "gaussian"
     # For now, continue to use 'vpa' initialization options for neutral species
     if spec.vpa_IC.initialization_option == "gaussian"
@@ -1341,7 +1351,7 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     return nothing
 end
 
-function init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composition)
+function init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composition, zero)
 
     begin_serial_region()
     @serial_region begin
@@ -1370,7 +1380,7 @@ function init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composi
                             v_transverse = sqrt(vzeta.grid[ivzeta]^2 + vr.grid[ivr]^2)
                             v_normal = abs(vz.grid[ivz])
                             v_tot = sqrt(v_normal^2 + v_transverse^2)
-                            if  v_tot > 0.0
+                            if  v_tot > zero
                                 prefac = v_normal/v_tot
                             else
                                 prefac = 0.0
@@ -1432,11 +1442,13 @@ various boundaries. Also initialise the Knudsen cosine distribution here so it c
 when initialising the neutral pdf.
 """
 function create_boundary_distributions(vz, vr, vzeta, vpa, vperp, z, composition)
+    zero = 1.0e-14
+
     #initialise knudsen distribution for neutral wall bc
     knudsen_cosine = allocate_shared_float(vz.n, vr.n, vzeta.n)
     #initialise knudsen distribution for neutral wall bc - can be done here as this only
     #depends on T_wall, which has already been set
-    init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composition)
+    init_knudsen_cosine!(knudsen_cosine, vz, vr, vzeta, vpa, vperp, composition, zero)
     #initialise fixed-in-time radial boundary condition based on initial condition values
     pdf_rboundary_ion = allocate_shared_float(vpa.n, vperp.n, z.n, 2,
                                                   composition.n_ion_species)
@@ -1459,8 +1471,8 @@ enforce boundary conditions in vpa and z on the evolved pdf;
 also enforce boundary conditions in z on all separately evolved velocity space moments of the pdf
 """
 function enforce_boundary_conditions!(f, f_r_bc, density, upar, ppar, moments, vpa_bc,
-        z_bc, r_bc, vpa, vperp, z, r, vpa_adv, z_adv, r_adv, composition, scratch_dummy,
-        r_diffusion, vpa_diffusion)
+        z_bc, r_bc, vpa, vperp, z, r, vpa_spectral, vperp_spectral, vpa_adv, z_adv, r_adv,
+        composition, scratch_dummy, r_diffusion, vpa_diffusion)
     if vpa.n > 1
         begin_s_r_z_vperp_region()
         @loop_s_r_z_vperp is ir iz ivperp begin
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index d39dc1360..0a7f1fd58 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -85,7 +85,7 @@ using .debugging
 using .external_sources
 using .input_structs
 using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!,
-                           enforce_boundary_conditions!
+                           enforce_boundary_conditions!, initialize_scratch_arrays!
 using .load_data: reload_evolving_fields!
 using .looping
 using .moment_constraints: hard_force_moment_constraints!
@@ -318,8 +318,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     end
 
     # create the "fields" structure that contains arrays
-    # for the electrostatic potential phi (and eventually the electromagnetic fields)
-    fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude, 
+    # for the electrostatic potential phi and the electromagnetic fields
+    fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude,
                              drive_input.frequency, drive_input.force_Er_zero_at_wall)
 
     # Allocate arrays and create the pdf and moments structs
@@ -416,17 +416,21 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         initialize_external_source_amplitude!(moments, external_source_settings, vperp,
                                               vzeta, vr, composition.n_neutral_species)
 
+        # Copy the reloaded values into the `scratch` struct
+        initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
+
         _block_synchronize()
     end
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
-    moments, spectral_objects, advance, manufactured_source_list =
-        setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
-            vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral,
-            r_spectral, composition, drive_input, moments, fields, t_input, collisions,
-            geometry, boundary_distributions, external_source_settings, num_diss_params,
-            manufactured_solns_input, restarting)
+    moments, fields, spectral_objects, advance, fp_arrays, manufactured_source_list =
+        setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z, r,
+            vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
+            z_spectral, r_spectral, composition, drive_input, moments, t_input,
+            collisions, species, geometry, boundary_distributions,
+            external_source_settings, num_diss_params, manufactured_solns_input,
+            advection_structs, scratch_dummy, restarting)
 
     # This is the closest we can get to the end time of the setup before writing it to the
     # output file
@@ -438,8 +442,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         moments.evolve_upar, moments.evolve_ppar, external_source_settings, input_dict,
         restart_time_index, previous_runs_info, time_for_setup)
     # write initial data to ascii files
-    write_data_to_ascii(moments, fields, vpa, vperp, z, r, code_time,
-        composition.n_ion_species, composition.n_neutral_species, ascii_io)
+    write_data_to_ascii(moments, fields, z, r, code_time, composition.n_ion_species,
+        composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
     write_all_moments_data_to_binary(moments, fields, code_time,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 0401ea349..6a41de159 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -266,12 +266,13 @@ this includes creating and populating structs
 for Chebyshev transforms, velocity space moments,
 EM fields, and advection terms
 """
-function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
-                             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
-                             z_spectral, r_spectral, composition, drive_input, moments,
-                             t_input, collisions, species, geometry,
+function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z, r,
+                             vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral,
+                             vperp_spectral, z_spectral, r_spectral, composition,
+                             drive_input, moments, t_input, collisions, species, geometry,
                              boundary_distributions, external_source_settings,
-                             num_diss_params, manufactured_solns_input, restarting)
+                             num_diss_params, manufactured_solns_input, advection_structs,
+                             scratch_dummy, restarting)
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
@@ -288,29 +289,20 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
 
     begin_serial_region()
 
-    # create an array of structs containing scratch arrays for the pdf and low-order moments
-    # that may be evolved separately via fluid equations
-    scratch = setup_scratch_arrays(moments, pdf.ion.norm, pdf.neutral.norm, t_input.n_rk_stages)
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
-    scratch_dummy = setup_dummy_and_buffer_arrays(r.n,z.n,vpa.n,vperp.n,vz.n,vr.n,vzeta.n,
-                                   composition.n_ion_species,n_neutral_species_alloc)
     # create arrays for Fokker-Planck collisions 
     if advance.explicit_weakform_fp_collisions
         fp_arrays = init_fokker_planck_collisions_weak_form(vpa,vperp,vpa_spectral,vperp_spectral; precompute_weights=true)
     else
         fp_arrays = nothing
     end
-    # create the "fields" structure that contains arrays
-    # for the electrostatic potential phi and eventually the electromagnetic fields
-    fields = setup_em_fields(z.n, r.n, drive_input.force_phi, drive_input.amplitude, drive_input.frequency, drive_input.force_Er_zero_at_wall)
-    # initialize the electrostatic potential
-    begin_serial_region()
     # update the derivatives of the electron moments as these may be needed when
     # computing the electrostatic potential (and components of the electric field)
     calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 
                                            z_spectral, num_diss_params, composition.electron_physics)
     # initialize the electrostatic potential
+    begin_serial_region()
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
     @serial_region begin
         # save the initial phi(z) for possible use later (e.g., if forcing phi)
@@ -325,15 +317,17 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
     calculate_moment_derivatives_neutral!(moments, scratch[1], scratch_dummy, z,
                                           z_spectral, num_diss_params)
 
+    r_advect = advection_structs.r_advect
+    z_advect = advection_structs.z_advect
+    vperp_advect = advection_structs.vperp_advect
+    vpa_advect = advection_structs.vpa_advect
+    neutral_r_advect = advection_structs.neutral_r_advect
+    neutral_z_advect = advection_structs.neutral_z_advect
+    neutral_vz_advect = advection_structs.neutral_vz_advect
     ##
     # ion particle advection only
     ##
 
-    # create structure r_advect whose members are the arrays needed to compute
-    # the advection term(s) appearing in the split part of the GK equation dealing
-    # with advection in r
-    begin_serial_region()
-    r_advect = setup_advection(n_ion_species, r, vpa, vperp, z)
     if r.n > 1
         # initialise the r advection speed
         begin_s_z_vperp_vpa_region()
@@ -345,11 +339,6 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         # enforce prescribed boundary condition in r on the distribution function f
     end
 
-    # create structure z_advect whose members are the arrays needed to compute
-    # the advection term(s) appearing in the split part of the GK equation dealing
-    # with advection in z
-    begin_serial_region()
-    z_advect = setup_advection(n_ion_species, z, vpa, vperp, r)
     if z.n > 1
         # initialise the z advection speed
         begin_s_r_vperp_vpa_region()
@@ -360,23 +349,13 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
                                    geometry)
         end
     end
-    begin_serial_region()
 
-    # create structure vpa_advect whose members are the arrays needed to compute
-    # the advection term(s) appearing in the split part of the GK equation dealing
-    # with advection in vpa
-    vpa_advect = setup_advection(n_ion_species, vpa, vperp, z, r)
     # initialise the vpa advection speed
     begin_s_r_z_vperp_region()
     update_speed_vpa!(vpa_advect, fields, scratch[1], moments, vpa, vperp, z, r,
                       composition, collisions, external_source_settings.ion, 0.0,
                       geometry)
 
-    # create structure vperp_advect whose members are the arrays needed to compute
-    # the advection term(s) appearing in the split part of the GK equation dealing
-    # with advection in vperp
-    begin_serial_region()
-    vperp_advect = setup_advection(n_ion_species, vperp, vpa, z, r)
     # initialise the vperp advection speed
     if vperp.n > 1
         begin_serial_region()
@@ -391,9 +370,6 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
     # Neutral particle advection
     ##
 
-    # create structure neutral_r_advect for neutral particle advection
-    begin_serial_region()
-    neutral_r_advect = setup_advection(n_neutral_species_alloc, r, vz, vr, vzeta, z)
     if n_neutral_species > 0 && r.n > 1
         # initialise the r advection speed
         begin_sn_z_vzeta_vr_vz_region()
@@ -402,9 +378,6 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         end
     end
 
-    # create structure neutral_z_advect for neutral particle advection
-    begin_serial_region()
-    neutral_z_advect = setup_advection(n_neutral_species_alloc, z, vz, vr, vzeta, r)
     if n_neutral_species > 0 && z.n > 1
         # initialise the z advection speed
         begin_sn_r_vzeta_vr_vz_region()
@@ -417,9 +390,6 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         end
     end
 
-    # create structure neutral_vz_advect for neutral particle advection
-    begin_serial_region()
-    neutral_vz_advect = setup_advection(n_neutral_species_alloc, vz, vr, vzeta, z, r)
     if n_neutral_species > 0
         # initialise the z advection speed
         @serial_region begin
@@ -440,7 +410,6 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
     # construct advect & spectral objects to compactify arguments
     ##
 
-    advect_objects = advect_object_struct(vpa_advect, vperp_advect, z_advect, r_advect, neutral_z_advect, neutral_r_advect, neutral_vz_advect)
     spectral_objects = spectral_object_struct(vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral, r_spectral)
     if(advance.manufactured_solns_test)
         manufactured_source_list = manufactured_sources(manufactured_solns_input, r, z,
@@ -513,16 +482,16 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
         calculate_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
         calculate_electron_upar_from_charge_conservation!(moments.electron.upar, moments.electron.upar_updated,
                                                           moments.electron.dens, moments.ion.upar, moments.ion.dens,
-                                                          composition.electron_physics)
+                                                          composition.electron_physics, r, z)
         # compute the updated electron temperature
         # NB: not currently necessary, as initial vth is not directly dependent on ion quantities
         @. moments.electron.temp = moments.electron.vth^2
         # as the electron temperature has now been updated, set the appropriate flag
-        moments.electron.temp_updated = true
+        moments.electron.temp_updated[] = true
         # compute the updated electron parallel pressure
         @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
         # as the electron ppar has now been updated, set the appropriate flag
-        moments.electron.ppar_updated = true
+        moments.electron.ppar_updated[] = true
         # calculate the zed derivative of the initial electron temperature, potentially
         # needed in the following calculation of the electron parallel friction force and
         # parallel heat flux
@@ -568,8 +537,8 @@ function setup_time_advance!(pdf, scratch, vz, vr, vzeta, vpa, vperp, z, r, vz_s
     # Ensure all processes are synchronized at the end of the setup
     _block_synchronize()
 
-    return moments, fields, spectral_objects, advect_objects,
-           advance, fp_arrays, scratch_dummy, manufactured_source_list
+    return moments, fields, spectral_objects, advance, fp_arrays, scratch_dummy,
+           manufactured_source_list
 end
 
 """
@@ -876,35 +845,6 @@ function normalize_pdf!(pdf, moments, scratch)
     return nothing
 end
 
-"""
-initialize the array of structs containing scratch arrays for the normalised pdf and low-order moments
-that may be evolved separately via fluid equations
-"""
-function initialize_scratch_arrays!(scratch, moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
-    # populate each of the structs
-    for istage ∈ 1:n_rk_stages+1
-        @serial_region begin
-            # initialise the scratch arrays for the ion pdf and velocity moments
-            scratch[istage].pdf .= pdf_ion_in
-            scratch[istage].density .= moments.ion.dens
-            scratch[istage].upar .= moments.ion.upar
-            scratch[istage].ppar .= moments.ion.ppar
-            # initialise the scratch arrays for the electron pdf and velocity moments
-            scratch[istage].pdf_electron .= pdf_electron_in
-            scratch[istage].electron_density .= moments.electron.dens
-            scratch[istage].electron_upar .= moments.electron.upar
-            scratch[istage].electron_ppar .= moments.electron.ppar
-            scratch[istage].electron_temp .= moments.electron.temp
-            # initialise the scratch arrays for the neutral velocity moments and pdf
-            scratch[istage].pdf_neutral .= pdf_neutral_in
-            scratch[istage].density_neutral .= moments.neutral.dens
-            scratch[istage].uz_neutral .= moments.neutral.uz
-            scratch[istage].pz_neutral .= moments.neutral.pz
-        end
-    end
-    return nothing
-end
-
 """
 given the number of Runge Kutta stages that are requested,
 returns the needed Runge Kutta coefficients;
@@ -1055,9 +995,8 @@ function time_advance!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyro
                           Dates.format(now(), dateformat"H:MM:SS"))
                 end
             end
-            write_data_to_ascii(moments, fields, vpa, vperp, z, r, t,
-                                composition.n_ion_species, composition.n_neutral_species,
-                                ascii_io)
+            write_data_to_ascii(moments, fields, z, r, t, composition.n_ion_species,
+                                composition.n_neutral_species, ascii_io)
             write_all_moments_data_to_binary(moments, fields, t,
                                              composition.n_ion_species,
                                              composition.n_neutral_species, io_moments,
@@ -1365,7 +1304,7 @@ or update them by taking the appropriate velocity moment of the evolved pdf
 """
 function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, vr, vzeta,
                     vpa, vperp, z, r, spectral_objects, advect_objects, rk_coefs, istage, composition,
-                    geometry, num_diss_params, advance, scratch_dummy)
+                    collisions, geometry, num_diss_params, advance, scratch_dummy)
     begin_s_r_z_region()
 
     new_scratch = scratch[istage+1]
@@ -1440,73 +1379,10 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                   num_diss_params)
 
-    ##
-    # update the neutral particle distribution and moments
-    ##
-
-    if composition.n_neutral_species > 0
-        begin_sn_r_z_region()
-        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
-            new_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn] = ( rk_coefs[1]*pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn]
-             + rk_coefs[2]*old_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn] + rk_coefs[3]*new_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn])
-        end
-        # use Runge Kutta to update any velocity moments evolved separately from the pdf
-        rk_update_evolved_moments_neutral!(new_scratch, old_scratch, moments, rk_coefs)
-
-        # Ensure there are no negative values in the pdf before applying boundary
-        # conditions, so that negative deviations do not mess up the integral-constraint
-        # corrections in the sheath boundary conditions.
-        force_minimum_pdf_value_neutral!(new_scratch.pdf_neutral, num_diss_params)
-
-        # Enforce boundary conditions in z and vpa on the distribution function.
-        # Must be done after Runge Kutta update so that the boundary condition applied to
-        # the updated pdf is consistent with the updated moments - otherwise different upar
-        # between 'pdf', 'old_scratch' and 'new_scratch' might mean a point that should be
-        # set to zero at the sheath boundary according to the final upar has a non-zero
-        # contribution from one or more of the terms.
-        # NB: probably need to do the same for the evolved moments
-        # Note, so far vr and vzeta do not need advect objects, so pass `nothing` for
-        # those as a placeholder
-        enforce_neutral_boundary_conditions!(new_scratch.pdf_neutral, new_scratch.pdf,
-            boundary_distributions, new_scratch.density_neutral, new_scratch.uz_neutral,
-            new_scratch.pz_neutral, moments, new_scratch.density, new_scratch.upar,
-            fields.Er, vzeta_spectral, vr_spectral, vz_spectral, neutral_r_advect,
-            neutral_z_advect, nothing, nothing, neutral_vz_advect, r, z, vzeta, vr, vz,
-            composition, geometry, scratch_dummy, advance.r_diffusion,
-            advance.vz_diffusion)
-
-        if moments.evolve_density && moments.enforce_conservation
-            begin_sn_r_z_region()
-            @loop_sn_r_z isn ir iz begin
-                @views hard_force_moment_constraints_neutral!(
-                    new_scratch.pdf_neutral[:,:,:,iz,ir,isn], moments, vz)
-            end
-        end
-
-        # update remaining velocity moments that are calculable from the evolved pdf
-        update_derived_moments_neutral!(new_scratch, moments, vz, vr, vzeta, z, r,
-                                        composition)
-        # update the thermal speed
-        begin_sn_r_z_region()
-        @loop_sn_r_z isn ir iz begin
-            moments.neutral.vth[iz,ir,isn] = sqrt(2.0*new_scratch.pz_neutral[iz,ir,isn]/new_scratch.density_neutral[iz,ir,isn])
-        end
-
-        # update the parallel heat flux
-        update_neutral_qz!(moments.neutral.qz, moments.neutral.qz_updated,
-                           new_scratch.density_neutral, new_scratch.uz_neutral,
-                           moments.neutral.vth, new_scratch.pdf_neutral, vz, vr, vzeta, z,
-                           r, composition, moments.evolve_density, moments.evolve_upar,
-                           moments.evolve_ppar)
-
-        calculate_moment_derivatives_neutral!(moments, new_scratch, scratch_dummy, z,
-                                              z_spectral, num_diss_params)
-    end
-
     # update the lowest three electron moments (density, upar and ppar)
     calculate_electron_density!(new_scratch.electron_density, moments.electron.dens_updated, new_scratch.density)
     calculate_electron_upar_from_charge_conservation!(new_scratch.electron_upar, moments.electron.upar_updated,
-        new_scratch.electron_density, new_scratch.upar, new_scratch.density, composition.electron_physics)
+        new_scratch.electron_density, new_scratch.upar, new_scratch.density, composition.electron_physics, r, z)
     # if electron model is braginskii_fluid, then ppar is evolved via the energy equation
     # and is already updated;
     # otherwise update assuming electron temperature is fixed in time
@@ -1519,10 +1395,10 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
     end
     @. moments.electron.temp = 2 * new_scratch.electron_ppar / new_scratch.electron_density
-    moments.electron.temp_updated = true
+    moments.electron.temp_updated[] = true
     @. moments.electron.vth = sqrt(moments.electron.temp)
     # regardless of electron model, electron ppar is now updated
-    moments.electron.ppar_updated = true
+    moments.electron.ppar_updated[] = true
     # calculate the corresponding zed derivatives of the moments
     calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                            num_diss_params, composition.electron_physics)
@@ -1565,9 +1441,10 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         enforce_neutral_boundary_conditions!(new_scratch.pdf_neutral, new_scratch.pdf,
             boundary_distributions, new_scratch.density_neutral, new_scratch.uz_neutral,
             new_scratch.pz_neutral, moments, new_scratch.density, new_scratch.upar,
-            fields.Er, neutral_r_advect, neutral_z_advect, nothing, nothing,
-            neutral_vz_advect, r, z, vzeta, vr, vz, composition, geometry, scratch_dummy,
-            advance.r_diffusion, advance.vz_diffusion)
+            fields.Er, vzeta_spectral, vr_spectral, vz_spectral, neutral_r_advect,
+            neutral_z_advect, nothing, nothing, neutral_vz_advect, r, z, vzeta, vr, vz,
+            composition, geometry, scratch_dummy, advance.r_diffusion,
+            advance.vz_diffusion)
 
         if moments.evolve_density && moments.enforce_conservation
             begin_sn_r_z_region()
@@ -1768,8 +1645,8 @@ function ssp_rk!(pdf, scratch, t, t_input, vz, vr, vzeta, vpa, vperp, gyrophase,
             external_source_settings, num_diss_params, advance, fp_arrays, istage)
         @views rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, vr,
                           vzeta, vpa, vperp, z, r, spectral_objects, advect_objects,
-                          advance.rk_coefs[:,istage], istage, composition, geometry,
-                          num_diss_params, advance, scratch_dummy)
+                          advance.rk_coefs[:,istage], istage, composition, collisions,
+                          geometry, num_diss_params, advance, scratch_dummy)
     end
 
     istage = n_rk_stages+1
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 968966a31..4ba5d081e 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -1973,10 +1973,10 @@ function reset_moments_status!(moments)
     moments.neutral.pzeta_updated .= false
     moments.neutral.pr_updated .= false
     moments.neutral.qz_updated .= false
-    moments.electron.dens_updated = false
-    moments.electron.upar_updated = false
-    moments.electron.ppar_updated = false
-    moments.electron.qpar_updated = false
+    moments.electron.dens_updated[] = false
+    moments.electron.upar_updated[] = false
+    moments.electron.ppar_updated[] = false
+    moments.electron.qpar_updated[] = false
 end
 
 end
diff --git a/moment_kinetics/test/Krook_collisions_tests.jl b/moment_kinetics/test/Krook_collisions_tests.jl
index ec247545a..ef2528814 100644
--- a/moment_kinetics/test/Krook_collisions_tests.jl
+++ b/moment_kinetics/test/Krook_collisions_tests.jl
@@ -10,7 +10,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_ion_particle_moments_data, load_pdf_data,
+                                 load_ion_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -224,7 +224,7 @@ function run_test(test_input, rtol, atol; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_particle_moments_data(fid)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
             z, z_spectral = load_coordinate_data(fid, "z")
 
diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index 3d066c67f..ac3507e3a 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -8,7 +8,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_ion_particle_moments_data, load_pdf_data,
+                                 load_ion_moments_data, load_pdf_data,
                                  load_time_data, load_species_data
 using moment_kinetics.type_definitions: mk_float
 
@@ -249,7 +249,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
 
             # load velocity moments data
             n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, 
-            pperp_ion_zrst, qpar_ion_zrst, v_t_ion_zrst, dSdt_zrst = load_ion_particle_moments_data(fid,extended_moments=true)
+            pperp_ion_zrst, qpar_ion_zrst, v_t_ion_zrst, dSdt_zrst = load_ion_moments_data(fid,extended_moments=true)
             
             close(fid)
             
diff --git a/moment_kinetics/test/nonlinear_sound_wave_tests.jl b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
index 71f32dd1d..f6e5d7334 100644
--- a/moment_kinetics/test/nonlinear_sound_wave_tests.jl
+++ b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
@@ -8,7 +8,7 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_ion_particle_moments_data, load_pdf_data,
+                                 load_ion_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa
@@ -89,7 +89,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
 
             # load velocity moments data
-            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_particle_moments_data(fid)
+            n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
             z, z_spectral = load_coordinate_data(fid, "z")
 
@@ -214,7 +214,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
                 @test isapprox(expected.ppar_ion[:, tind], newgrid_ppar_ion[:,1], rtol=rtol)
 
                 newgrid_vth_ion = @. sqrt(2.0*newgrid_ppar_ion/newgrid_n_ion)
-                newgrid_f_ion = interpolate_to_grid_vpa(expected.vpa, newgrid_f_ion, vpa, vpa_spectral)
+                newgrid_f_ion = interpolate_to_grid_z(expected.z, f_ion[:, :, :, tind], z, z_spectral)
                 temp = newgrid_f_ion
                 newgrid_f_ion = fill(NaN, length(expected.vpa),
                                      size(newgrid_f_ion, 2),
diff --git a/moment_kinetics/test/restart_interpolation_tests.jl b/moment_kinetics/test/restart_interpolation_tests.jl
index cbdacded2..6308890ff 100644
--- a/moment_kinetics/test/restart_interpolation_tests.jl
+++ b/moment_kinetics/test/restart_interpolation_tests.jl
@@ -12,7 +12,7 @@ using moment_kinetics.file_io: io_has_parallel
 using moment_kinetics.input_structs: grid_input, advection_input, hdf5
 using moment_kinetics.load_data: open_readonly_output_file, load_coordinate_data,
                                  load_species_data, load_fields_data,
-                                 load_ion_particle_moments_data, load_pdf_data,
+                                 load_ion_moments_data, load_pdf_data,
                                  load_neutral_particle_moments_data,
                                  load_neutral_pdf_data, load_time_data, load_species_data
 using moment_kinetics.interpolation: interpolate_to_grid_z, interpolate_to_grid_vpa

From f6287ee16a0d4139271311f560c6986560f6f6b2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 20 Feb 2024 15:13:37 +0000
Subject: [PATCH 109/428] Move pdf and moments structs to
 `moment_kinetics_structs`

Allows them to be used in other modules that are defined before
`initial_conditions` and/or `velocity_moments`.
---
 moment_kinetics/src/initial_conditions.jl     |  48 +---
 .../src/moment_kinetics_structs.jl            | 269 ++++++++++++++++++
 moment_kinetics/src/velocity_moments.jl       | 222 +--------------
 3 files changed, 273 insertions(+), 266 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index a124c2014..8e94ebe91 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -26,13 +26,13 @@ using ..looping
 using ..em_fields: update_phi!
 using ..file_io: setup_initial_electron_io, write_initial_electron_state,
                  finish_initial_electron_io
-using ..moment_kinetics_structs: scratch_pdf
+using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, pdf_struct, moments_struct,
+                                 boundary_distributions_struct
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
 using ..velocity_moments: create_moments_ion, create_moments_electron, create_moments_neutral
 using ..velocity_moments: update_qpar!
-using ..velocity_moments: moments_ion_substruct, moments_electron_substruct, moments_neutral_substruct
 using ..velocity_moments: update_neutral_density!, update_neutral_pz!, update_neutral_pr!, update_neutral_pzeta!
 using ..velocity_moments: update_neutral_uz!, update_neutral_ur!, update_neutral_uzeta!, update_neutral_qz!
 using ..velocity_moments: update_ppar!, update_upar!, update_density!, update_pperp!, update_vth!, reset_moments_status!
@@ -48,50 +48,6 @@ using ..manufactured_solns: manufactured_solutions
 
 using MPI
 
-"""
-"""
-struct pdf_substruct{n_distribution}
-    norm::MPISharedArray{mk_float,n_distribution}
-    buffer::MPISharedArray{mk_float,n_distribution} # for collision operator terms when pdfs must be interpolated onto different velocity space grids
-end
-
-# struct of structs neatly contains i+n info?
-struct pdf_struct
-    #ion particles: s + r + z + vperp + vpa
-    ion::pdf_substruct{5}
-    # electron particles: r + z + vperp + vpa
-    electron::pdf_substruct{4}
-    #neutral particles: s + r + z + vzeta + vr + vz
-    neutral::pdf_substruct{6}
-end
-
-struct moments_struct
-    ion::moments_ion_substruct
-    electron::moments_electron_substruct
-    neutral::moments_neutral_substruct
-    # flag that indicates if the density should be evolved via continuity equation
-    evolve_density::Bool
-    # flag that indicates if particle number should be conserved for each species
-    # effects like ionisation or net particle flux from the domain would lead to
-    # non-conservation
-    particle_number_conserved::Bool
-    # flag that indicates if exact particle conservation should be enforced
-    enforce_conservation::Bool
-    # flag that indicates if the parallel flow should be evolved via force balance
-    evolve_upar::Bool
-    # flag that indicates if the parallel pressure should be evolved via the energy equation
-    evolve_ppar::Bool
-end
-
-struct boundary_distributions_struct
-    # knudsen cosine distribution for imposing the neutral wall boundary condition
-    knudsen::MPISharedArray{mk_float,3}
-    # ion particle r boundary values (vpa,vperp,z,r,s)
-    pdf_rboundary_ion::MPISharedArray{mk_float,5}
-    # neutral particle r boundary values (vz,vr,vzeta,z,r,s)
-    pdf_rboundary_neutral::MPISharedArray{mk_float,6}
-end
-
 """
 Creates the structs for the pdf and the velocity-space moments
 """
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index 9e4901e71..38dd77469 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -53,6 +53,275 @@ struct em_fields_struct
     force_Er_zero_at_wall::Bool
 end
 
+"""
+"""
+struct moments_ion_substruct
+    # this is the particle density
+    dens::MPISharedArray{mk_float,3}
+    # flag that keeps track of if the density needs updating before use
+    # Note: may not be set for all species on this process, but this process only ever
+    # sets/uses the value for the same subset of species. This means dens_update does
+    # not need to be a shared memory array.
+    dens_updated::Vector{Bool}
+    # this is the parallel flow
+    upar::MPISharedArray{mk_float,3}
+    # flag that keeps track of whether or not upar needs updating before use
+    # Note: may not be set for all species on this process, but this process only ever
+    # sets/uses the value for the same subset of species. This means upar_update does
+    # not need to be a shared memory array.
+    upar_updated::Vector{Bool}
+    # this is the parallel pressure
+    ppar::MPISharedArray{mk_float,3}
+    # flag that keeps track of whether or not ppar needs updating before use
+    # Note: may not be set for all species on this process, but this process only ever
+    # sets/uses the value for the same subset of species. This means ppar_update does
+    # not need to be a shared memory array.
+    ppar_updated::Vector{Bool}
+    # this is the perpendicular pressure
+    pperp::MPISharedArray{mk_float,3}
+    # this is the parallel heat flux
+    qpar::MPISharedArray{mk_float,3}
+    # flag that keeps track of whether or not qpar needs updating before use
+    # Note: may not be set for all species on this process, but this process only ever
+    # sets/uses the value for the same subset of species. This means qpar_update does
+    # not need to be a shared memory array.
+    qpar_updated::Vector{Bool}
+    # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
+    vth::MPISharedArray{mk_float,3}
+    # generalised Chodura integrals for the lower and upper plates
+    chodura_integral_lower::MPISharedArray{mk_float,2}
+    chodura_integral_upper::MPISharedArray{mk_float,2}
+    # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
+    # a factor of vth for each power of wpa in velocity space integrals.
+    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
+    # and it is one otherwise
+    v_norm_fac::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the upwinded z-derivative of the particle density
+    ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the particle density
+    d2dens_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the parallel flow
+    dupar_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the upwinded z-derivative of the parallel flow
+    dupar_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the parallel flow
+    d2upar_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the parallel pressure
+    dppar_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the upwinded z-derivative of the parallel pressure
+    dppar_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the parallel pressure
+    d2ppar_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the parallel heat flux
+    dqpar_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
+    dvth_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the entropy production dS/dt = - int (ln f sum_s' C_ss' [f_s,f_s']) d^3 v
+    dSdt::MPISharedArray{mk_float,3}
+    # Spatially varying amplitude of the external source term
+    external_source_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the density moment of the external source term
+    external_source_density_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel momentum moment of the external source
+    # term
+    external_source_momentum_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel pressure moment of the external source
+    # term
+    external_source_pressure_amplitude::MPISharedArray{mk_float,2}
+    # Integral term for the PID controller of the external source term
+    external_source_controller_integral::MPISharedArray{mk_float,2}
+end
+
+"""
+moments_electron_substruct is a struct that contains moment information for electrons
+"""
+struct moments_electron_substruct
+    # this is the particle density
+    dens::MPISharedArray{mk_float,2}
+    # flag that keeps track of if the density needs updating before use
+    dens_updated::Ref{Bool}
+    # this is the parallel flow
+    upar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not upar needs updating before use
+    upar_updated::Ref{Bool}
+    # this is the parallel pressure
+    ppar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not ppar needs updating before use
+    ppar_updated::Ref{Bool}
+    # this is the temperature
+    temp::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not temp needs updating before use
+    temp_updated::Ref{Bool}
+    # this is the parallel heat flux
+    qpar::MPISharedArray{mk_float,2}
+    # flag that keeps track of whether or not qpar needs updating before use
+    qpar_updated::Ref{Bool}
+    # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
+    vth::MPISharedArray{mk_float,2}
+    # this is the parallel friction force between ions and electrons
+    parallel_friction::MPISharedArray{mk_float,2}
+    # this is the electron heat source
+    heat_source::MPISharedArray{mk_float,2}
+    # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
+    # a factor of vth for each power of wpa in velocity space integrals.
+    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
+    # and it is one otherwise
+    v_norm_fac::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel flow
+    dupar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel pressure
+    dppar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the parallel pressure
+    dppar_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the second-z-derivative of the parallel pressure
+    d2ppar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel heat flux
+    dqpar_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the parallel temperature Tpar = ppar/dens
+    dT_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the upwinded z-derivative of the temperature Tpar = ppar/dens
+    dT_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
+    # this is the z-derivative of the electron thermal speed vth = sqrt(2*Tpar/m)
+    dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
+end
+
+"""
+"""
+struct moments_neutral_substruct
+    # this is the particle density
+    dens::MPISharedArray{mk_float,3}
+    # flag that keeps track of if the density needs updating before use
+    # Note: may not be set for all species on this process, but this process only ever
+    # sets/uses the value for the same subset of species. This means dens_update does
+    # not need to be a shared memory array.
+    dens_updated::Vector{Bool}
+    # this is the particle mean velocity in z
+    uz::MPISharedArray{mk_float,3}
+    # flag that keeps track of if uz needs updating before use
+    uz_updated::Vector{Bool}
+    # this is the particle mean velocity in r
+    ur::MPISharedArray{mk_float,3}
+    # flag that keeps track of if ur needs updating before use
+    ur_updated::Vector{Bool}
+    # this is the particle mean velocity in zeta
+    uzeta::MPISharedArray{mk_float,3}
+    # flag that keeps track of if uzeta needs updating before use
+    uzeta_updated::Vector{Bool}
+    # this is the zz particle pressure tensor component
+    pz::MPISharedArray{mk_float,3}
+    # flag that keeps track of if pz needs updating before use
+    pz_updated::Vector{Bool}
+    # this is the rr particle pressure tensor component
+    pr::MPISharedArray{mk_float,3}
+    # flag that keeps track of if pr needs updating before use
+    pr_updated::Vector{Bool}
+    # this is the zetazeta particle pressure tensor component
+    pzeta::MPISharedArray{mk_float,3}
+    # flag that keeps track of if pzeta needs updating before use
+    pzeta_updated::Vector{Bool}
+    # this is the total (isotropic) particle pressure
+    ptot::MPISharedArray{mk_float,3}
+    # this is the heat flux along z
+    qz::MPISharedArray{mk_float,3}
+    # flag that keeps track of if qz needs updating before use
+    qz_updated::Vector{Bool}
+    # this is the thermal speed based on the temperature T = ptot/dens: vth = sqrt(2*T/m)
+    vth::MPISharedArray{mk_float,3}
+    # if evolve_ppar = true, then the velocity variable is (vz - uz)/vth, which introduces
+    # a factor of vth for each power of wz in velocity space integrals.
+    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
+    # and it is one otherwise
+    v_norm_fac::MPISharedArray{mk_float,3}
+    # this is the z-derivative of the particle density
+    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the particle density
+    ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the particle density
+    d2dens_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the particle mean velocity in z
+    duz_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the upwinded z-derivative of the particle mean velocity in z
+    duz_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the particle mean velocity in z
+    d2uz_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the zz particle pressure tensor component
+    dpz_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the upwinded z-derivative of the zz particle pressure tensor component
+    dpz_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the second-z-derivative of the zz particle pressure tensor component
+    d2pz_dz2::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the thermal speed based on the temperature T = ptot/dens: vth = sqrt(2*T/m)
+    dvth_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # this is the z-derivative of the heat flux along z
+    dqz_dz::Union{MPISharedArray{mk_float,3},Nothing}
+    # Spatially varying amplitude of the external source term
+    external_source_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the density moment of the external source term
+    external_source_density_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel momentum moment of the external source
+    # term
+    external_source_momentum_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel pressure moment of the external source
+    # term
+    external_source_pressure_amplitude::MPISharedArray{mk_float,2}
+    # Integral term for the PID controller of the external source term
+    external_source_controller_integral::MPISharedArray{mk_float,2}
+end
+
+"""
+"""
+struct pdf_substruct{n_distribution}
+    norm::MPISharedArray{mk_float,n_distribution}
+    buffer::MPISharedArray{mk_float,n_distribution} # for collision operator terms when pdfs must be interpolated onto different velocity space grids
+end
+
+# struct of structs neatly contains i+n info?
+"""
+"""
+struct pdf_struct
+    #ion particles: s + r + z + vperp + vpa
+    ion::pdf_substruct{5}
+    # electron particles: r + z + vperp + vpa
+    electron::pdf_substruct{4}
+    #neutral particles: s + r + z + vzeta + vr + vz
+    neutral::pdf_substruct{6}
+end
+
+"""
+"""
+struct moments_struct
+    ion::moments_ion_substruct
+    electron::moments_electron_substruct
+    neutral::moments_neutral_substruct
+    # flag that indicates if the density should be evolved via continuity equation
+    evolve_density::Bool
+    # flag that indicates if particle number should be conserved for each species
+    # effects like ionisation or net particle flux from the domain would lead to
+    # non-conservation
+    particle_number_conserved::Bool
+    # flag that indicates if exact particle conservation should be enforced
+    enforce_conservation::Bool
+    # flag that indicates if the parallel flow should be evolved via force balance
+    evolve_upar::Bool
+    # flag that indicates if the parallel pressure should be evolved via the energy equation
+    evolve_ppar::Bool
+end
+
+"""
+"""
+struct boundary_distributions_struct
+    # knudsen cosine distribution for imposing the neutral wall boundary condition
+    knudsen::MPISharedArray{mk_float,3}
+    # ion particle r boundary values (vpa,vperp,z,r,s)
+    pdf_rboundary_ion::MPISharedArray{mk_float,5}
+    # neutral particle r boundary values (vz,vr,vzeta,z,r,s)
+    pdf_rboundary_neutral::MPISharedArray{mk_float,6}
+end
+
 """
 discretization_info for one dimension
 
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 4ba5d081e..d223d3319 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -14,7 +14,6 @@ export update_pperp!
 export update_qpar!
 export update_vth!
 export reset_moments_status!
-export moments_ion_substruct, moments_electron_substruct, moments_neutral_substruct
 export update_neutral_density!
 export update_neutral_uz!
 export update_neutral_ur!
@@ -42,231 +41,14 @@ using ..derivatives: derivative_z!
 using ..derivatives: derivative_r!
 using ..looping
 using ..input_structs: braginskii_fluid, kinetic_electrons
+using ..moment_kinetics_structs: moments_ion_substruct, moments_electron_substruct,
+                                 moments_neutral_substruct
 
 #global tmpsum1 = 0.0
 #global tmpsum2 = 0.0
 #global dens_hist = zeros(17,1)
 #global n_hist = 0
 
-"""
-"""
-struct moments_ion_substruct
-    # this is the particle density
-    dens::MPISharedArray{mk_float,3}
-    # flag that keeps track of if the density needs updating before use
-    # Note: may not be set for all species on this process, but this process only ever
-    # sets/uses the value for the same subset of species. This means dens_update does
-    # not need to be a shared memory array.
-    dens_updated::Vector{Bool}
-    # this is the parallel flow
-    upar::MPISharedArray{mk_float,3}
-    # flag that keeps track of whether or not upar needs updating before use
-    # Note: may not be set for all species on this process, but this process only ever
-    # sets/uses the value for the same subset of species. This means upar_update does
-    # not need to be a shared memory array.
-    upar_updated::Vector{Bool}
-    # this is the parallel pressure
-    ppar::MPISharedArray{mk_float,3}
-    # flag that keeps track of whether or not ppar needs updating before use
-    # Note: may not be set for all species on this process, but this process only ever
-    # sets/uses the value for the same subset of species. This means ppar_update does
-    # not need to be a shared memory array.
-    ppar_updated::Vector{Bool}
-    # this is the perpendicular pressure
-    pperp::MPISharedArray{mk_float,3}
-    # this is the parallel heat flux
-    qpar::MPISharedArray{mk_float,3}
-    # flag that keeps track of whether or not qpar needs updating before use
-    # Note: may not be set for all species on this process, but this process only ever
-    # sets/uses the value for the same subset of species. This means qpar_update does
-    # not need to be a shared memory array.
-    qpar_updated::Vector{Bool}
-    # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
-    vth::MPISharedArray{mk_float,3}
-    # generalised Chodura integrals for the lower and upper plates
-    chodura_integral_lower::MPISharedArray{mk_float,2}
-    chodura_integral_upper::MPISharedArray{mk_float,2}
-    # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
-    # a factor of vth for each power of wpa in velocity space integrals.
-    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
-    # and it is one otherwise
-    v_norm_fac::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the particle density
-    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the upwinded z-derivative of the particle density
-    ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the particle density
-    d2dens_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the parallel flow
-    dupar_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the upwinded z-derivative of the parallel flow
-    dupar_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the parallel flow
-    d2upar_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the parallel pressure
-    dppar_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the upwinded z-derivative of the parallel pressure
-    dppar_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the parallel pressure
-    d2ppar_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the parallel heat flux
-    dqpar_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
-    dvth_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the entropy production dS/dt = - int (ln f sum_s' C_ss' [f_s,f_s']) d^3 v
-    dSdt::MPISharedArray{mk_float,3}
-    # Spatially varying amplitude of the external source term
-    external_source_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the density moment of the external source term
-    external_source_density_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the parallel momentum moment of the external source
-    # term
-    external_source_momentum_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the parallel pressure moment of the external source
-    # term
-    external_source_pressure_amplitude::MPISharedArray{mk_float,2}
-    # Integral term for the PID controller of the external source term
-    external_source_controller_integral::MPISharedArray{mk_float,2}
-end
-
-"""
-moments_electron_substruct is a struct that contains moment information for electrons
-"""
-struct moments_electron_substruct
-    # this is the particle density
-    dens::MPISharedArray{mk_float,2}
-    # flag that keeps track of if the density needs updating before use
-    dens_updated::Ref{Bool}
-    # this is the parallel flow
-    upar::MPISharedArray{mk_float,2}
-    # flag that keeps track of whether or not upar needs updating before use
-    upar_updated::Ref{Bool}
-    # this is the parallel pressure
-    ppar::MPISharedArray{mk_float,2}
-    # flag that keeps track of whether or not ppar needs updating before use
-    ppar_updated::Ref{Bool}
-    # this is the temperature
-    temp::MPISharedArray{mk_float,2}
-    # flag that keeps track of whether or not temp needs updating before use
-    temp_updated::Ref{Bool}
-    # this is the parallel heat flux
-    qpar::MPISharedArray{mk_float,2}
-    # flag that keeps track of whether or not qpar needs updating before use
-    qpar_updated::Ref{Bool}
-    # this is the thermal speed based on the parallel temperature Tpar = ppar/dens: vth = sqrt(2*Tpar/m)
-    vth::MPISharedArray{mk_float,2}
-    # this is the parallel friction force between ions and electrons
-    parallel_friction::MPISharedArray{mk_float,2}
-    # this is the electron heat source
-    heat_source::MPISharedArray{mk_float,2}
-    # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
-    # a factor of vth for each power of wpa in velocity space integrals.
-    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
-    # and it is one otherwise
-    v_norm_fac::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the particle density
-    ddens_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the parallel flow
-    dupar_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the parallel pressure
-    dppar_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the upwinded z-derivative of the parallel pressure
-    dppar_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the second-z-derivative of the parallel pressure
-    d2ppar_dz2::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the parallel heat flux
-    dqpar_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the parallel temperature Tpar = ppar/dens
-    dT_dz::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the upwinded z-derivative of the temperature Tpar = ppar/dens
-    dT_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
-    # this is the z-derivative of the electron thermal speed vth = sqrt(2*Tpar/m)
-    dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
-end
-
-"""
-"""
-struct moments_neutral_substruct
-    # this is the particle density
-    dens::MPISharedArray{mk_float,3}
-    # flag that keeps track of if the density needs updating before use
-    # Note: may not be set for all species on this process, but this process only ever
-    # sets/uses the value for the same subset of species. This means dens_update does
-    # not need to be a shared memory array.
-    dens_updated::Vector{Bool}
-    # this is the particle mean velocity in z
-    uz::MPISharedArray{mk_float,3}
-    # flag that keeps track of if uz needs updating before use
-    uz_updated::Vector{Bool}
-    # this is the particle mean velocity in r
-    ur::MPISharedArray{mk_float,3}
-    # flag that keeps track of if ur needs updating before use
-    ur_updated::Vector{Bool}
-    # this is the particle mean velocity in zeta
-    uzeta::MPISharedArray{mk_float,3}
-    # flag that keeps track of if uzeta needs updating before use
-    uzeta_updated::Vector{Bool}
-    # this is the zz particle pressure tensor component
-    pz::MPISharedArray{mk_float,3}
-    # flag that keeps track of if pz needs updating before use
-    pz_updated::Vector{Bool}
-    # this is the rr particle pressure tensor component
-    pr::MPISharedArray{mk_float,3}
-    # flag that keeps track of if pr needs updating before use
-    pr_updated::Vector{Bool}
-    # this is the zetazeta particle pressure tensor component
-    pzeta::MPISharedArray{mk_float,3}
-    # flag that keeps track of if pzeta needs updating before use
-    pzeta_updated::Vector{Bool}
-    # this is the total (isotropic) particle pressure
-    ptot::MPISharedArray{mk_float,3}
-    # this is the heat flux along z
-    qz::MPISharedArray{mk_float,3}
-    # flag that keeps track of if qz needs updating before use
-    qz_updated::Vector{Bool}
-    # this is the thermal speed based on the temperature T = ptot/dens: vth = sqrt(2*T/m)
-    vth::MPISharedArray{mk_float,3}
-    # if evolve_ppar = true, then the velocity variable is (vz - uz)/vth, which introduces
-    # a factor of vth for each power of wz in velocity space integrals.
-    # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
-    # and it is one otherwise
-    v_norm_fac::MPISharedArray{mk_float,3}
-    # this is the z-derivative of the particle density
-    ddens_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the particle density
-    ddens_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the particle density
-    d2dens_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the particle mean velocity in z
-    duz_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the upwinded z-derivative of the particle mean velocity in z
-    duz_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the particle mean velocity in z
-    d2uz_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the zz particle pressure tensor component
-    dpz_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the upwinded z-derivative of the zz particle pressure tensor component
-    dpz_dz_upwind::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the second-z-derivative of the zz particle pressure tensor component
-    d2pz_dz2::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the thermal speed based on the temperature T = ptot/dens: vth = sqrt(2*T/m)
-    dvth_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # this is the z-derivative of the heat flux along z
-    dqz_dz::Union{MPISharedArray{mk_float,3},Nothing}
-    # Spatially varying amplitude of the external source term
-    external_source_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the density moment of the external source term
-    external_source_density_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the parallel momentum moment of the external source
-    # term
-    external_source_momentum_amplitude::MPISharedArray{mk_float,2}
-    # Spatially varying amplitude of the parallel pressure moment of the external source
-    # term
-    external_source_pressure_amplitude::MPISharedArray{mk_float,2}
-    # Integral term for the PID controller of the external source term
-    external_source_controller_integral::MPISharedArray{mk_float,2}
-end
-
 """
 """
 function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,

From b37c27fd95bfc5d382714bb2efc1a34c47514077 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Feb 2024 13:41:46 +0000
Subject: [PATCH 110/428] Don't allocate arrays for electron pdf when not using
 kinetic electrons

---
 .../src/electron_fluid_equations.jl           |  8 ++-
 moment_kinetics/src/file_io.jl                |  6 +-
 moment_kinetics/src/initial_conditions.jl     | 55 ++++++++++++-------
 .../src/moment_kinetics_structs.jl            |  2 +-
 moment_kinetics/src/time_advance.jl           |  2 +-
 5 files changed, 47 insertions(+), 26 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 8071e773f..64d999c9a 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -11,6 +11,7 @@ using ..communication
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
 using ..input_structs: braginskii_fluid, kinetic_electrons
+using ..moment_kinetics_structs: pdf_substruct
 using ..velocity_moments: integrate_over_vspace
 
 using MPI
@@ -297,7 +298,12 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth
             end
         elseif electron_model == kinetic_electrons
             # use the modified electron pdf to calculate the electron heat flux
-            calculate_electron_qpar_from_pdf!(qpar_e, ppar_e, vth_e, pdf, vpa)
+            if isa(pdf, pdf_substruct)
+                electron_pdf = pdf.norm
+            else
+                electron_pdf = pdf
+            end
+            calculate_electron_qpar_from_pdf!(qpar_e, ppar_e, vth_e, electron_pdf, vpa)
         end
     end
     # qpar has been updated
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 011b1948f..ec3aaad66 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1712,8 +1712,10 @@ function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
         # add the distribution function data at this time slice to the output file
         write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_dfns, t_idx, r, z,
                                       vperp, vpa)
-        write_electron_dfns_data_to_binary(pdf.electron.norm, io_dfns, t_idx, r, z, vperp,
-                                           vpa)
+        if pdf.electron !== nothing
+            write_electron_dfns_data_to_binary(pdf.electron.norm, io_dfns, t_idx, r, z,
+                                               vperp, vpa)
+        end
         write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species, io_dfns,
                                           t_idx, r, z, vzeta, vr, vz)
 
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 8e94ebe91..16354aeae 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -92,20 +92,24 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
     # that may be evolved separately via fluid equations
-    scratch = allocate_scratch_arrays(moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    scratch = allocate_scratch_arrays(moments, pdf, t_input.n_rk_stages)
 
     return pdf, moments, boundary_distributions, scratch
 end
 
-function allocate_scratch_arrays(moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
+function allocate_scratch_arrays(moments, pdf, n_rk_stages)
     # create n_rk_stages+1 structs, each of which will contain one pdf,
     # one density, and one parallel flow array
     scratch = Vector{scratch_pdf{5,3,4,2,6,3}}(undef, n_rk_stages+1)
-    pdf_dims = size(pdf_ion_in)
+    pdf_dims = size(pdf.ion.norm)
     moment_ion_dims = size(moments.ion.dens)
-    pdf_electron_dims = size(pdf_electron_in)
+    if pdf.electron === nothing
+        pdf_electron_dims = (0,0,0,0)
+    else
+        pdf_electron_dims = size(pdf.electron.norm)
+    end
     moment_electron_dims = size(moments.electron.dens)
-    pdf_neutral_dims = size(pdf_neutral_in)
+    pdf_neutral_dims = size(pdf.neutral.norm)
     moment_neutral_dims = size(moments.neutral.dens)
     # populate each of the structs
     for istage ∈ 1:n_rk_stages+1
@@ -152,13 +156,18 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     pdf_ion_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n, composition.n_neutral_species) # n.b. n_species is n_neutral_species here
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
-    pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
-    # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
-    # to emulate ion and neutral behaviour
-    pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
+    if composition.electron_physics == kinetic_electrons
+        pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
+        # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
+        # to emulate ion and neutral behaviour
+        pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
+        electron_substruct = pdf_substruct(pdf_electron_norm, pdf_electron_buffer)
+    else
+        electron_substruct = nothing
+    end
 
     return pdf_struct(pdf_substruct(pdf_ion_norm, pdf_ion_buffer),
-                      pdf_substruct(pdf_electron_norm, pdf_electron_buffer),
+                      electron_substruct,
                       pdf_substruct(pdf_neutral_norm, pdf_neutral_buffer))
 
 end
@@ -301,13 +310,15 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
         scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
         scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # Initialise the array for the electron pd
-    init_electron_pdf_over_density_and_boundary_phi!(
-        pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
-        moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
+    if composition.electron_physics == kinetic_electrons
+        # Initialise the array for the electron pdf
+        init_electron_pdf_over_density_and_boundary_phi!(
+            pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
+            moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
+    end
     # calculate the electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
-    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
         moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
         collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
     # calculate the zed derivative of the initial electron parallel heat flux
@@ -324,7 +335,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     # arrays need to exist and be otherwise initialised in order to compute the initial
     # electron pdf. The electron arrays will be updated as necessary by
     # initialize_electron_pdf!().
-    initialize_scratch_arrays!(scratch, moments, pdf.ion.norm, pdf.electron.norm, pdf.neutral.norm, t_input.n_rk_stages)
+    initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
     # get the initial electrostatic potential and parallel electric field
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
 
@@ -343,24 +354,26 @@ end
 initialize the array of structs containing scratch arrays for the normalised pdf and low-order moments
 that may be evolved separately via fluid equations
 """
-function initialize_scratch_arrays!(scratch, moments, pdf_ion_in, pdf_electron_in, pdf_neutral_in, n_rk_stages)
+function initialize_scratch_arrays!(scratch, moments, pdf, n_rk_stages)
     # populate each of the structs
     begin_serial_region()
     @serial_region begin
         for istage ∈ 1:n_rk_stages+1
             # initialise the scratch arrays for the ion pdf and velocity moments
-            scratch[istage].pdf .= pdf_ion_in
+            scratch[istage].pdf .= pdf.ion.norm
             scratch[istage].density .= moments.ion.dens
             scratch[istage].upar .= moments.ion.upar
             scratch[istage].ppar .= moments.ion.ppar
             # initialise the scratch arrays for the electron pdf and velocity moments
-            scratch[istage].pdf_electron .= pdf_electron_in
+            if pdf.electron !== nothing
+                scratch[istage].pdf_electron .= pdf.electron.norm
+            end
             scratch[istage].electron_density .= moments.electron.dens
             scratch[istage].electron_upar .= moments.electron.upar
             scratch[istage].electron_ppar .= moments.electron.ppar
             scratch[istage].electron_temp .= moments.electron.temp
             # initialise the scratch arrays for the neutral velocity moments and pdf
-            scratch[istage].pdf_neutral .= pdf_neutral_in
+            scratch[istage].pdf_neutral .= pdf.neutral.norm
             scratch[istage].density_neutral .= moments.neutral.dens
             scratch[istage].uz_neutral .= moments.neutral.uz
             scratch[istage].pz_neutral .= moments.neutral.pz
@@ -486,7 +499,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         end
 
         moments.electron.qpar_updated[] = false
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm, 
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
             moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
             moments.electron.dT_dz, moments.ion.upar, 
             collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index 38dd77469..97354ee5e 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -286,7 +286,7 @@ struct pdf_struct
     #ion particles: s + r + z + vperp + vpa
     ion::pdf_substruct{5}
     # electron particles: r + z + vperp + vpa
-    electron::pdf_substruct{4}
+    electron::Union{pdf_substruct{4},Nothing}
     #neutral particles: s + r + z + vzeta + vr + vz
     neutral::pdf_substruct{6}
 end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 6a41de159..107f27009 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -499,7 +499,7 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
             scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
             scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         # calculate the electron parallel heat flux
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron.norm,
+        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
             moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
             collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
         # calculate the electron-ion parallel friction force

From a1e3e9ff568f8652ef2925985234ec8b264f7112 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 19 Feb 2024 13:43:37 +0000
Subject: [PATCH 111/428] Refactor reloading functions

Define reload_moments(), etc., functions at the module level instead of
locally inside reload_evolving_fields!(). This will allow (some of)
these functions to be used separately to reload electron variables.

Also implements reloading of electron variables.
---
 moment_kinetics/src/load_data.jl | 2999 ++++++++++++++++++------------
 1 file changed, 1792 insertions(+), 1207 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index ae7066fda..b80e60cf2 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -555,68 +555,11 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
             previous_runs_info = load_run_info_history(fid)
 
             restart_n_ion_species, restart_n_neutral_species = load_species_data(fid)
-            if parallel_io
-                restart_z, restart_z_spectral, _ =
-                    load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank)
-                restart_r, restart_r_spectral, _ =
-                    load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank)
-                restart_vperp, restart_vperp_spectral, _ =
-                    load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank)
-                restart_vpa, restart_vpa_spectral, _ =
-                    load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank)
-                restart_vzeta, restart_vzeta_spectral, _ =
-                    load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank)
-                restart_vr, restart_vr_spectral, _ =
-                    load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank)
-                restart_vz, restart_vz_spectral, _ =
-                    load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank)
-            else
-                restart_z, restart_z_spectral, _ = load_coordinate_data(fid, "z")
-                restart_r, restart_r_spectral, _ = load_coordinate_data(fid, "r")
-                restart_vperp, restart_vperp_spectral, _ =
-                    load_coordinate_data(fid, "vperp")
-                restart_vpa, restart_vpa_spectral, _ = load_coordinate_data(fid, "vpa")
-                restart_vzeta, restart_vzeta_spectral, _ =
-                    load_coordinate_data(fid, "vzeta")
-                restart_vr, restart_vr_spectral, _ = load_coordinate_data(fid, "vr")
-                restart_vz, restart_vz_spectral, _ = load_coordinate_data(fid, "vz")
-
-                if restart_r.nrank != r.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now r.nrank=$(r.nrank), but we are trying to "
-                          * "restart from files ith restart_r.nrank=$(restart_r.nrank).")
-                end
-                if restart_z.nrank != z.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now z.nrank=$(z.nrank), but we are trying to "
-                          * "restart from files ith restart_z.nrank=$(restart_z.nrank).")
-                end
-                if restart_vperp.nrank != vperp.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now vperp.nrank=$(vperp.nrank), but we are trying to "
-                          * "restart from files ith restart_vperp.nrank=$(restart_vperp.nrank).")
-                end
-                if restart_vpa.nrank != vpa.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now vpa.nrank=$(vpa.nrank), but we are trying to "
-                          * "restart from files ith restart_vpa.nrank=$(restart_vpa.nrank).")
-                end
-                if restart_vzeta.nrank != vzeta.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now vzeta.nrank=$(vzeta.nrank), but we are trying to "
-                          * "restart from files ith restart_vzeta.nrank=$(restart_vzeta.nrank).")
-                end
-                if restart_vr.nrank != vr.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now vr.nrank=$(vr.nrank), but we are trying to "
-                          * "restart from files ith restart_vr.nrank=$(restart_vr.nrank).")
-                end
-                if restart_vz.nrank != vz.nrank
-                    error("Not using parallel I/O, and distributed MPI layout has "
-                          * "changed: now vz.nrank=$(vz.nrank), but we are trying to "
-                          * "restart from files ith restart_vz.nrank=$(restart_vz.nrank).")
-                end
-            end
+            restart_r, restart_r_spectral, restart_z, restart_z_spectral, restart_vperp,
+                restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
+                restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
+                restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
+                                                               vzeta, vr, vz, parallel_io)
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(
@@ -640,74 +583,41 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
 
             code_time = load_slice(dynamic, "time", time_index)
 
-            if parallel_io
-                function get_range(coord)
-                    if coord.irank == coord.nrank - 1
-                        return coord.global_io_range
-                    else
-                        # Need to modify the range to load the end-point that is duplicated on
-                        # the next process
-                        this_range = coord.global_io_range
-                        return this_range.start:(this_range.stop+1)
-                    end
-                end
-                r_range = get_range(restart_r)
-                z_range = get_range(restart_z)
-                vperp_range = get_range(restart_vperp)
-                vpa_range = get_range(restart_vpa)
-                vzeta_range = get_range(restart_vzeta)
-                vr_range = get_range(restart_vr)
-                vz_range = get_range(restart_vz)
-            else
-                r_range = (:)
-                z_range = (:)
-                vperp_range = (:)
-                vpa_range = (:)
-                vzeta_range = (:)
-                vr_range = (:)
-                vz_range = (:)
-            end
-
-            function load_moment(var_name)
-                moment = load_slice(dynamic, var_name, z_range, r_range, :, time_index)
-                orig_nz, orig_nr, nspecies = size(moment)
-                if interpolation_needed["r"]
-                    new_moment = allocate_float(orig_nz, r.n, nspecies)
-                    for is ∈ 1:nspecies, iz ∈ 1:orig_nz
-                        @views interpolate_to_grid_1d!(new_moment[iz,:,is], r.grid,
-                                                       moment[iz,:,is], restart_r,
-                                                       restart_r_spectral)
-                    end
-                    moment = new_moment
-                end
-                if interpolation_needed["z"]
-                    new_moment = allocate_float(z.n, r.n, nspecies)
-                    for is ∈ 1:nspecies, ir ∈ 1:r.n
-                        @views interpolate_to_grid_1d!(new_moment[:,ir,is], z.grid,
-                                                       moment[:,ir,is], restart_z,
-                                                       restart_z_spectral)
-                    end
-                    moment = new_moment
-                end
-                return moment
-            end
+            r_range, z_range, vperp_range, vpa_range, vzeta_range, vr_range, vz_range =
+                get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp,
+                                  restart_vpa, restart_vzeta, restart_vr, restart_vz)
 
-            moments.ion.dens .= load_moment("density")
+            moments.ion.dens .= reload_moment("density", dynamic, time_index, r, z,
+                                              r_range, z_range, restart_r,
+                                              restart_r_spectral, restart_z,
+                                              restart_z_spectral, interpolation_needed)
             moments.ion.dens_updated .= true
-            moments.ion.upar .= load_moment("parallel_flow")
+            moments.ion.upar .= reload_moment("parallel_flow", dynamic, time_index, r, z,
+                                              r_range, z_range, restart_r,
+                                              restart_r_spectral, restart_z,
+                                              restart_z_spectral, interpolation_needed)
             moments.ion.upar_updated .= true
-            moments.ion.ppar .= load_moment("parallel_pressure")
+            moments.ion.ppar .= reload_moment("parallel_pressure", dynamic, time_index, r,
+                                              z, r_range, z_range, restart_r,
+                                              restart_r_spectral, restart_z,
+                                              restart_z_spectral, interpolation_needed)
             moments.ion.ppar_updated .= true
-            moments.ion.qpar .= load_moment("parallel_heat_flux")
+            moments.ion.qpar .= reload_moment("parallel_heat_flux", dynamic, time_index,
+                                              r, z, r_range, z_range, restart_r,
+                                              restart_r_spectral, restart_z,
+                                              restart_z_spectral, interpolation_needed)
             moments.ion.qpar_updated .= true
-            moments.ion.vth .= load_moment("thermal_speed")
+            moments.ion.vth .= reload_moment("thermal_speed", dynamic, time_index, r, z,
+                                             r_range, z_range, restart_r,
+                                             restart_r_spectral, restart_z,
+                                             restart_z_spectral, interpolation_needed)
             if parallel_io || z.irank == 0
                 moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower",
-                                                                     r_range, :, time_index)
+                                                                 r_range, :, time_index)
             end
             if parallel_io || z.irank == z.nrank - 1
                 moments.ion.chodura_integral_upper .= load_slice(dynamic, "chodura_integral_upper",
-                                                                     r_range, :, time_index)
+                                                                 r_range, :, time_index)
             end
 
             if "external_source_controller_integral" ∈ get_variable_keys(dynamic) &&
@@ -716,542 +626,135 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                     load_slice(dynamic, "external_source_controller_integral", time_index)
             elseif length(moments.ion.external_source_controller_integral) > 1
                 moments.ion.external_source_controller_integral .=
-                    load_moment("external_source_controller_integral")
+                    reload_moment("external_source_controller_integral", dynamic,
+                                  time_index, r, z, r_range, z_range, restart_r,
+                                  restart_r_spectral, restart_z, restart_z_spectral,
+                                  interpolation_needed)
             end
 
-            function load_ion_pdf()
-                this_pdf = load_slice(dynamic, "f", vpa_range, vperp_range, z_range,
-                                      r_range, :, time_index)
-                orig_nvpa, orig_nvperp, orig_nz, orig_nr, nspecies = size(this_pdf)
-                if interpolation_needed["r"]
-                    new_pdf = allocate_float(orig_nvpa, orig_nvperp, orig_nz, r.n, nspecies)
-                    for is ∈ 1:nspecies, iz ∈ 1:orig_nz, ivperp ∈ 1:orig_nvperp,
-                            ivpa ∈ 1:orig_nvpa
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[ivpa,ivperp,iz,:,is], r.grid,
-                            this_pdf[ivpa,ivperp,iz,:,is], restart_r,
-                            restart_r_spectral)
-                    end
-                    this_pdf = new_pdf
-                end
-                if interpolation_needed["z"]
-                    new_pdf = allocate_float(orig_nvpa, orig_nvperp, z.n, r.n, nspecies)
-                    for is ∈ 1:nspecies, ir ∈ 1:r.n, ivperp ∈ 1:orig_nvperp,
-                            ivpa ∈ 1:orig_nvpa
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[ivpa,ivperp,:,ir,is], z.grid,
-                            this_pdf[ivpa,ivperp,:,ir,is], restart_z,
-                            restart_z_spectral)
-                    end
-                    this_pdf = new_pdf
-                end
-
-                # Current moment-kinetic implementation is only 1V, so no need to handle a
-                # normalised vperp coordinate. This will need to change when 2V
-                # moment-kinetics is implemented.
-                if interpolation_needed["vperp"]
-                    new_pdf = allocate_float(orig_nvpa, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivpa ∈ 1:orig_nvpa
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[ivpa,:,iz,ir,is], vperp.grid,
-                            this_pdf[ivpa,:,iz,ir,is], restart_vperp,
-                            restart_vperp_spectral)
-                    end
-                    this_pdf = new_pdf
-                end
-
-                if (
-                    (moments.evolve_density == restart_evolve_density &&
-                     moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
-                     restart_evolve_ppar)
-                    || (!moments.evolve_upar && !restart_evolve_upar &&
-                        !moments.evolve_ppar && !restart_evolve_ppar)
-                   )
-                    # No chages to velocity-space coordinates, so just interpolate from
-                    # one grid to the other
-                    if interpolation_needed["vpa"]
-                        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivperp,iz,ir,is], vpa.grid,
-                                this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                                restart_vpa_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa + upar
-                    # => old_wpa = new_wpa - upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa*vth
-                    # => old_wpa = new_wpa/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa*vth + upar
-                    # => old_wpa = (new_wpa - upar)/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. (vpa.grid - moments.ion.upar[iz,ir,is]) /
-                               moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa
-                    # => old_wpa = new_wpa + upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa*vth
-                    # => old_wpa = (new_wpa + upar)/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. (vpa.grid + moments.ion.upar[iz,ir,is]) /
-                               moments.ion.vth
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa*vth + upar
-                    # => old_wpa = new_wpa/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa
-                    # => old_wpa = new_wpa*vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa + upar
-                    # => old_wpa = new_wpa*vth - upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] -
-                                              moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa*vth + upar
-                    # => old_wpa = new_wpa - upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. vpa.grid -
-                               moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa
-                    # => old_wpa = new_wpa*vth + upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] +
-                                              moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa + upar
-                    # => old_wpa = new_wpa*vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa*vth
-                    # => old_wpa = new_wpa + upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. vpa.grid +
-                               moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,ir,is], restart_vpa,
-                            restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                else
-                    # This should never happen, as all combinations of evolve_* options
-                    # should be handled above.
-                    error("Unsupported combination of moment-kinetic options:"
-                          * " evolve_density=", moments.evolve_density
-                          * " evolve_upar=", moments.evolve_upar
-                          * " evolve_ppar=", moments.evolve_ppar
-                          * " restart_evolve_density=", restart_evolve_density
-                          * " restart_evolve_upar=", restart_evolve_upar
-                          * " restart_evolve_ppar=", restart_evolve_ppar)
-                end
-                if moments.evolve_density && !restart_evolve_density
-                    # Need to normalise by density
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] ./= moments.ion.dens[iz,ir,is]
-                    end
-                elseif !moments.evolve_density && restart_evolve_density
-                    # Need to unnormalise by density
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] .*= moments.ion.dens[iz,ir,is]
-                    end
-                end
-                if moments.evolve_ppar && !restart_evolve_ppar
-                    # Need to normalise by vth
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] .*= moments.ion.vth[iz,ir,is]
-                    end
-                elseif !moments.evolve_ppar && restart_evolve_ppar
-                    # Need to unnormalise by vth
-                    for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,ir,is] ./= moments.ion.vth[iz,ir,is]
-                    end
-                end
-
-                return this_pdf
+            pdf.ion.norm .= reload_ion_pdf(dynamic, time_index, moments, r, z, vperp, vpa, r_range,
+                                           z_range, vperp_range, vpa_range, restart_r,
+                                           restart_r_spectral, restart_z,
+                                           restart_z_spectral, restart_vperp,
+                                           restart_vperp_spectral, restart_vpa,
+                                           restart_vpa_spectral, interpolation_needed,
+                                           restart_evolve_density, restart_evolve_upar,
+                                           restart_evolve_ppar)
+            if z.irank == 0
+                moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower", :, :,
+                                                                 time_index)
+            else
+                moments.ion.chodura_integral_lower .= 0.0
             end
-
-            pdf.ion.norm .= load_ion_pdf()
-                if z.irank == 0
-                    moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower", :, :,
-                                                  time_index)
-                else
-                    moments.ion.chodura_integral_lower .= 0.0
-                end
-                if z.irank == z.nrank - 1
-                    moments.ion.chodura_integral_upper .= load_slice(dynamic, "chodura_integral_upper", :, :,
-                                                  time_index)
-                else
-                    moments.ion.chodura_integral_upper .= 0.0
-                end
-            boundary_distributions_io = get_group(fid, "boundary_distributions")
-
-            function load_ion_boundary_pdf(var_name, ir)
-                this_pdf = load_slice(boundary_distributions_io, var_name, vpa_range,
-                                      vperp_range, z_range, :)
-                orig_nvpa, orig_nvperp, orig_nz, nspecies = size(this_pdf)
-                if interpolation_needed["z"]
-                    new_pdf = allocate_float(orig_nvpa, orig_nvperp, z.n, nspecies)
-                    for is ∈ 1:nspecies, ivperp ∈ 1:orig_nvperp,
-                            ivpa ∈ 1:orig_nvpa
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[ivpa,ivperp,:,is], z.grid,
-                            this_pdf[ivpa,ivperp,:,is], restart_z,
-                            restart_z_spectral)
-                    end
-                    this_pdf = new_pdf
-                end
-
-                # Current moment-kinetic implementation is only 1V, so no need to handle a
-                # normalised vperp coordinate. This will need to change when 2V
-                # moment-kinetics is implemented.
-                if interpolation_needed["vperp"]
-                    new_pdf = allocate_float(orig_nvpa, vperp.n, z.n, nspecies)
-                    for is ∈ 1:nspecies, iz ∈ 1:z.n, ivpa ∈ 1:orig_nvpa
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[ivpa,:,iz,is], vperp.grid,
-                            this_pdf[ivpa,:,iz,is], restart_vperp,
-                            restart_vperp_spectral)
-                    end
-                    this_pdf = new_pdf
-                end
-
-                if (
-                    (moments.evolve_density == restart_evolve_density &&
-                     moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
-                     restart_evolve_ppar)
-                    || (!moments.evolve_upar && !restart_evolve_upar &&
-                        !moments.evolve_ppar && !restart_evolve_ppar)
-                   )
-                    # No chages to velocity-space coordinates, so just interpolate from
-                    # one grid to the other
-                    if interpolation_needed["vpa"]
-                        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivperp,iz,is], vpa.grid,
-                                this_pdf[:,ivperp,iz,is], restart_vpa,
-                                restart_vpa_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa + upar
-                    # => old_wpa = new_wpa - upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa*vth
-                    # => old_wpa = new_wpa/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa = old_wpa*vth + upar
-                    # => old_wpa = (new_wpa - upar)/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. (vpa.grid - moments.ion.upar[iz,ir,is]) /
-                               moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa
-                    # => old_wpa = new_wpa + upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa*vth
-                    # => old_wpa = (new_wpa + upar)/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. (vpa.grid + moments.ion.upar[iz,ir,is]) /
-                               moments.ion.vth
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa + upar = old_wpa*vth + upar
-                    # => old_wpa = new_wpa/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa
-                    # => old_wpa = new_wpa*vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa + upar
-                    # => old_wpa = new_wpa*vth - upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] -
-                                              moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa*vth = old_wpa*vth + upar
-                    # => old_wpa = new_wpa - upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. vpa.grid -
-                               moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa
-                    # => old_wpa = new_wpa*vth + upar
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] +
-                                              moments.ion.upar[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        restart_evolve_upar && !restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa + upar
-                    # => old_wpa = new_wpa*vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                elseif (moments.evolve_upar && moments.evolve_ppar &&
-                        !restart_evolve_upar && restart_evolve_ppar)
-                    # vpa = new_wpa*vth + upar = old_wpa*vth
-                    # => old_wpa = new_wpa + upar/vth
-                    new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
-                    for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
-                        restart_vpa_vals =
-                            @. vpa.grid +
-                               moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
-                        @views interpolate_to_grid_1d!(
-                            new_pdf[:,ivperp,iz,is], restart_vpa_vals,
-                            this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
-                    end
-                    this_pdf = new_pdf
-                else
-                    # This should never happen, as all combinations of evolve_* options
-                    # should be handled above.
-                    error("Unsupported combination of moment-kinetic options:"
-                          * " evolve_density=", moments.evolve_density
-                          * " evolve_upar=", moments.evolve_upar
-                          * " evolve_ppar=", moments.evolve_ppar
-                          * " restart_evolve_density=", restart_evolve_density
-                          * " restart_evolve_upar=", restart_evolve_upar
-                          * " restart_evolve_ppar=", restart_evolve_ppar)
-                end
-                if moments.evolve_density && !restart_evolve_density
-                    # Need to normalise by density
-                    for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] ./= moments.ion.dens[iz,ir,is]
-                    end
-                elseif !moments.evolve_density && restart_evolve_density
-                    # Need to unnormalise by density
-                    for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] .*= moments.ion.dens[iz,ir,is]
-                    end
-                end
-                if moments.evolve_ppar && !restart_evolve_ppar
-                    # Need to normalise by vth
-                    for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] .*= moments.ion.vth[iz,ir,is]
-                    end
-                elseif !moments.evolve_ppar && restart_evolve_ppar
-                    # Need to unnormalise by vth
-                    for is ∈ nspecies, iz ∈ 1:z.n
-                        this_pdf[:,:,iz,is] ./= moments.ion.vth[iz,ir,is]
-                    end
-                end
-
-                return this_pdf
+            if z.irank == z.nrank - 1
+                moments.ion.chodura_integral_upper .= load_slice(dynamic, "chodura_integral_upper", :, :,
+                                                                 time_index)
+            else
+                moments.ion.chodura_integral_upper .= 0.0
             end
+            boundary_distributions_io = get_group(fid, "boundary_distributions")
 
             boundary_distributions.pdf_rboundary_ion[:,:,:,1,:] .=
-                load_ion_boundary_pdf("pdf_rboundary_ion_left", 1)
+                reload_ion_boundary_pdf(boundary_distributions_io,
+                                        "pdf_rboundary_ion_left", 1, moments, z, vperp,
+                                        vpa, z_range, vperp_range, vpa_range, restart_z,
+                                        restart_z_spectral, restart_vperp,
+                                        restart_vperp_spectral, restart_vpa,
+                                        restart_vpa_spectral, interpolation_needed,
+                                        restart_evolve_density, restart_evolve_upar,
+                                        restart_evolve_ppar)
             boundary_distributions.pdf_rboundary_ion[:,:,:,2,:] .=
-                load_ion_boundary_pdf("pdf_rboundary_ion_right", r.n)
+                reload_ion_boundary_pdf(boundary_distributions_io,
+                                        "pdf_rboundary_ion_right", r.n, moments, z, vperp,
+                                        vpa, z_range, vperp_range, vpa_range, restart_z,
+                                        restart_z_spectral, restart_vperp,
+                                        restart_vperp_spectral, restart_vpa,
+                                        restart_vpa_spectral, interpolation_needed,
+                                        restart_evolve_density, restart_evolve_upar,
+                                        restart_evolve_ppar)
+
+            moments.electron.dens .= reload_electron_moment("electron_density", dynamic,
+                                                            time_index, r, z, r_range,
+                                                            z_range, restart_r,
+                                                            restart_r_spectral, restart_z,
+                                                            restart_z_spectral,
+                                                            interpolation_needed)
+            moments.electron.dens_updated[] = true
+            moments.electron.upar .= reload_electron_moment("electron_parallel_flow",
+                                                            dynamic, time_index, r, z,
+                                                            r_range, z_range, restart_r,
+                                                            restart_r_spectral, restart_z,
+                                                            restart_z_spectral,
+                                                            interpolation_needed)
+            moments.electron.upar_updated[] = true
+            moments.electron.ppar .= reload_electron_moment("electron_parallel_pressure",
+                                                            dynamic, time_index, r, z,
+                                                            r_range, z_range, restart_r,
+                                                            restart_r_spectral, restart_z,
+                                                            restart_z_spectral,
+                                                            interpolation_needed)
+            moments.electron.ppar_updated[] = true
+            moments.electron.qpar .= reload_electron_moment("electron_parallel_heat_flux",
+                                                            dynamic, time_index, r, z,
+                                                            r_range, z_range, restart_r,
+                                                            restart_r_spectral, restart_z,
+                                                            restart_z_spectral,
+                                                            interpolation_needed)
+            moments.electron.qpar_updated[] = true
+            moments.electron.vth .= reload_electron_moment("electron_thermal_speed",
+                                                           dynamic, time_index, r, z,
+                                                           r_range, z_range, restart_r,
+                                                           restart_r_spectral, restart_z,
+                                                           restart_z_spectral,
+                                                           interpolation_needed)
+
+            # For now, electrons are always fully moment_kinetic
+            restart_electron_evolve_density, restart_electron_evolve_upar,
+                restart_electron_evolve_ppar = true, true, true
+            electron_evolve_density, electron_evolve_upar, electron_evolve_ppar =
+                true, true, true
+            if pdf.electron !== nothing
+                pdf.electron.norm .=
+                    reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
+                                        r_range, z_range, vperp_range, vpa_range,
+                                        restart_r, restart_r_spectral, restart_z,
+                                        restart_z_spectral, restart_vperp,
+                                        restart_vperp_spectral, restart_vpa,
+                                        restart_vpa_spectral, interpolation_needed,
+                                        restart_evolve_density, restart_evolve_upar,
+                                        restart_evolve_ppar)
+            end
 
             if composition.n_neutral_species > 0
-                moments.neutral.dens .= load_moment("density_neutral")
+                moments.neutral.dens .= reload_moment("density_neutral", dynamic,
+                                                      time_index, r, z, r_range, z_range,
+                                                      restart_r, restart_r_spectral,
+                                                      restart_z, restart_z_spectral,
+                                                      interpolation_needed)
                 moments.neutral.dens_updated .= true
-                moments.neutral.uz .= load_moment("uz_neutral")
+                moments.neutral.uz .= reload_moment("uz_neutral", dynamic, time_index, r,
+                                                    z, r_range, z_range, restart_r,
+                                                    restart_r_spectral, restart_z,
+                                                    restart_z_spectral,
+                                                    interpolation_needed)
                 moments.neutral.uz_updated .= true
-                moments.neutral.pz .= load_moment("pz_neutral")
+                moments.neutral.pz .= reload_moment("pz_neutral", dynamic, time_index, r,
+                                                    z, r_range, z_range, restart_r,
+                                                    restart_r_spectral, restart_z,
+                                                    restart_z_spectral,
+                                                    interpolation_needed)
                 moments.neutral.pz_updated .= true
-                moments.neutral.qz .= load_moment("qz_neutral")
+                moments.neutral.qz .= reload_moment("qz_neutral", dynamic, time_index, r,
+                                                    z, r_range, z_range, restart_r,
+                                                    restart_r_spectral, restart_z,
+                                                    restart_z_spectral,
+                                                    interpolation_needed)
                 moments.neutral.qz_updated .= true
-                moments.neutral.vth .= load_moment("thermal_speed_neutral")
+                moments.neutral.vth .= reload_moment("thermal_speed_neutral", dynamic,
+                                                     time_index, r, z, r_range, z_range,
+                                                     restart_r, restart_r_spectral,
+                                                     restart_z, restart_z_spectral,
+                                                     interpolation_needed)
 
                 if "external_source_neutral_controller_integral" ∈ get_variable_keys(dynamic) &&
                         length(moments.neutral.external_source_controller_integral) == 1
@@ -1261,580 +764,151 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                                    time_index)
                 elseif length(moments.neutral.external_source_controller_integral) > 1
                     moments.neutral.external_source_controller_integral .=
-                        load_moment("external_source_neutral_controller_integral")
+                        reload_moment("external_source_neutral_controller_integral",
+                                      dynamic, time_index, r, z, r_range, z_range,
+                                      restart_r, restart_r_spectral, restart_z,
+                                      restart_z_spectral, interpolation_needed)
                 end
 
-                function load_neutral_pdf()
-                    this_pdf = load_slice(dynamic, "f_neutral", vz_range, vr_range,
-                                          vzeta_range, z_range, r_range, :, time_index)
-                    orig_nvz, orig_nvr, orig_nvzeta, orig_nz, orig_nr, nspecies =
-                        size(this_pdf)
-                    if interpolation_needed["r"]
-                        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, orig_nz,
-                                                 r.n, nspecies)
-                        for is ∈ 1:nspecies, iz ∈ 1:orig_nz, ivzeta ∈ 1:orig_nvzeta,
-                                ivr ∈ 1:orig_nvr, ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,ivr,ivzeta,iz,:,is], r.grid,
-                                this_pdf[ivz,ivr,ivzeta,iz,:,is], restart_r,
-                                restart_r_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-                    if interpolation_needed["z"]
-                        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, z.n,
-                                                 r.n, nspecies)
-                        for is ∈ 1:nspecies, ir ∈ 1:r.n, ivzeta ∈ 1:orig_nvzeta,
-                                ivr ∈ 1:orig_nvr, ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,ivr,ivzeta,:,ir,is], z.grid,
-                                this_pdf[ivz,ivr,ivzeta,:,ir,is], restart_z,
-                                restart_z_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-
-                    # Current moment-kinetic implementation is only 1V, so no need
-                    # to handle normalised vzeta or vr coordinates. This will need
-                    # to change when/if 3V moment-kinetics is implemented.
-                    if interpolation_needed["vzeta"]
-                        new_pdf = allocate_float(orig_nvz, orig_nvr, vzeta.n, z.n, r.n,
-                                                 nspecies)
-                        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:orig_nvr,
-                                ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,ivr,:,iz,ir,is], vzeta.grid,
-                                this_pdf[ivz,ivr,:,iz,ir,is], restart_vzeta,
-                                restart_vzeta_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-                    if interpolation_needed["vr"]
-                        new_pdf = allocate_float(orig_nvz, vr.n, vzeta.n, z.n, r.n,
-                                                 nspecies)
-                        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n,
-                                ivzeta ∈ 1:vzeta.n, ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,:,ivzeta,iz,ir,is], vr.grid,
-                                this_pdf[ivz,:,ivzeta,iz,ir,is], restart_vr,
-                                restart_vr_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
+                pdf.neutral.norm .=
+                    reload_neutral_pdf(dynamic, time_index, moments, r, z, vzeta, vr, vz,
+                                       r_range, z_range, vzeta_range, vr_range, vz_range,
+                                       restart_r, restart_r_spectral, restart_z,
+                                       restart_z_spectral, restart_vzeta,
+                                       restart_vzeta_spectral, restart_vr,
+                                       restart_vr_spectral, restart_vz,
+                                       restart_vz_spectral, interpolation_needed,
+                                       restart_evolve_density, restart_evolve_upar,
+                                       restart_evolve_ppar)
 
-                    if (
-                        (moments.evolve_density == restart_evolve_density &&
-                         moments.evolve_upar == restart_evolve_upar &&
-                         moments.evolve_ppar == restart_evolve_ppar)
-                        || (!moments.evolve_upar && !restart_evolve_upar &&
-                            !moments.evolve_ppar && !restart_evolve_ppar)
-                       )
-                        # No chages to velocity-space coordinates, so just interpolate from
-                        # one grid to the other
-                        if interpolation_needed["vz"]
-                            new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                            for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                    ivzeta ∈ 1:vzeta.n
-                                @views interpolate_to_grid_1d!(
-                                    new_pdf[:,ivr,ivzeta,iz,ir,is], vz.grid,
-                                    this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                    restart_vz_spectral)
-                            end
-                            this_pdf = new_pdf
-                        end
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa + upar
-                        # => old_wpa = new_wpa - upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n,
-                                ivr ∈ 1:vr.n
-                            restart_vz_vals = vz.grid .- moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa*vth
-                        # => old_wpa = new_wpa/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa*vth + upar
-                        # => old_wpa = (new_wpa - upar)/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. (vz.grid - moments.neutral.uz[iz,ir,is]) /
-                                   moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa
-                        # => old_wpa = new_wpa + upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .+ moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz, restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa*vth
-                        # => old_wpa = (new_wpa + upar)/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. (vz.grid + moments.neutral.uz[iz,ir,is]) /
-                                    moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa*vth + upar
-                        # => old_wpa = new_wpa/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa
-                        # => old_wpa = new_wpa*vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa + upar
-                        # => old_wpa = new_wpa*vth - upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid * moments.neutral.vth[iz,ir,is] -
-                                   moments.neutral.upar[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa*vth + upar
-                        # => old_wpa = new_wpa - upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid -
-                                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa
-                        # => old_wpa = new_wpa*vth + upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid * moments.neutral.vth[iz,ir,is] +
-                                   moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa + upar
-                        # => old_wpa = new_wpa*vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa*vth
-                        # => old_wpa = new_wpa + upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid +
-                                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    else
-                        # This should never happen, as all combinations of evolve_* options
-                        # should be handled above.
-                        error("Unsupported combination of moment-kinetic options:"
-                              * " evolve_density=", moments.evolve_density
-                              * " evolve_upar=", moments.evolve_upar
-                              * " evolve_ppar=", moments.evolve_ppar
-                              * " restart_evolve_density=", restart_evolve_density
-                              * " restart_evolve_upar=", restart_evolve_upar
-                              * " restart_evolve_ppar=", restart_evolve_ppar)
-                    end
-                    if moments.evolve_density && !restart_evolve_density
-                        # Need to normalise by density
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,ir,is] ./= moments.neutral.dens[iz,ir,is]
-                        end
-                    elseif !moments.evolve_density && restart_evolve_density
-                        # Need to unnormalise by density
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,ir,is] .*= moments.neutral.dens[iz,ir,is]
-                        end
-                    end
-                    if moments.evolve_ppar && !restart_evolve_ppar
-                        # Need to normalise by vth
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,ir,is] .*= moments.neutral.vth[iz,ir,is]
-                        end
-                    elseif !moments.evolve_ppar && restart_evolve_ppar
-                        # Need to unnormalise by vth
-                        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,ir,is] ./= moments.neutral.vth[iz,ir,is]
-                        end
-                    end
+                boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:] .=
+                    reload_neutral_boundary_pdf(boundary_distributions_io,
+                                                "pdf_rboundary_neutral_left", 1, moments,
+                                                z, vzeta, vr, vz, z_range, vzeta_range,
+                                                vr_range, vz_range, restart_z,
+                                                restart_z_spectral, restart_vzeta,
+                                                restart_vzeta_spectral, restart_vr,
+                                                restart_vr_spectral, restart_vz,
+                                                restart_vz_spectral, interpolation_needed,
+                                                restart_evolve_density,
+                                                restart_evolve_upar, restart_evolve_ppar)
+                boundary_distributions.pdf_rboundary_neutral[:,:,:,:,2,:] .=
+                    reload_neutral_boundary_pdf(boundary_distributions_io,
+                                                "pdf_rboundary_neutral_right", r.n,
+                                                moments, z, vzeta, vr, vz, z_range,
+                                                vzeta_range, vr_range, vz_range,
+                                                restart_z, restart_z_spectral,
+                                                restart_vzeta, restart_vzeta_spectral,
+                                                restart_vr, restart_vr_spectral,
+                                                restart_vz, restart_vz_spectral,
+                                                interpolation_needed,
+                                                restart_evolve_density,
+                                                restart_evolve_upar, restart_evolve_ppar)
+            end
+        finally
+            close(fid)
+        end
+    end
 
-                    return this_pdf
-                end
+    return code_time, previous_runs_info, time_index
+end
 
-                pdf.neutral.norm .= load_neutral_pdf()
-
-                function load_neutral_boundary_pdf(var_name, ir)
-                    this_pdf = load_slice(boundary_distributions_io, var_name, vz_range,
-                                          vr_range, vzeta_range, z_range, :)
-                    orig_nvz, orig_nvr, orig_nvzeta, orig_nz, nspecies = size(this_pdf)
-                    if interpolation_needed["z"]
-                        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, z.n,
-                                                 nspecies)
-                        for is ∈ 1:nspecies, ivzeta ∈ 1:orig_nvzeta, ivr ∈ 1:orig_nvr,
-                                ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,ivr,ivzeta,:,is], z.grid,
-                                this_pdf[ivz,ivr,ivzeta,:,is], restart_z,
-                                restart_z_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
+"""
+Reload electron pdf and moments from an existing output file.
+"""
+function reload_electron_data!(pdf, moments, restart_prefix_iblock, time_index, geometry,
+                               r, z, vpa, vperp, vzeta, vr, vz)
+    code_time = 0.0
+    previous_runs_info = nothing
+    begin_serial_region()
+    @serial_region begin
+        fid = open_readonly_output_file(restart_prefix_iblock[1], "dfns";
+                                        iblock=restart_prefix_iblock[2])
+        try # finally to make sure to close file0
+            overview = get_group(fid, "overview")
+            dynamic = get_group(fid, "dynamic_data")
+            parallel_io = load_variable(overview, "parallel_io")
+            if time_index < 0
+                time_index, _, _ = load_time_data(fid)
+            end
+            #restart_evolve_density, restart_evolve_upar, restart_evolve_ppar =
+            #    load_mk_options(fid)
+            # For now, electrons are always fully moment_kinetic
+            restart_evolve_density, restart_evolve_upar, restart_evolve_ppar = true, true,
+                                                                               true
+            evolve_density, evolve_upar, evolve_ppar = true, true, true
 
-                    # Current moment-kinetic implementation is only 1V, so no need
-                    # to handle normalised vzeta or vr coordinates. This will need
-                    # to change when/if 3V moment-kinetics is implemented.
-                    if interpolation_needed["vzeta"]
-                        new_pdf = allocate_float(orig_nvz, orig_nvr, vzeta.n, z.n,
-                                                 nspecies)
-                        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivr ∈ 1:orig_nvr,
-                                ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,ivr,:,iz,is], vzeta.grid,
-                                this_pdf[ivz,ivr,:,iz,is], restart_vzeta,
-                                restart_vzeta_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
-                    if interpolation_needed["vr"]
-                        new_pdf = allocate_float(orig_nvz, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n,
-                                ivz ∈ 1:orig_nvz
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[ivz,:,ivzeta,iz,is], vr.grid,
-                                this_pdf[ivz,:,ivzeta,iz,is], restart_vr,
-                                restart_vr_spectral)
-                        end
-                        this_pdf = new_pdf
-                    end
+            previous_runs_info = load_run_info_history(fid)
 
-                    if (
-                        (moments.evolve_density == restart_evolve_density &&
-                         moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
-                         restart_evolve_ppar)
-                        || (!moments.evolve_upar && !restart_evolve_upar &&
-                            !moments.evolve_ppar && !restart_evolve_ppar)
-                       )
-                        # No chages to velocity-space coordinates, so just interpolate from
-                        # one grid to the other
-                        if interpolation_needed["vz"]
-                            new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                            for is ∈ 1:nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
-                                    ivzeta ∈ 1:vzeta.n
-                                @views interpolate_to_grid_1d!(
-                                    new_pdf[:,ivr,ivzeta,iz,is], vz.grid,
-                                    this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                    restart_vz_spectral)
-                            end
-                            this_pdf = new_pdf
-                        end
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa + upar
-                        # => old_wpa = new_wpa - upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n, ivr ∈ 1:vr.n
-                            restart_vz_vals = vz.grid .- moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa*vth
-                        # => old_wpa = new_wpa/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa = old_wpa*vth + upar
-                        # => old_wpa = (new_wpa - upar)/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. (vz.grid - moments.neutral.uz[iz,ir,is]) /
-                                   moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa
-                        # => old_wpa = new_wpa + upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .+
-                                              moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa*vth
-                        # => old_wpa = (new_wpa + upar)/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. (vz.grid + moments.neutral.uz[iz,ir,is]) /
-                                   moments.neutral.vth
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && !moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa + upar = old_wpa*vth + upar
-                        # => old_wpa = new_wpa/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa
-                        # => old_wpa = new_wpa*vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .*
-                                              moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa + upar
-                        # => old_wpa = new_wpa*vth - upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid * moments.neutral.vth[iz,ir,is] -
-                                   moments.neutral.upar[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (!moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa*vth = old_wpa*vth + upar
-                        # => old_wpa = new_wpa - upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid -
-                                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa
-                        # => old_wpa = new_wpa*vth + upar
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid * moments.neutral.vth[iz,ir,is] +
-                                   moments.neutral.uz[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            restart_evolve_upar && !restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa + upar
-                        # => old_wpa = new_wpa*vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                        this_pdf = new_pdf
-                    elseif (moments.evolve_upar && moments.evolve_ppar &&
-                            !restart_evolve_upar && restart_evolve_ppar)
-                        # vpa = new_wpa*vth + upar = old_wpa*vth
-                        # => old_wpa = new_wpa + upar/vth
-                        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
-                        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
-                            restart_vz_vals =
-                                @. vz.grid +
-                                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
-                            @views interpolate_to_grid_1d!(
-                                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
-                                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
-                                restart_vz_spectral)
-                        end
-                    else
-                        # This should never happen, as all combinations of evolve_* options
-                        # should be handled above.
-                        error("Unsupported combination of moment-kinetic options:"
-                              * " evolve_density=", moments.evolve_density
-                              * " evolve_upar=", moments.evolve_upar
-                              * " evolve_ppar=", moments.evolve_ppar
-                              * " restart_evolve_density=", restart_evolve_density
-                              * " restart_evolve_upar=", restart_evolve_upar
-                              * " restart_evolve_ppar=", restart_evolve_ppar)
-                    end
-                    if moments.evolve_density && !restart_evolve_density
-                        # Need to normalise by density
-                        for is ∈ nspecies, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,is] ./= moments.neutral.dens[iz,ir,is]
-                        end
-                    elseif !moments.evolve_density && restart_evolve_density
-                        # Need to unnormalise by density
-                        for is ∈ nspecies, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,is] .*= moments.neutral.dens[iz,ir,is]
-                        end
-                    end
-                    if moments.evolve_ppar && !restart_evolve_ppar
-                        # Need to normalise by vth
-                        for is ∈ nspecies, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,is] .*= moments.neutral.vth[iz,ir,is]
-                        end
-                    elseif !moments.evolve_ppar && restart_evolve_ppar
-                        # Need to unnormalise by vth
-                        for is ∈ nspecies, iz ∈ 1:z.n
-                            this_pdf[:,:,:,iz,is] ./= moments.neutral.vth[iz,ir,is]
-                        end
-                    end
+            restart_n_ion_species, restart_n_neutral_species = load_species_data(fid)
+            restart_r, restart_r_spectral, restart_z, restart_z_spectral, restart_vperp,
+                restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
+                restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
+                restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
+                                                               vzeta, vr, vz, parallel_io)
 
-                    return this_pdf
-                end
+            # Test whether any interpolation is needed
+            interpolation_needed = Dict(
+                x.name => x.n != restart_x.n || !all(isapprox.(x.grid, restart_x.grid))
+                for (x, restart_x) ∈ ((z, restart_z), (r, restart_r),
+                                      (vperp, restart_vperp), (vpa, restart_vpa)))
 
-                boundary_distributions.pdf_rboundary_neutral[:,:,:,:,1,:] .=
-                    load_neutral_boundary_pdf("pdf_rboundary_neutral_left", 1)
-                boundary_distributions.pdf_rboundary_neutral[:,:,:,:,2,:] .=
-                    load_neutral_boundary_pdf("pdf_rboundary_neutral_right", r.n)
+            neutral_1V = (vzeta.n_global == 1 && vr.n_global == 1)
+            restart_neutral_1V = (restart_vzeta.n_global == 1 && restart_vr.n_global == 1)
+            if geometry.bzeta != 0.0 && ((neutral1V && !restart_neutral_1V) ||
+                                         (!neutral1V && restart_neutral_1V))
+                # One but not the other of the run being restarted from and this run are
+                # 1V, but the interpolation below does not allow for vz and vpa being in
+                # different directions. Therefore interpolation between 1V and 3V cases
+                # only works (at the moment!) if bzeta=0.
+                error("Interpolation between 1V and 3V neutrals not yet supported when "
+                      * "bzeta!=0.")
             end
+
+            code_time = load_slice(dynamic, "time", time_index)
+
+            r_range, z_range, vperp_range, vpa_range, vzeta_range, vr_range, vz_range =
+                get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp,
+                                  restart_vpa, restart_vzeta, restart_vr, restart_vz)
+
+            moments.electron.dens .=
+                reload_electron_moment("electron_density", dynamic, time_index, r, z,
+                                       r_range, z_range, restart_r, restart_r_spectral,
+                                       restart_z, restart_z_spectral,
+                                       interpolation_needed)
+            moments.electron.dens_updated[] = true
+            moments.electron.upar .=
+                reload_electron_moment("electron_parallel_flow", dynamic, time_index, r,
+                                       z, r_range, z_range, restart_r, restart_r_spectral,
+                                       restart_z, restart_z_spectral,
+                                       interpolation_needed)
+            moments.electron.upar_updated[] = true
+            moments.electron.ppar .=
+                reload_electron_moment("electron_parallel_pressure", dynamic, time_index,
+                                       r, z, r_range, z_range, restart_r,
+                                       restart_r_spectral, restart_z, restart_z_spectral,
+                                       interpolation_needed)
+            moments.electron.ppar_updated[] = true
+            moments.electron.qpar .=
+                reload_electron_moment("electron_parallel_heat_flux", dynamic, time_index,
+                                       r, z, r_range, z_range, restart_r,
+                                       restart_r_spectral, restart_z, restart_z_spectral,
+                                       interpolation_needed)
+            moments.electron.qpar_updated[] = true
+            moments.electron.vth .=
+                reload_electron_moment("electron_thermal_speed", dynamic, time_index, r,
+                                       z, r_range, z_range, restart_r, restart_r_spectral,
+                                       restart_z, restart_z_spectral,
+                                       interpolation_needed)
+
+            pdf.electron.norm .=
+                reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
+                                    r_range, z_range, vperp_range, vpa_range, restart_r,
+                                    restart_r_spectral, restart_z, restart_z_spectral,
+                                    restart_vperp, restart_vperp_spectral, restart_vpa,
+                                    restart_vpa_spectral, interpolation_needed,
+                                    restart_evolve_density, restart_evolve_upar,
+                                    restart_evolve_ppar)
         finally
             close(fid)
         end
@@ -1843,6 +917,1517 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
     return code_time, previous_runs_info, time_index
 end
 
+function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io)
+    if parallel_io
+        restart_z, restart_z_spectral, _ =
+            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank)
+        restart_r, restart_r_spectral, _ =
+            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank)
+        restart_vperp, restart_vperp_spectral, _ =
+            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank)
+        restart_vpa, restart_vpa_spectral, _ =
+            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank)
+        restart_vzeta, restart_vzeta_spectral, _ =
+            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank)
+        restart_vr, restart_vr_spectral, _ =
+            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank)
+        restart_vz, restart_vz_spectral, _ =
+            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank)
+    else
+        restart_z, restart_z_spectral, _ = load_coordinate_data(fid, "z")
+        restart_r, restart_r_spectral, _ = load_coordinate_data(fid, "r")
+        restart_vperp, restart_vperp_spectral, _ =
+            load_coordinate_data(fid, "vperp")
+        restart_vpa, restart_vpa_spectral, _ = load_coordinate_data(fid, "vpa")
+        restart_vzeta, restart_vzeta_spectral, _ =
+            load_coordinate_data(fid, "vzeta")
+        restart_vr, restart_vr_spectral, _ = load_coordinate_data(fid, "vr")
+        restart_vz, restart_vz_spectral, _ = load_coordinate_data(fid, "vz")
+
+        if restart_r.nrank != r.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now r.nrank=$(r.nrank), but we are trying to "
+                  * "restart from files ith restart_r.nrank=$(restart_r.nrank).")
+        end
+        if restart_z.nrank != z.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now z.nrank=$(z.nrank), but we are trying to "
+                  * "restart from files ith restart_z.nrank=$(restart_z.nrank).")
+        end
+        if restart_vperp.nrank != vperp.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now vperp.nrank=$(vperp.nrank), but we are trying to "
+                  * "restart from files ith restart_vperp.nrank=$(restart_vperp.nrank).")
+        end
+        if restart_vpa.nrank != vpa.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now vpa.nrank=$(vpa.nrank), but we are trying to "
+                  * "restart from files ith restart_vpa.nrank=$(restart_vpa.nrank).")
+        end
+        if restart_vzeta.nrank != vzeta.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now vzeta.nrank=$(vzeta.nrank), but we are trying to "
+                  * "restart from files ith restart_vzeta.nrank=$(restart_vzeta.nrank).")
+        end
+        if restart_vr.nrank != vr.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now vr.nrank=$(vr.nrank), but we are trying to "
+                  * "restart from files ith restart_vr.nrank=$(restart_vr.nrank).")
+        end
+        if restart_vz.nrank != vz.nrank
+            error("Not using parallel I/O, and distributed MPI layout has "
+                  * "changed: now vz.nrank=$(vz.nrank), but we are trying to "
+                  * "restart from files ith restart_vz.nrank=$(restart_vz.nrank).")
+        end
+    end
+
+    return restart_r, restart_r_spectral, restart_z, restart_z_spectral, restart_vperp,
+           restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
+           restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
+           restart_vz_spectral
+end
+
+function get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp, restart_vpa,
+                           restart_vzeta, restart_vr, restart_vz)
+    if parallel_io
+        function get_range(coord)
+            if coord.irank == coord.nrank - 1
+                return coord.global_io_range
+            else
+                # Need to modify the range to load the end-point that is duplicated on
+                # the next process
+                this_range = coord.global_io_range
+                return this_range.start:(this_range.stop+1)
+            end
+        end
+        r_range = get_range(restart_r)
+        z_range = get_range(restart_z)
+        vperp_range = get_range(restart_vperp)
+        vpa_range = get_range(restart_vpa)
+        vzeta_range = get_range(restart_vzeta)
+        vr_range = get_range(restart_vr)
+        vz_range = get_range(restart_vz)
+    else
+        r_range = (:)
+        z_range = (:)
+        vperp_range = (:)
+        vpa_range = (:)
+        vzeta_range = (:)
+        vr_range = (:)
+        vz_range = (:)
+    end
+    return r_range, z_range, vperp_range, vpa_range, vzeta_range, vr_range, vz_range
+end
+
+function reload_moment(var_name, dynamic, time_index, r, z, r_range, z_range, restart_r,
+                       restart_r_spectral, restart_z, restart_z_spectral,
+                       interpolation_needed)
+    moment = load_slice(dynamic, var_name, z_range, r_range, :, time_index)
+    orig_nz, orig_nr, nspecies = size(moment)
+    if interpolation_needed["r"]
+        new_moment = allocate_float(orig_nz, r.n, nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:orig_nz
+            @views interpolate_to_grid_1d!(new_moment[iz,:,is], r.grid,
+                                           moment[iz,:,is], restart_r,
+                                           restart_r_spectral)
+        end
+        moment = new_moment
+    end
+    if interpolation_needed["z"]
+        new_moment = allocate_float(z.n, r.n, nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n
+            @views interpolate_to_grid_1d!(new_moment[:,ir,is], z.grid,
+                                           moment[:,ir,is], restart_z,
+                                           restart_z_spectral)
+        end
+        moment = new_moment
+    end
+    return moment
+end
+
+function reload_electron_moment(var_name, dynamic, time_index, r, z, r_range, z_range,
+                                restart_r, restart_r_spectral, restart_z,
+                                restart_z_spectral, interpolation_needed)
+    moment = load_slice(dynamic, var_name, z_range, r_range, time_index)
+    orig_nz, orig_nr = size(moment)
+    if interpolation_needed["r"]
+        new_moment = allocate_float(orig_nz, r.n)
+        for iz ∈ 1:orig_nz
+            @views interpolate_to_grid_1d!(new_moment[iz,:], r.grid, moment[iz,:],
+                                           restart_r, restart_r_spectral)
+        end
+        moment = new_moment
+    end
+    if interpolation_needed["z"]
+        new_moment = allocate_float(z.n, r.n)
+        for ir ∈ 1:r.n
+            @views interpolate_to_grid_1d!(new_moment[:,ir], z.grid, moment[:,ir],
+                                           restart_z, restart_z_spectral)
+        end
+        moment = new_moment
+    end
+    return moment
+end
+
+function reload_ion_pdf(dynamic, time_index, moments, r, z, vperp, vpa, r_range, z_range,
+                        vperp_range, vpa_range, restart_r, restart_r_spectral, restart_z,
+                        restart_z_spectral, restart_vperp, restart_vperp_spectral,
+                        restart_vpa, restart_vpa_spectral, interpolation_needed,
+                        restart_evolve_density, restart_evolve_upar, restart_evolve_ppar)
+
+    this_pdf = load_slice(dynamic, "f", vpa_range, vperp_range, z_range,
+        r_range, :, time_index)
+    orig_nvpa, orig_nvperp, orig_nz, orig_nr, nspecies = size(this_pdf)
+    if interpolation_needed["r"]
+        new_pdf = allocate_float(orig_nvpa, orig_nvperp, orig_nz, r.n, nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:orig_nz, ivperp ∈ 1:orig_nvperp,
+            ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,ivperp,iz,:,is], r.grid,
+                       this_pdf[ivpa,ivperp,iz,:,is], restart_r,
+                       restart_r_spectral)
+        end
+        this_pdf = new_pdf
+    end
+    if interpolation_needed["z"]
+        new_pdf = allocate_float(orig_nvpa, orig_nvperp, z.n, r.n, nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n, ivperp ∈ 1:orig_nvperp,
+            ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,ivperp,:,ir,is], z.grid,
+                       this_pdf[ivpa,ivperp,:,ir,is], restart_z,
+                       restart_z_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    # Current moment-kinetic implementation is only 1V, so no need to handle a
+    # normalised vperp coordinate. This will need to change when 2V
+    # moment-kinetics is implemented.
+    if interpolation_needed["vperp"]
+        new_pdf = allocate_float(orig_nvpa, vperp.n, z.n, r.n, nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,:,iz,ir,is], vperp.grid,
+                       this_pdf[ivpa,:,iz,ir,is], restart_vperp,
+                       restart_vperp_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    if (
+        (moments.evolve_density == restart_evolve_density &&
+         moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
+         restart_evolve_ppar)
+        || (!moments.evolve_upar && !restart_evolve_upar &&
+            !moments.evolve_ppar && !restart_evolve_ppar)
+       )
+        # No chages to velocity-space coordinates, so just interpolate from
+        # one grid to the other
+        if interpolation_needed["vpa"]
+            new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+            for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+                @views interpolate_to_grid_1d!(
+                           new_pdf[:,ivperp,iz,ir,is], vpa.grid,
+                           this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                           restart_vpa_spectral)
+            end
+            this_pdf = new_pdf
+        end
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa + upar
+        # => old_wpa = new_wpa - upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth + upar
+        # => old_wpa = (new_wpa - upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid - moments.ion.upar[iz,ir,is]) / moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa
+        # => old_wpa = new_wpa + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth
+        # => old_wpa = (new_wpa + upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid + moments.ion.upar[iz,ir,is]) / moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth + upar
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa + upar
+        # => old_wpa = new_wpa*vth - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] - moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa*vth + upar
+        # => old_wpa = new_wpa - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid - moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa
+        # => old_wpa = new_wpa*vth + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] + moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa + upar
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa*vth
+        # => old_wpa = new_wpa + upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid + moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir,is], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir,is], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    else
+        # This should never happen, as all combinations of evolve_* options
+        # should be handled above.
+        error("Unsupported combination of moment-kinetic options:"
+              * " evolve_density=", moments.evolve_density
+              * " evolve_upar=", moments.evolve_upar
+              * " evolve_ppar=", moments.evolve_ppar
+              * " restart_evolve_density=", restart_evolve_density
+              * " restart_evolve_upar=", restart_evolve_upar
+              * " restart_evolve_ppar=", restart_evolve_ppar)
+    end
+    if moments.evolve_density && !restart_evolve_density
+        # Need to normalise by density
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir,is] ./= moments.ion.dens[iz,ir,is]
+        end
+    elseif !moments.evolve_density && restart_evolve_density
+        # Need to unnormalise by density
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir,is] .*= moments.ion.dens[iz,ir,is]
+        end
+    end
+    if moments.evolve_ppar && !restart_evolve_ppar
+        # Need to normalise by vth
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir,is] .*= moments.ion.vth[iz,ir,is]
+        end
+    elseif !moments.evolve_ppar && restart_evolve_ppar
+        # Need to unnormalise by vth
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir,is] ./= moments.ion.vth[iz,ir,is]
+        end
+    end
+
+    return this_pdf
+end
+
+function reload_ion_boundary_pdf(boundary_distributions_io, var_name, ir, moments, z,
+                                 vperp, vpa, z_range, vperp_range, vpa_range, restart_z,
+                                 restart_z_spectral, restart_vperp,
+                                 restart_vperp_spectral, restart_vpa,
+                                 restart_vpa_spectral, interpolation_needed,
+                                 restart_evolve_density, restart_evolve_upar,
+                                 restart_evolve_ppar)
+    this_pdf = load_slice(boundary_distributions_io, var_name, vpa_range,
+        vperp_range, z_range, :)
+    orig_nvpa, orig_nvperp, orig_nz, nspecies = size(this_pdf)
+    if interpolation_needed["z"]
+        new_pdf = allocate_float(orig_nvpa, orig_nvperp, z.n, nspecies)
+        for is ∈ 1:nspecies, ivperp ∈ 1:orig_nvperp,
+            ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[ivpa,ivperp,:,is], z.grid,
+                                           this_pdf[ivpa,ivperp,:,is], restart_z,
+                                           restart_z_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    # Current moment-kinetic implementation is only 1V, so no need to handle a
+    # normalised vperp coordinate. This will need to change when 2V
+    # moment-kinetics is implemented.
+    if interpolation_needed["vperp"]
+        new_pdf = allocate_float(orig_nvpa, vperp.n, z.n, nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[ivpa,:,iz,is], vperp.grid,
+                                           this_pdf[ivpa,:,iz,is], restart_vperp,
+                                           restart_vperp_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    if (
+        (moments.evolve_density == restart_evolve_density &&
+         moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
+         restart_evolve_ppar)
+        || (!moments.evolve_upar && !restart_evolve_upar &&
+            !moments.evolve_ppar && !restart_evolve_ppar)
+       )
+        # No chages to velocity-space coordinates, so just interpolate from
+        # one grid to the other
+        if interpolation_needed["vpa"]
+            new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+            for is ∈ 1:nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+                @views interpolate_to_grid_1d!(
+                                               new_pdf[:,ivperp,iz,is], vpa.grid,
+                                               this_pdf[:,ivperp,iz,is], restart_vpa,
+                                               restart_vpa_spectral)
+            end
+            this_pdf = new_pdf
+        end
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa + upar
+        # => old_wpa = new_wpa - upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .- moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth + upar
+        # => old_wpa = (new_wpa - upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid - moments.ion.upar[iz,ir,is]) /
+            moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa
+        # => old_wpa = new_wpa + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .+ moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth
+        # => old_wpa = (new_wpa + upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid + moments.ion.upar[iz,ir,is]) /
+            moments.ion.vth
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth + upar
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa + upar
+        # => old_wpa = new_wpa*vth - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] -
+            moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa*vth + upar
+        # => old_wpa = new_wpa - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid -
+            moments.ion.upar[iz,ir,is]/moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa
+        # => old_wpa = new_wpa*vth + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.ion.vth[iz,ir,is] +
+            moments.ion.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa + upar
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa*vth
+        # => old_wpa = new_wpa + upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid +
+            moments.ion.upar[iz,ir,is] / moments.ion.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                                           new_pdf[:,ivperp,iz,is], restart_vpa_vals,
+                                           this_pdf[:,ivperp,iz,is], restart_vpa, restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    else
+        # This should never happen, as all combinations of evolve_* options
+        # should be handled above.
+        error("Unsupported combination of moment-kinetic options:"
+              * " evolve_density=", moments.evolve_density
+              * " evolve_upar=", moments.evolve_upar
+              * " evolve_ppar=", moments.evolve_ppar
+              * " restart_evolve_density=", restart_evolve_density
+              * " restart_evolve_upar=", restart_evolve_upar
+              * " restart_evolve_ppar=", restart_evolve_ppar)
+    end
+    if moments.evolve_density && !restart_evolve_density
+        # Need to normalise by density
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,iz,is] ./= moments.ion.dens[iz,ir,is]
+        end
+    elseif !moments.evolve_density && restart_evolve_density
+        # Need to unnormalise by density
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,iz,is] .*= moments.ion.dens[iz,ir,is]
+        end
+    end
+    if moments.evolve_ppar && !restart_evolve_ppar
+        # Need to normalise by vth
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,iz,is] .*= moments.ion.vth[iz,ir,is]
+        end
+    elseif !moments.evolve_ppar && restart_evolve_ppar
+        # Need to unnormalise by vth
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,iz,is] ./= moments.ion.vth[iz,ir,is]
+        end
+    end
+
+    return this_pdf
+end
+
+function reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa, r_range,
+                             z_range, vperp_range, vpa_range, restart_r,
+                             restart_r_spectral, restart_z, restart_z_spectral,
+                             restart_vperp, restart_vperp_spectral, restart_vpa,
+                             restart_vpa_spectral, interpolation_needed,
+                             restart_evolve_density, restart_evolve_upar,
+                             restart_evolve_ppar)
+
+    # Currently, electrons are always fully moment-kinetic
+    evolve_density = true
+    evolve_upar = true
+    evolve_ppar = true
+
+    this_pdf = load_slice(dynamic, "f_electron", vpa_range, vperp_range, z_range, r_range,
+                          time_index)
+    orig_nvpa, orig_nvperp, orig_nz, orig_nr = size(this_pdf)
+    if interpolation_needed["r"]
+        new_pdf = allocate_float(orig_nvpa, orig_nvperp, orig_nz, r.n)
+        for iz ∈ 1:orig_nz, ivperp ∈ 1:orig_nvperp,
+            ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,ivperp,iz,:], r.grid,
+                       this_pdf[ivpa,ivperp,iz,:], restart_r,
+                       restart_r_spectral)
+        end
+        this_pdf = new_pdf
+    end
+    if interpolation_needed["z"]
+        new_pdf = allocate_float(orig_nvpa, orig_nvperp, z.n, r.n)
+        for ir ∈ 1:r.n, ivperp ∈ 1:orig_nvperp,
+            ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,ivperp,:,ir], z.grid,
+                       this_pdf[ivpa,ivperp,:,ir], restart_z,
+                       restart_z_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    # Current moment-kinetic implementation is only 1V, so no need to handle a
+    # normalised vperp coordinate. This will need to change when 2V
+    # moment-kinetics is implemented.
+    if interpolation_needed["vperp"]
+        new_pdf = allocate_float(orig_nvpa, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivpa ∈ 1:orig_nvpa
+            @views interpolate_to_grid_1d!(
+                       new_pdf[ivpa,:,iz,ir], vperp.grid,
+                       this_pdf[ivpa,:,iz,ir], restart_vperp,
+                       restart_vperp_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    if (
+        (evolve_density == restart_evolve_density &&
+         evolve_upar == restart_evolve_upar && evolve_ppar ==
+         restart_evolve_ppar)
+        || (!evolve_upar && !restart_evolve_upar &&
+            !evolve_ppar && !restart_evolve_ppar)
+       )
+        # No chages to velocity-space coordinates, so just interpolate from
+        # one grid to the other
+        if interpolation_needed["vpa"]
+            new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+            for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+                @views interpolate_to_grid_1d!(
+                           new_pdf[:,ivperp,iz,ir], vpa.grid,
+                           this_pdf[:,ivperp,iz,ir], restart_vpa,
+                           restart_vpa_spectral)
+            end
+            this_pdf = new_pdf
+        end
+    elseif (!evolve_upar && !evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa + upar
+        # => old_wpa = new_wpa - upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .- moments.electron.upar[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!evolve_upar && !evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!evolve_upar && !evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth + upar
+        # => old_wpa = (new_wpa - upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid - moments.electron.upar[iz,ir]) /
+                moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && !evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa
+        # => old_wpa = new_wpa + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .+ moments.electron.upar[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && !evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth
+        # => old_wpa = (new_wpa + upar)/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. (vpa.grid + moments.electron.upar[iz,ir]) /
+                moments.electron.vth
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && !evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth + upar
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid ./ moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!evolve_upar && evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!evolve_upar && evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa + upar
+        # => old_wpa = new_wpa*vth - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.electron.vth[iz,ir] -
+                                  moments.electron.upar[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!evolve_upar && evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa*vth + upar
+        # => old_wpa = new_wpa - upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid -
+            moments.electron.upar[iz,ir]/moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa
+        # => old_wpa = new_wpa*vth + upar
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = @. vpa.grid * moments.electron.vth[iz,ir] +
+                                  moments.electron.upar[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa + upar
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals = vpa.grid .* moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (evolve_upar && evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa*vth
+        # => old_wpa = new_wpa + upar/vth
+        new_pdf = allocate_float(vpa.n, vperp.n, z.n, r.n)
+        for ir ∈ 1:r.n, iz ∈ 1:z.n, ivperp ∈ 1:vperp.n
+            restart_vpa_vals =
+            @. vpa.grid +
+            moments.electron.upar[iz,ir] / moments.electron.vth[iz,ir]
+            @views interpolate_to_grid_1d!(
+                       new_pdf[:,ivperp,iz,ir], restart_vpa_vals,
+                       this_pdf[:,ivperp,iz,ir], restart_vpa,
+                       restart_vpa_spectral)
+        end
+        this_pdf = new_pdf
+    else
+        # This should never happen, as all combinations of evolve_* options
+        # should be handled above.
+        error("Unsupported combination of moment-kinetic options:"
+              * " evolve_density=", evolve_density
+              * " evolve_upar=", evolve_upar
+              * " evolve_ppar=", evolve_ppar
+              * " restart_evolve_density=", restart_evolve_density
+              * " restart_evolve_upar=", restart_evolve_upar
+              * " restart_evolve_ppar=", restart_evolve_ppar)
+    end
+    if evolve_density && !restart_evolve_density
+        # Need to normalise by density
+        for ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir] ./= moments.electron.dens[iz,ir]
+        end
+    elseif !evolve_density && restart_evolve_density
+        # Need to unnormalise by density
+        for ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir] .*= moments.electron.dens[iz,ir]
+        end
+    end
+    if evolve_ppar && !restart_evolve_ppar
+        # Need to normalise by vth
+        for ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir] .*= moments.electron.vth[iz,ir]
+        end
+    elseif !evolve_ppar && restart_evolve_ppar
+        # Need to unnormalise by vth
+        for ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,iz,ir] ./= moments.electron.vth[iz,ir]
+        end
+    end
+
+    return this_pdf
+end
+
+function reload_neutral_pdf(dynamic, time_index, moments, r, z, vzeta, vr, vz, r_range,
+                            z_range, vzeta_range, vr_range, vz_range, restart_r,
+                            restart_r_spectral, restart_z, restart_z_spectral,
+                            restart_vzeta, restart_vzeta_spectral, restart_vr,
+                            restart_vr_spectral, restart_vz, restart_vz_spectral,
+                            interpolation_needed, restart_evolve_density,
+                            restart_evolve_upar, restart_evolve_ppar)
+    this_pdf = load_slice(dynamic, "f_neutral", vz_range, vr_range,
+                          vzeta_range, z_range, r_range, :, time_index)
+    orig_nvz, orig_nvr, orig_nvzeta, orig_nz, orig_nr, nspecies =
+        size(this_pdf)
+    if interpolation_needed["r"]
+        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, orig_nz,
+                                 r.n, nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:orig_nz, ivzeta ∈ 1:orig_nvzeta,
+                ivr ∈ 1:orig_nvr, ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,ivr,ivzeta,iz,:,is], r.grid,
+                this_pdf[ivz,ivr,ivzeta,iz,:,is], restart_r,
+                restart_r_spectral)
+        end
+        this_pdf = new_pdf
+    end
+    if interpolation_needed["z"]
+        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, z.n,
+                                 r.n, nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n, ivzeta ∈ 1:orig_nvzeta,
+                ivr ∈ 1:orig_nvr, ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,ivr,ivzeta,:,ir,is], z.grid,
+                this_pdf[ivz,ivr,ivzeta,:,ir,is], restart_z,
+                restart_z_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    # Current moment-kinetic implementation is only 1V, so no need
+    # to handle normalised vzeta or vr coordinates. This will need
+    # to change when/if 3V moment-kinetics is implemented.
+    if interpolation_needed["vzeta"]
+        new_pdf = allocate_float(orig_nvz, orig_nvr, vzeta.n, z.n, r.n,
+                                 nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:orig_nvr,
+                ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,ivr,:,iz,ir,is], vzeta.grid,
+                this_pdf[ivz,ivr,:,iz,ir,is], restart_vzeta,
+                restart_vzeta_spectral)
+        end
+        this_pdf = new_pdf
+    end
+    if interpolation_needed["vr"]
+        new_pdf = allocate_float(orig_nvz, vr.n, vzeta.n, z.n, r.n,
+                                 nspecies)
+        for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n,
+                ivzeta ∈ 1:vzeta.n, ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,:,ivzeta,iz,ir,is], vr.grid,
+                this_pdf[ivz,:,ivzeta,iz,ir,is], restart_vr,
+                restart_vr_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    if (
+        (moments.evolve_density == restart_evolve_density &&
+         moments.evolve_upar == restart_evolve_upar &&
+         moments.evolve_ppar == restart_evolve_ppar)
+        || (!moments.evolve_upar && !restart_evolve_upar &&
+            !moments.evolve_ppar && !restart_evolve_ppar)
+       )
+        # No chages to velocity-space coordinates, so just interpolate from
+        # one grid to the other
+        if interpolation_needed["vz"]
+            new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+            for is ∈ 1:nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                    ivzeta ∈ 1:vzeta.n
+                @views interpolate_to_grid_1d!(
+                    new_pdf[:,ivr,ivzeta,iz,ir,is], vz.grid,
+                    this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                    restart_vz_spectral)
+            end
+            this_pdf = new_pdf
+        end
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa + upar
+        # => old_wpa = new_wpa - upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n,
+                ivr ∈ 1:vr.n
+            restart_vz_vals = vz.grid .- moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth + upar
+        # => old_wpa = (new_wpa - upar)/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. (vz.grid - moments.neutral.uz[iz,ir,is]) /
+                   moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa
+        # => old_wpa = new_wpa + upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .+ moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz, restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth
+        # => old_wpa = (new_wpa + upar)/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. (vz.grid + moments.neutral.uz[iz,ir,is]) /
+                    moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth + upar
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa + upar
+        # => old_wpa = new_wpa*vth - upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid * moments.neutral.vth[iz,ir,is] -
+                   moments.neutral.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa*vth + upar
+        # => old_wpa = new_wpa - upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid -
+                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa
+        # => old_wpa = new_wpa*vth + upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid * moments.neutral.vth[iz,ir,is] +
+                   moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa + upar
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa*vth
+        # => old_wpa = new_wpa + upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, r.n, nspecies)
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid +
+                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,ir,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    else
+        # This should never happen, as all combinations of evolve_* options
+        # should be handled above.
+        error("Unsupported combination of moment-kinetic options:"
+              * " evolve_density=", moments.evolve_density
+              * " evolve_upar=", moments.evolve_upar
+              * " evolve_ppar=", moments.evolve_ppar
+              * " restart_evolve_density=", restart_evolve_density
+              * " restart_evolve_upar=", restart_evolve_upar
+              * " restart_evolve_ppar=", restart_evolve_ppar)
+    end
+    if moments.evolve_density && !restart_evolve_density
+        # Need to normalise by density
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,ir,is] ./= moments.neutral.dens[iz,ir,is]
+        end
+    elseif !moments.evolve_density && restart_evolve_density
+        # Need to unnormalise by density
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,ir,is] .*= moments.neutral.dens[iz,ir,is]
+        end
+    end
+    if moments.evolve_ppar && !restart_evolve_ppar
+        # Need to normalise by vth
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,ir,is] .*= moments.neutral.vth[iz,ir,is]
+        end
+    elseif !moments.evolve_ppar && restart_evolve_ppar
+        # Need to unnormalise by vth
+        for is ∈ nspecies, ir ∈ 1:r.n, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,ir,is] ./= moments.neutral.vth[iz,ir,is]
+        end
+    end
+
+    return this_pdf
+end
+
+function reload_neutral_boundary_pdf(boundary_distributions_io, var_name, ir, moments, z,
+                                     vzeta, vr, vz, z_range, vzeta_range, vr_range,
+                                     vz_range, restart_z, restart_z_spectral,
+                                     restart_vzeta, restart_vzeta_spectral, restart_vr,
+                                     restart_vr_spectral, restart_vz, restart_vz_spectral,
+                                     interpolation_needed, restart_evolve_density,
+                                     restart_evolve_upar, restart_evolve_ppar)
+    this_pdf = load_slice(boundary_distributions_io, var_name, vz_range,
+                          vr_range, vzeta_range, z_range, :)
+    orig_nvz, orig_nvr, orig_nvzeta, orig_nz, nspecies = size(this_pdf)
+    if interpolation_needed["z"]
+        new_pdf = allocate_float(orig_nvz, orig_nvr, orig_nvzeta, z.n,
+                                 nspecies)
+        for is ∈ 1:nspecies, ivzeta ∈ 1:orig_nvzeta, ivr ∈ 1:orig_nvr,
+                ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,ivr,ivzeta,:,is], z.grid,
+                this_pdf[ivz,ivr,ivzeta,:,is], restart_z,
+                restart_z_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    # Current moment-kinetic implementation is only 1V, so no need
+    # to handle normalised vzeta or vr coordinates. This will need
+    # to change when/if 3V moment-kinetics is implemented.
+    if interpolation_needed["vzeta"]
+        new_pdf = allocate_float(orig_nvz, orig_nvr, vzeta.n, z.n,
+                                 nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivr ∈ 1:orig_nvr,
+                ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,ivr,:,iz,is], vzeta.grid,
+                this_pdf[ivz,ivr,:,iz,is], restart_vzeta,
+                restart_vzeta_spectral)
+        end
+        this_pdf = new_pdf
+    end
+    if interpolation_needed["vr"]
+        new_pdf = allocate_float(orig_nvz, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ 1:nspecies, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n,
+                ivz ∈ 1:orig_nvz
+            @views interpolate_to_grid_1d!(
+                new_pdf[ivz,:,ivzeta,iz,is], vr.grid,
+                this_pdf[ivz,:,ivzeta,iz,is], restart_vr,
+                restart_vr_spectral)
+        end
+        this_pdf = new_pdf
+    end
+
+    if (
+        (moments.evolve_density == restart_evolve_density &&
+         moments.evolve_upar == restart_evolve_upar && moments.evolve_ppar ==
+         restart_evolve_ppar)
+        || (!moments.evolve_upar && !restart_evolve_upar &&
+            !moments.evolve_ppar && !restart_evolve_ppar)
+       )
+        # No chages to velocity-space coordinates, so just interpolate from
+        # one grid to the other
+        if interpolation_needed["vz"]
+            new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+            for is ∈ 1:nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n,
+                    ivzeta ∈ 1:vzeta.n
+                @views interpolate_to_grid_1d!(
+                    new_pdf[:,ivr,ivzeta,iz,is], vz.grid,
+                    this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                    restart_vz_spectral)
+            end
+            this_pdf = new_pdf
+        end
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa + upar
+        # => old_wpa = new_wpa - upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivzeta ∈ 1:vzeta.n, ivr ∈ 1:vr.n
+            restart_vz_vals = vz.grid .- moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa = old_wpa*vth + upar
+        # => old_wpa = (new_wpa - upar)/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. (vz.grid - moments.neutral.uz[iz,ir,is]) /
+                   moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa
+        # => old_wpa = new_wpa + upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .+
+                              moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth
+        # => old_wpa = (new_wpa + upar)/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. (vz.grid + moments.neutral.uz[iz,ir,is]) /
+                   moments.neutral.vth
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && !moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa + upar = old_wpa*vth + upar
+        # => old_wpa = new_wpa/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid ./ moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .*
+                              moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa + upar
+        # => old_wpa = new_wpa*vth - upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid * moments.neutral.vth[iz,ir,is] -
+                   moments.neutral.upar[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (!moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth = old_wpa*vth + upar
+        # => old_wpa = new_wpa - upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid -
+                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa
+        # => old_wpa = new_wpa*vth + upar
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid * moments.neutral.vth[iz,ir,is] +
+                   moments.neutral.uz[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            restart_evolve_upar && !restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa + upar
+        # => old_wpa = new_wpa*vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals = vz.grid .* moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+        this_pdf = new_pdf
+    elseif (moments.evolve_upar && moments.evolve_ppar &&
+            !restart_evolve_upar && restart_evolve_ppar)
+        # vpa = new_wpa*vth + upar = old_wpa*vth
+        # => old_wpa = new_wpa + upar/vth
+        new_pdf = allocate_float(vz.n, vr.n, vzeta.n, z.n, nspecies)
+        for is ∈ nspecies, iz ∈ 1:z.n, ivr ∈ 1:vr.n, ivzeta ∈ 1:vzeta.n
+            restart_vz_vals =
+                @. vz.grid +
+                   moments.neutral.uz[iz,ir,is]/moments.neutral.vth[iz,ir,is]
+            @views interpolate_to_grid_1d!(
+                new_pdf[:,ivr,ivzeta,iz,is], restart_vz_vals,
+                this_pdf[:,ivr,ivzeta,iz,is], restart_vz,
+                restart_vz_spectral)
+        end
+    else
+        # This should never happen, as all combinations of evolve_* options
+        # should be handled above.
+        error("Unsupported combination of moment-kinetic options:"
+              * " evolve_density=", moments.evolve_density
+              * " evolve_upar=", moments.evolve_upar
+              * " evolve_ppar=", moments.evolve_ppar
+              * " restart_evolve_density=", restart_evolve_density
+              * " restart_evolve_upar=", restart_evolve_upar
+              * " restart_evolve_ppar=", restart_evolve_ppar)
+    end
+    if moments.evolve_density && !restart_evolve_density
+        # Need to normalise by density
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,is] ./= moments.neutral.dens[iz,ir,is]
+        end
+    elseif !moments.evolve_density && restart_evolve_density
+        # Need to unnormalise by density
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,is] .*= moments.neutral.dens[iz,ir,is]
+        end
+    end
+    if moments.evolve_ppar && !restart_evolve_ppar
+        # Need to normalise by vth
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,is] .*= moments.neutral.vth[iz,ir,is]
+        end
+    elseif !moments.evolve_ppar && restart_evolve_ppar
+        # Need to unnormalise by vth
+        for is ∈ nspecies, iz ∈ 1:z.n
+            this_pdf[:,:,:,iz,is] ./= moments.neutral.vth[iz,ir,is]
+        end
+    end
+
+    return this_pdf
+end
+
 """
 Read a slice of an ion distribution function
 

From ec4d029df1512a994dc10c913d52edfaac7af372 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 20 Feb 2024 15:42:45 +0000
Subject: [PATCH 112/428] Broadcast code_time to all processes in a block

Previously the 'time' was only correct (in a restarted simulation) on
the root process of each block. The output written to file would have
been correct, but the value in memory on the non-root processes would
have been incorrect.
---
 moment_kinetics/src/moment_kinetics.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 0a7f1fd58..a5ce28dad 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -421,6 +421,11 @@ function setup_moment_kinetics(input_dict::AbstractDict;
 
         _block_synchronize()
     end
+
+    # Broadcast code_time from the root process of each shared-memory block (on which it
+    # might have been loaded from a restart file).
+    code_time = MPI.Bcast(code_time, 0, comm_block[])
+
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 

From 35d9555896b4b66963b94d9d2fdf1d24c981ecb7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 20 Feb 2024 14:31:09 +0000
Subject: [PATCH 113/428] Move restart file setup functionality to utility
 functions

...will allow the functions to be reused to restart the initialization
of kinetic electrons.
---
 moment_kinetics/src/moment_kinetics.jl | 124 ++-------------------
 moment_kinetics/src/utils.jl           | 144 +++++++++++++++++++++++++
 2 files changed, 151 insertions(+), 117 deletions(-)

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index a5ce28dad..a4344901e 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -54,7 +54,6 @@ include("neutral_r_advection.jl")
 include("neutral_z_advection.jl")
 include("neutral_vz_advection.jl")
 include("electron_kinetic_equation.jl")
-include("initial_conditions.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("krook_collisions.jl")
@@ -64,15 +63,15 @@ include("force_balance.jl")
 include("source_terms.jl")
 include("numerical_dissipation.jl")
 include("moment_kinetics_input.jl")
+include("utils.jl")
+include("initial_conditions.jl")
 include("load_data.jl")
 include("parameter_scans.jl")
 include("analysis.jl")
-include("utils.jl")
 include("time_advance.jl")
 
 using TimerOutputs
 using Dates
-using Glob
 using Primes
 
 using .file_io: setup_file_io, finish_file_io
@@ -93,7 +92,8 @@ using .looping: debug_setup_loop_ranges_split_one_combination!
 using .moment_kinetics_input: mk_input, read_input_file
 using .time_advance: setup_time_advance!, time_advance!
 using .type_definitions: mk_int
-using .utils: to_minutes
+using .utils: to_minutes, get_default_restart_filename,
+              get_prefix_iblock_and_move_existing_file
 using .em_fields: setup_em_fields
 using .time_advance: setup_dummy_and_buffer_arrays
 using .time_advance: allocate_advection_structs
@@ -195,75 +195,6 @@ function run_moment_kinetics()
                         restart_time_index=restart_time_index)
 end
 
-"""
-Append a number to the filename, to get a new, non-existing filename to backup the file
-to.
-"""
-function get_backup_filename(filename)
-    if !isfile(filename)
-        error("Requested to restart from $filename, but this file does not exist")
-    end
-    counter = 1
-    temp, extension = splitext(filename)
-    extension = extension[2:end]
-    temp, iblock_or_type = splitext(temp)
-    iblock_or_type = iblock_or_type[2:end]
-    iblock = nothing
-    basename = nothing
-    type = nothing
-    if iblock_or_type == "dfns"
-        iblock = nothing
-        type = iblock_or_type
-        basename = temp
-        parallel_io = true
-    else
-        # Filename had an iblock, so we are not using parallel I/O, but actually want to
-        # use the iblock for this block, not necessarily for the exact file that was
-        # passed.
-        iblock = iblock_index[]
-        basename, type = splitext(temp)
-        type = type[2:end]
-        parallel_io = false
-    end
-    if type != "dfns"
-        error("Must pass the '.dfns.h5' output file for restarting. Got $filename.")
-    end
-    backup_dfns_filename = ""
-    if parallel_io
-        # Using parallel I/O
-        while true
-            backup_dfns_filename = "$(basename)_$(counter).$(type).$(extension)"
-            if !isfile(backup_dfns_filename)
-                break
-            end
-            counter += 1
-        end
-        # Create dfns_filename here even though it is the filename passed in, as
-        # parallel_io=false branch needs to get the right `iblock` for this block.
-        dfns_filename = "$(basename).dfns.$(extension)"
-        moments_filename = "$(basename).moments.$(extension)"
-        backup_moments_filename = "$(basename)_$(counter).moments.$(extension)"
-    else
-        while true
-            backup_dfns_filename = "$(basename)_$(counter).$(type).$(iblock).$(extension)"
-            if !isfile(backup_dfns_filename)
-                break
-            end
-            counter += 1
-        end
-        # Create dfns_filename here even though it is almost the filename passed in, in
-        # order to get the right `iblock` for this block.
-        dfns_filename = "$(basename).dfns.$(iblock).$(extension)"
-        moments_filename = "$(basename).moments.$(iblock).$(extension)"
-        backup_moments_filename = "$(basename)_$(counter).moments.$(iblock).$(extension)"
-    end
-    backup_dfns_filename == "" && error("Failed to find a name for backup file.")
-    backup_prefix_iblock = ("$(basename)_$(counter)", iblock)
-    original_prefix_iblock = (basename, iblock)
-    return dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
-           backup_moments_filename, backup_prefix_iblock, original_prefix_iblock
-end
-
 """
 Perform all the initialization steps for a run.
 
@@ -354,55 +285,14 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     else
         restarting = true
 
-        run_name = input_dict["run_name"]
-        base_directory = get(input_dict, "base_directory", "runs")
-        output_dir = joinpath(base_directory, run_name)
         if restart === true
-            run_name = input_dict["run_name"]
-            io_settings = get(input_dict, "output", Dict{String,Any}())
-            binary_format = get(io_settings, "binary_format", hdf5)
-            if binary_format === hdf5
-                ext = "h5"
-            elseif binary_format === netcdf
-                ext = "cdf"
-            else
-                error("Unrecognized binary_format '$binary_format'")
-            end
-            restart_filename_pattern = joinpath(output_dir, run_name * ".dfns*." * ext)
-            restart_filename_glob = glob(restart_filename_pattern)
-            if length(restart_filename_glob) == 0
-                error("No output file to restart from found matching the pattern "
-                      * "$restart_filename_pattern")
-            end
-            restart_filename = restart_filename_glob[1]
+            restart_filename = get_default_restart_filename(io_input, "dfns")
         else
             restart_filename = restart
         end
 
-        # Move the output file being restarted from to make sure it doesn't get
-        # overwritten.
-        dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
-        backup_moments_filename, backup_prefix_iblock, original_prefix_iblock =
-            get_backup_filename(restart_filename)
-
-        # Ensure every process got the filenames and checked files exist before moving
-        # files
-        MPI.Barrier(comm_world)
-
-        if abspath(output_dir) == abspath(dirname(dfns_filename))
-            # Only move the file if it is in our current run directory. Otherwise we are
-            # restarting from another run, and will not be overwriting the file.
-            if (parallel_io && global_rank[] == 0) || (!parallel_io && block_rank[] == 0)
-                mv(dfns_filename, backup_dfns_filename)
-                mv(moments_filename, backup_moments_filename)
-            end
-        else
-            # Reload from dfns_filename without moving the file
-            backup_prefix_iblock = original_prefix_iblock
-        end
-
-        # Ensure files have been moved before any process tries to read from them
-        MPI.Barrier(comm_world)
+        backup_prefix_iblock = get_prefix_iblock_and_move_existing_file(restart_filename,
+                                                                        io_input.output_dir)
 
         # Reload pdf and moments from an existing output file
         code_time, previous_runs_info, restart_time_index =
diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl
index c0e5ca190..7f3ad71f3 100644
--- a/moment_kinetics/src/utils.jl
+++ b/moment_kinetics/src/utils.jl
@@ -6,11 +6,15 @@ module utils
 export get_unnormalized_parameters, print_unnormalized_parameters, to_seconds, to_minutes,
        to_hours
 
+using ..communication
 using ..constants
+using ..input_structs
 using ..moment_kinetics_input: mk_input
 using ..reference_parameters
 
 using Dates
+using Glob
+using MPI
 using OrderedCollections
 using TOML
 using Unitful
@@ -122,4 +126,144 @@ Convert a time period `x` to seconds
 """
 to_hours(x::T) where {T<:TimePeriod} = x/convert(T, Hour(1))
 
+# Utility functions used for restarting
+
+"""
+Append a number to the filename, to get a new, non-existing filename to backup the file
+to.
+"""
+function get_backup_filename(filename)
+    if !isfile(filename)
+        error("Requested to restart from $filename, but this file does not exist")
+    end
+    counter = 1
+    temp, extension = splitext(filename)
+    extension = extension[2:end]
+    temp, iblock_or_type = splitext(temp)
+    iblock_or_type = iblock_or_type[2:end]
+    iblock = nothing
+    basename = nothing
+    type = nothing
+    if iblock_or_type == "dfns"
+        iblock = nothing
+        type = iblock_or_type
+        basename = temp
+        parallel_io = true
+    else
+        # Filename had an iblock, so we are not using parallel I/O, but actually want to
+        # use the iblock for this block, not necessarily for the exact file that was
+        # passed.
+        iblock = iblock_index[]
+        basename, type = splitext(temp)
+        type = type[2:end]
+        parallel_io = false
+    end
+    if type != "dfns"
+        error("Must pass the '.dfns.h5' output file for restarting. Got $filename.")
+    end
+    backup_dfns_filename = ""
+    if parallel_io
+        # Using parallel I/O
+        while true
+            backup_dfns_filename = "$(basename)_$(counter).$(type).$(extension)"
+            if !isfile(backup_dfns_filename)
+                break
+            end
+            counter += 1
+        end
+        # Create dfns_filename here even though it is the filename passed in, as
+        # parallel_io=false branch needs to get the right `iblock` for this block.
+        dfns_filename = "$(basename).dfns.$(extension)"
+        moments_filename = "$(basename).moments.$(extension)"
+        backup_moments_filename = "$(basename)_$(counter).moments.$(extension)"
+    else
+        while true
+            backup_dfns_filename = "$(basename)_$(counter).$(type).$(iblock).$(extension)"
+            if !isfile(backup_dfns_filename)
+                break
+            end
+            counter += 1
+        end
+        # Create dfns_filename here even though it is almost the filename passed in, in
+        # order to get the right `iblock` for this block.
+        dfns_filename = "$(basename).dfns.$(iblock).$(extension)"
+        moments_filename = "$(basename).moments.$(iblock).$(extension)"
+        backup_moments_filename = "$(basename)_$(counter).moments.$(iblock).$(extension)"
+    end
+    backup_dfns_filename == "" && error("Failed to find a name for backup file.")
+    backup_prefix_iblock = ("$(basename)_$(counter)", iblock)
+    original_prefix_iblock = (basename, iblock)
+    return dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
+           backup_moments_filename, backup_prefix_iblock, original_prefix_iblock
+end
+
+"""
+    get_default_restart_filename(io_input, prefix; error_if_no_file_found=true)
+
+Get the default name for the file to restart from, using the input from `io_input`.
+
+`prefix` gives the type of file to open, e.g. "moments", "dfns", or "initial_electron".
+
+If no matching file is found, raise an error unless `error_if_no_file_found=false` is
+passed, in which case no error is raised and instead the function returns `nothing`.
+"""
+function get_default_restart_filename(io_input, prefix; error_if_no_file_found=true)
+    binary_format = io_input.binary_format
+    if binary_format === hdf5
+        ext = "h5"
+    elseif binary_format === netcdf
+        ext = "cdf"
+    else
+        error("Unrecognized binary_format '$binary_format'")
+    end
+    restart_filename_pattern = joinpath(io_input.output_dir, io_input.run_name * ".$prefix*." * ext)
+    restart_filename_glob = glob(restart_filename_pattern)
+    if length(restart_filename_glob) == 0
+        if error_if_no_file_found
+            error("No '$prefix' output file to restart from found matching the pattern "
+                  * "$restart_filename_pattern")
+        end
+        restart_filename = nothing
+    else
+        restart_filename = restart_filename_glob[1]
+    end
+    return restart_filename
+end
+
+"""
+    get_prefix_iblock_and_move_existing_file(restart_filename, output_dir)
+
+Move `restart_filename` to a backup location (if it is in `output_dir`), returning a
+prefix and block-index (which might be `nothing`) which can be used to open the file for
+reloading variables.
+"""
+function get_prefix_iblock_and_move_existing_file(restart_filename, output_dir)
+    # Move the output file being restarted from to make sure it doesn't get
+    # overwritten.
+    dfns_filename, backup_dfns_filename, parallel_io, moments_filename,
+    backup_moments_filename, backup_prefix_iblock, original_prefix_iblock =
+        get_backup_filename(restart_filename)
+
+    # Ensure every process got the filenames and checked files exist before moving
+    # files
+    MPI.Barrier(comm_world)
+
+    if abspath(output_dir) == abspath(dirname(dfns_filename))
+        # Only move the file if it is in our current run directory. Otherwise we are
+        # restarting from another run, and will not be overwriting the file.
+        if (parallel_io && global_rank[] == 0) || (!parallel_io && block_rank[] == 0)
+            mv(dfns_filename, backup_dfns_filename)
+            mv(moments_filename, backup_moments_filename)
+        end
+    else
+        # Reload from dfns_filename without moving the file
+        backup_prefix_iblock = original_prefix_iblock
+    end
+
+    # Ensure files have been moved before any process tries to read from them
+    MPI.Barrier(comm_world)
+
+    return backup_prefix_iblock
+end
+
 end #utils

From 91cdf43c2b6678652e953e807a0008b6f0904d20 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 20 Feb 2024 15:03:48 +0000
Subject: [PATCH 114/428] Restart kinetic electrons from existing output, if
 available

---
 .../src/electron_kinetic_equation.jl          | 14 ++++---
 moment_kinetics/src/file_io.jl                |  6 +--
 moment_kinetics/src/initial_conditions.jl     | 37 +++++++++++++++----
 moment_kinetics/src/load_data.jl              |  2 +-
 moment_kinetics/src/moment_kinetics.jl        |  2 +-
 moment_kinetics/src/utils.jl                  | 24 ++++++++----
 6 files changed, 59 insertions(+), 26 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 326ba800c..aa3515b1f 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -50,7 +50,8 @@ OUTPUT:
 function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
         phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
         vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
-        num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing)
+        num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing,
+        initial_time=0.0)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -61,7 +62,8 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_u
         return update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
             moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
             r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
-            num_diss_params, max_electron_pdf_iterations; io_initial_electron=io_initial_electron)
+            num_diss_params, max_electron_pdf_iterations;
+            io_initial_electron=io_initial_electron, initial_time=initial_time)
     elseif solution_method == "shooting_method"
         return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
             qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
@@ -99,13 +101,15 @@ The electron kinetic equation is:
     scratch_dummy = dummy arrays to be used for temporary storage
     dt = time step size
     max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
+    io_initial_electron = info struct for binary file I/O
+    initial_time = initial value for the (pseudo-)time
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
 function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
     moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition,
     r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
-    num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing)
+    num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0)
 
     begin_r_z_region()
 
@@ -164,7 +168,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #n_ppar_subcycles = 1000
     #n_ppar_subcycles = 200
     n_ppar_subcycles = 1
-    time = 0.0
+    time = initial_time
 
     #z_speedup_fac = 20.0
     #z_speedup_fac = 5.0
@@ -248,7 +252,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     end
     # evolve (artificially) in time until the residual is less than the tolerance
     try
-    while !electron_pdf_converged && (iteration < max_electron_pdf_iterations)
+    while !electron_pdf_converged && (iteration <= max_electron_pdf_iterations)
         begin_r_z_region()
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index ec3aaad66..e932f8a37 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -260,7 +260,7 @@ open output file to save the initial electron pressure and distribution function
 function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition,
                                    collisions, evolve_density, evolve_upar, evolve_ppar,
                                    external_source_settings, input_dict,
-                                   previous_runs_info)
+                                   restart_time_index, previous_runs_info)
     begin_serial_region()
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -274,10 +274,6 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
         parallel_io = io_input.parallel_io
         io_comm = comm_inter_block[]
 
-        # dummy value for restart_time_index, which is not used for the initial electron
-        # state
-        restart_time_index = -1
-
         electrons_prefix = string(out_prefix, ".initial_electron")
         if !parallel_io
             electrons_prefix *= ".$(iblock_index[])"
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 16354aeae..ec4eca8f3 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -26,6 +26,7 @@ using ..looping
 using ..em_fields: update_phi!
 using ..file_io: setup_initial_electron_io, write_initial_electron_state,
                  finish_initial_electron_io
+using ..load_data: reload_electron_data!
 using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, pdf_struct, moments_struct,
                                  boundary_distributions_struct
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
@@ -43,6 +44,7 @@ using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..electron_kinetic_equation: update_electron_pdf!, enforce_boundary_condition_on_electron_pdf!
 using ..input_structs: boltzmann_electron_response_with_simple_sheath, kinetic_electrons
 using ..derivatives: derivative_z!
+using ..utils: get_default_restart_filename, get_prefix_iblock_and_move_existing_file
 
 using ..manufactured_solns: manufactured_solutions
 
@@ -344,7 +346,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                              vzeta, vr, vz, z_spectral, vpa_spectral,
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
-                             collisions, composition, external_source_settings,
+                             collisions, composition, geometry, external_source_settings,
                              num_diss_params, t_input.dt, io_input, input_dict)
 
     return nothing
@@ -470,19 +472,38 @@ end
 
 function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
                                   vz, z_spectral, vpa_spectral, z_advect, vpa_advect,
-                                  scratch_dummy, collisions, composition,
+                                  scratch_dummy, collisions, composition, geometry,
                                   external_source_settings, num_diss_params, dt, io_input,
                                   input_dict)
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
         begin_serial_region()
+        restart_filename = get_default_restart_filename(io_input, "initial_electron";
+                                                        error_if_no_file_found=false)
+        if restart_filename === nothing
+            # No file to restart from
+            previous_runs_info = nothing
+            code_time = 0.0
+            restart_time_index = -1
+        else
+            # Previously-created electron distribution function exists, so use it as
+            # the initial guess.
+            backup_prefix_iblock =
+                get_prefix_iblock_and_move_existing_file(restart_filename,
+                                                         io_input.output_dir)
+
+            # Reload pdf and moments from an existing output file
+            code_time, previous_runs_info, restart_time_index =
+                reload_electron_data!(pdf, moments, backup_prefix_iblock, -1,
+                                      geometry, r, z, vpa, vperp, vzeta, vr, vz)
+
+            # Broadcast code_time from the root process of each shared-memory block (on which it
+            # might have been loaded from a restart file).
+            code_time = MPI.Bcast(code_time, 0, comm_block[])
+        end
         io_initial_electron = nothing
         @serial_region begin
-            if false && isfile(initial_electron_output_filename)
-            else
-                previous_runs_info = nothing
-            end
             # Setup I/O for initial electron state
             io_initial_electron = setup_initial_electron_io(io_input, vz, vr, vzeta, vpa,
                                                             vperp, z, r, composition,
@@ -492,6 +513,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                                             moments.evolve_ppar,
                                                             external_source_settings,
                                                             input_dict,
+                                                            restart_time_index,
                                                             previous_runs_info)
 
             # update the electron pdf in the fvec struct
@@ -526,7 +548,8 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                         z_spectral, vpa_spectral, z_advect, vpa_advect,
                                         scratch_dummy, dt, collisions, composition,
                                         num_diss_params, max_electron_pdf_iterations;
-                                        io_initial_electron=io_initial_electron)
+                                        io_initial_electron=io_initial_electron,
+                                        initial_time=code_time)
 
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index b80e60cf2..d6b3f4507 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -822,7 +822,7 @@ function reload_electron_data!(pdf, moments, restart_prefix_iblock, time_index,
     previous_runs_info = nothing
     begin_serial_region()
     @serial_region begin
-        fid = open_readonly_output_file(restart_prefix_iblock[1], "dfns";
+        fid = open_readonly_output_file(restart_prefix_iblock[1], "initial_electron";
                                         iblock=restart_prefix_iblock[2])
         try # finally to make sure to close file0
             overview = get_group(fid, "overview")
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index a4344901e..248b0038f 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -64,8 +64,8 @@ include("source_terms.jl")
 include("numerical_dissipation.jl")
 include("moment_kinetics_input.jl")
 include("utils.jl")
-include("initial_conditions.jl")
 include("load_data.jl")
+include("initial_conditions.jl")
 include("parameter_scans.jl")
 include("analysis.jl")
 include("time_advance.jl")
diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl
index 7f3ad71f3..f399b8bbd 100644
--- a/moment_kinetics/src/utils.jl
+++ b/moment_kinetics/src/utils.jl
@@ -144,7 +144,7 @@ function get_backup_filename(filename)
     iblock = nothing
     basename = nothing
     type = nothing
-    if iblock_or_type == "dfns"
+    if iblock_or_type ∈ ("dfns", "initial_electron")
         iblock = nothing
         type = iblock_or_type
         basename = temp
@@ -158,7 +158,7 @@ function get_backup_filename(filename)
         type = type[2:end]
         parallel_io = false
     end
-    if type != "dfns"
+    if type ∉ ("dfns", "initial_electron")
         error("Must pass the '.dfns.h5' output file for restarting. Got $filename.")
     end
     backup_dfns_filename = ""
@@ -173,8 +173,12 @@ function get_backup_filename(filename)
         end
         # Create dfns_filename here even though it is the filename passed in, as
         # parallel_io=false branch needs to get the right `iblock` for this block.
-        dfns_filename = "$(basename).dfns.$(extension)"
-        moments_filename = "$(basename).moments.$(extension)"
+        dfns_filename = "$(basename).$(type).$(extension)"
+        if type == "dfns"
+            moments_filename = "$(basename).moments.$(extension)"
+        else
+            moments_filename = nothing
+        end
         backup_moments_filename = "$(basename)_$(counter).moments.$(extension)"
     else
         while true
@@ -186,8 +190,12 @@ function get_backup_filename(filename)
         end
         # Create dfns_filename here even though it is almost the filename passed in, in
         # order to get the right `iblock` for this block.
-        dfns_filename = "$(basename).dfns.$(iblock).$(extension)"
-        moments_filename = "$(basename).moments.$(iblock).$(extension)"
+        dfns_filename = "$(basename).$(type).$(iblock).$(extension)"
+        if type == "dfns"
+            moments_filename = "$(basename).moments.$(iblock).$(extension)"
+        else
+            moments_filename = nothing
+        end
         backup_moments_filename = "$(basename)_$(counter).moments.$(iblock).$(extension)"
     end
     backup_dfns_filename == "" && error("Failed to find a name for backup file.")
@@ -253,7 +261,9 @@ function get_prefix_iblock_and_move_existing_file(restart_filename, output_dir)
         # restarting from another run, and will not be overwriting the file.
         if (parallel_io && global_rank[] == 0) || (!parallel_io && block_rank[] == 0)
             mv(dfns_filename, backup_dfns_filename)
-            mv(moments_filename, backup_moments_filename)
+            if moments_filename !== nothing
+                mv(moments_filename, backup_moments_filename)
+            end
         end
     else
         # Reload from dfns_filename without moving the file

From 8abc30f369d16fdcc53b07b3bc6e8b777d53e88f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 12:13:02 +0000
Subject: [PATCH 115/428] Stop with error if initialization of electron pdf
 does not converge

---
 .../src/electron_kinetic_equation.jl           | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index aa3515b1f..6d25c2a3c 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -17,7 +17,7 @@ using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!
 using ..electron_vpa_advection: electron_vpa_advection!
-using ..file_io: write_initial_electron_state
+using ..file_io: write_initial_electron_state, finish_initial_electron_io
 using ..moment_constraints: hard_force_moment_constraints!
 using ..velocity_moments: integrate_over_vspace
 
@@ -453,9 +453,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     begin_serial_region()
     @serial_region begin
         if !electron_pdf_converged
-            # need to exit or handle this appropriately
-            println("!!!max number of iterations for electron pdf update exceeded!!!")
-            println("Stopping at ", Dates.format(now(), dateformat"H:MM:SS"))
             @loop_vpa ivpa begin
                 @loop_z iz begin
                     println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
@@ -469,6 +466,19 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         close(io_vth)
         close(io_pdf)
         close(io_pdf_stages)
+        if !electron_pdf_converged
+            # need to exit or handle this appropriately
+
+            if io_initial_electron !== nothing
+                output_counter += 1
+                write_initial_electron_state(pdf, moments, time, io_initial_electron,
+                                             output_counter, r, z, vperp, vpa)
+                finish_initial_electron_io(io_initial_electron)
+            end
+
+            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+        end
     end
     return time, output_counter
 end

From 5d0e2369b08d7e36968d9c6c55d778626eba988f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 13:23:40 +0000
Subject: [PATCH 116/428] Update electron vth before using

---
 moment_kinetics/src/electron_kinetic_equation.jl | 8 ++++----
 moment_kinetics/src/initial_conditions.jl        | 5 +++++
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 6d25c2a3c..2f43e5d41 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -146,11 +146,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #@. ddens_dz = (dppar_dz / ppar)*0.9
 
     @loop_r_z ir iz begin
+        # update the electron thermal speed using the updated electron parallel pressure
+        vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
         # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
         # and parallel pressure
         dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
-        # update the electron thermal speed itself using the updated electron parallel pressure
-        vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
     end
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
@@ -401,11 +401,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                  buffer_r_4, z_spectral, z)
             begin_r_z_region()
             @loop_r_z ir iz begin
+                # update the electron thermal speed using the updated electron parallel pressure
+                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
                 # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
                 # and parallel pressure
                 dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
-                # update the electron thermal speed itself using the updated electron parallel pressure
-                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
             end
         end
 
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index ec4eca8f3..acb2fb7af 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -520,6 +520,11 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
             fvec.pdf_electron .= pdf.electron.norm
         end
 
+        @loop_r_z ir iz begin
+            # update the electron thermal speed using the updated electron parallel pressure
+            moments.electron.vth[iz,ir] = sqrt(abs(2.0 * moments.electron.ppar[iz,ir] / (moments.electron.dens[iz,ir] * composition.me_over_mi)))
+        end
+
         moments.electron.qpar_updated[] = false
         calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
             moments.electron.ppar, moments.electron.upar, moments.electron.vth, 

From 107edf8a8aa1ddbaf52987e4086fe033e79036cd Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 17:06:46 +0000
Subject: [PATCH 117/428] Move ppar update to beginning of electron pseudo-time
 loop

This ensures that the updates of the distribution function and ppar are
done consistently at the same time level. Previously the distribution
function at one time level was being updated using ppar from the next
time level.
---
 .../src/electron_kinetic_equation.jl          | 132 +++++++++---------
 1 file changed, 66 insertions(+), 66 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 2f43e5d41..e8bae88b7 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -253,6 +253,72 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     # evolve (artificially) in time until the residual is less than the tolerance
     try
     while !electron_pdf_converged && (iteration <= max_electron_pdf_iterations)
+        #dt_energy = dt_electron * 10.0
+
+        # get an updated iterate of the electron parallel pressure
+        begin_r_z_region()
+        #ppar .= ppar_old
+        wpa3_moment = @view scratch_dummy.buffer_zrs_1[:,:,1]
+        @loop_r_z ir iz begin
+            wpa3_moment[iz,ir] = qpar[iz,ir] / vthe[iz,ir]^3
+        end
+        for i in 1:n_ppar_subcycles
+            @loop_r_z ir iz begin
+                qpar[iz,ir] = vthe[iz,ir]^3 * wpa3_moment[iz,ir]
+                dummy_zr[iz,ir] = -upar[iz,ir]
+            end
+            @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
+                                 buffer_r_4, z_spectral, z)
+            # Compute the upwinded z-derivative of the electron parallel pressure for the
+            # electron energy equation
+            @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
+                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                             buffer_r_5, buffer_r_6, z_spectral, z)
+            # centred second derivative for dissipation
+            if num_diss_params.moment_dissipation_coefficient > 0.0
+                @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
+                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+                begin_serial_region()
+                @serial_region begin
+                    if z.irank == 0
+                        moments.electron.d2ppar_dz2[1,:] .= 0.0
+                    end
+                    if z.irank == z.nrank - 1
+                        moments.electron.d2ppar_dz2[end,:] .= 0.0
+                    end
+                end
+            end
+
+            #dt_energy = dt_electron
+            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
+
+            # Apply same 'speed up' hack to ppar that we do to the distribution function,
+            # but without the wpa dependence.
+            @loop_r_z ir iz begin
+                zval = z.grid[iz]
+                znorm = 2.0*zval/Lz
+                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
+                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
+                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
+            end
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                fvec.electron_ppar[iz,ir] = ppar[iz,ir]
+            end
+
+            # compute the z-derivative of the updated electron parallel pressure
+            @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
+                                 buffer_r_4, z_spectral, z)
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                # update the electron thermal speed using the updated electron parallel pressure
+                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
+                # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
+                # and parallel pressure
+                dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
+            end
+        end
+
         begin_r_z_region()
         # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
@@ -343,72 +409,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             end
         end
 
-        #dt_energy = dt_electron * 10.0
-
-        # get an updated iterate of the electron parallel pressure
-        begin_r_z_region()
-        #ppar .= ppar_old
-        wpa3_moment = @view scratch_dummy.buffer_zrs_1[:,:,1]
-        @loop_r_z ir iz begin
-            wpa3_moment[iz,ir] = qpar[iz,ir] / vthe[iz,ir]^3
-        end
-        for i in 1:n_ppar_subcycles
-            @loop_r_z ir iz begin
-                qpar[iz,ir] = vthe[iz,ir]^3 * wpa3_moment[iz,ir]
-                dummy_zr[iz,ir] = -upar[iz,ir]
-            end
-            @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
-                                 buffer_r_4, z_spectral, z)
-            # Compute the upwinded z-derivative of the electron parallel pressure for the 
-            # electron energy equation
-            @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
-                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
-                             buffer_r_5, buffer_r_6, z_spectral, z)
-            # centred second derivative for dissipation
-            if num_diss_params.moment_dissipation_coefficient > 0.0
-                @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
-                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-                begin_serial_region()
-                @serial_region begin
-                    if z.irank == 0
-                        moments.electron.d2ppar_dz2[1,:] .= 0.0
-                    end
-                    if z.irank == z.nrank - 1
-                        moments.electron.d2ppar_dz2[end,:] .= 0.0
-                    end
-                end
-            end
-
-            #dt_energy = dt_electron
-            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
-
-            # Apply same 'speed up' hack to ppar that we do to the distribution function,
-            # but without the wpa dependence.
-            @loop_r_z ir iz begin
-                zval = z.grid[iz]
-                znorm = 2.0*zval/Lz
-                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
-                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
-                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
-            end
-            begin_r_z_region()
-            @loop_r_z ir iz begin
-                fvec.electron_ppar[iz,ir] = ppar[iz,ir]
-            end
-        
-            # compute the z-derivative of the updated electron parallel pressure
-            @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
-                                 buffer_r_4, z_spectral, z)
-            begin_r_z_region()
-            @loop_r_z ir iz begin
-                # update the electron thermal speed using the updated electron parallel pressure
-                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
-                # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
-                # and parallel pressure
-                dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
-            end
-        end
-
         #@views derivative_z!(dqpar_dz, qpar, 
         #        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
         #        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)

From 173f8f9e402313daf9de1b80165c6dfde5a71c56 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 17:41:53 +0000
Subject: [PATCH 118/428] Stop all processes by raising error when electron
 initilization fails

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

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index e8bae88b7..2c5647cfc 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -476,10 +476,12 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 finish_initial_electron_io(io_initial_electron)
             end
 
-            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
-                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
         end
     end
+    if !electron_pdf_converged
+        error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+              * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+    end
     return time, output_counter
 end
 

From 452081bf1f3617783d574d2407929af8ce67656b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 19:09:04 +0000
Subject: [PATCH 119/428] Initialize fvec.electron_ppar

fvec.ppar is used within electron_energy_equation(), so needs to be
initialized.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 2c5647cfc..740dd7dcd 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -151,6 +151,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
         # and parallel pressure
         dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
+        fvec.electron_ppar[iz,ir] = ppar[iz,ir]
     end
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation

From 53febfc95474b66132442cfd45bfd22174d6f7e8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 21 Feb 2024 20:10:41 +0000
Subject: [PATCH 120/428] Load coordinates without shared-memory scratch arrays

...when loading from the restart file and in `get_run_info()`.

The restart coordinates are created inside an `@serial_region begin`
where shared-memory arrays cannot be allocated, so `ignore_MPI=true`
must be passed through to `define_coordinate()`.

When calling `get_run_info()` MPI may not be set up, so trying to
allocate a shared-memory array may be an error.
---
 moment_kinetics/src/load_data.jl | 62 ++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 27 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index d6b3f4507..2b4a9ac2e 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -189,8 +189,13 @@ If `printout` is set to `true` a message will be printed when this function is c
 If `irank` and `nrank` are passed, then the `coord` and `spectral` objects returned will
 be set up for the parallelisation specified by `irank` and `nrank`, rather than the one
 implied by the output file.
+
+If `ignore_MPI=true` is passed, the returned coordinates will be created without shared
+memory scratch arrays (`ignore_MPI=true` will be passed through to
+[`define_coordinate`](@ref)).
 """
-function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing)
+function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing,
+                              ignore_MPI=false)
     if printout
         println("Loading $name coordinate data...")
     end
@@ -273,7 +278,7 @@ function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=no
                        discretization, fd_option, cheb_option, bc, advection_input("", 0.0, 0.0, 0.0),
                        MPI.COMM_NULL, element_spacing_option)
 
-    coord, spectral = define_coordinate(input, parallel_io)
+    coord, spectral = define_coordinate(input, parallel_io; ignore_MPI=ignore_MPI)
 
     return coord, spectral, chunk_size
 end
@@ -920,29 +925,29 @@ end
 function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io)
     if parallel_io
         restart_z, restart_z_spectral, _ =
-            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank)
+            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank, ignore_MPI=true)
         restart_r, restart_r_spectral, _ =
-            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank)
+            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank, ignore_MPI=true)
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank)
+            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank, ignore_MPI=true)
         restart_vpa, restart_vpa_spectral, _ =
-            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank)
+            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank, ignore_MPI=true)
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank)
+            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank, ignore_MPI=true)
         restart_vr, restart_vr_spectral, _ =
-            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank)
+            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank, ignore_MPI=true)
         restart_vz, restart_vz_spectral, _ =
-            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank)
+            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank, ignore_MPI=true)
     else
-        restart_z, restart_z_spectral, _ = load_coordinate_data(fid, "z")
-        restart_r, restart_r_spectral, _ = load_coordinate_data(fid, "r")
+        restart_z, restart_z_spectral, _ = load_coordinate_data(fid, "z"; ignore_MPI=true)
+        restart_r, restart_r_spectral, _ = load_coordinate_data(fid, "r"; ignore_MPI=true)
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp")
-        restart_vpa, restart_vpa_spectral, _ = load_coordinate_data(fid, "vpa")
+            load_coordinate_data(fid, "vperp"; ignore_MPI=true)
+        restart_vpa, restart_vpa_spectral, _ = load_coordinate_data(fid, "vpa"; ignore_MPI=true)
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta")
-        restart_vr, restart_vr_spectral, _ = load_coordinate_data(fid, "vr")
-        restart_vz, restart_vz_spectral, _ = load_coordinate_data(fid, "vz")
+            load_coordinate_data(fid, "vzeta"; ignore_MPI=true)
+        restart_vr, restart_vr_spectral, _ = load_coordinate_data(fid, "vr"; ignore_MPI=true)
+        restart_vz, restart_vz_spectral, _ = load_coordinate_data(fid, "vz"; ignore_MPI=true)
 
         if restart_r.nrank != r.nrank
             error("Not using parallel I/O, and distributed MPI layout has "
@@ -3222,24 +3227,25 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
     evolve_density, evolve_upar, evolve_ppar = load_mk_options(file_final_restart)
 
     z_local, z_local_spectral, z_chunk_size =
-        load_coordinate_data(file_final_restart, "z")
+        load_coordinate_data(file_final_restart, "z"; ignore_MPI=true)
     r_local, r_local_spectral, r_chunk_size =
-        load_coordinate_data(file_final_restart, "r")
-    r, r_spectral, z, z_spectral = construct_global_zr_coords(r_local, z_local)
+        load_coordinate_data(file_final_restart, "r"; ignore_MPI=true)
+    r, r_spectral, z, z_spectral = construct_global_zr_coords(r_local, z_local;
+                                                              ignore_MPI=true)
 
     if dfns
         vperp, vperp_spectral, vperp_chunk_size =
-            load_coordinate_data(file_final_restart, "vperp")
+            load_coordinate_data(file_final_restart, "vperp"; ignore_MPI=true)
         vpa, vpa_spectral, vpa_chunk_size =
-            load_coordinate_data(file_final_restart, "vpa")
+            load_coordinate_data(file_final_restart, "vpa"; ignore_MPI=true)
 
         if n_neutral_species > 0
             vzeta, vzeta_spectral, vzeta_chunk_size =
-                load_coordinate_data(file_final_restart, "vzeta")
+                load_coordinate_data(file_final_restart, "vzeta"; ignore_MPI=true)
             vr, vr_spectral, vr_chunk_size =
-                load_coordinate_data(file_final_restart, "vr")
+                load_coordinate_data(file_final_restart, "vr"; ignore_MPI=true)
             vz, vz_spectral, vz_chunk_size =
-                load_coordinate_data(file_final_restart, "vz")
+                load_coordinate_data(file_final_restart, "vz"; ignore_MPI=true)
         else
             dummy_adv_input = advection_input("default", 1.0, 0.0, 0.0)
             dummy_comm = MPI.COMM_NULL
@@ -3785,7 +3791,7 @@ end
 
 """
 """
-function construct_global_zr_coords(r_local, z_local)
+function construct_global_zr_coords(r_local, z_local; ignore_MPI=false)
 
     function make_global_input(coord_local)
         return grid_input(coord_local.name, coord_local.ngrid,
@@ -3794,8 +3800,10 @@ function construct_global_zr_coords(r_local, z_local)
             coord_local.advection, MPI.COMM_NULL, coord_local.element_spacing_option)
     end
 
-    r_global, r_global_spectral = define_coordinate(make_global_input(r_local))
-    z_global, z_global_spectral = define_coordinate(make_global_input(z_local))
+    r_global, r_global_spectral = define_coordinate(make_global_input(r_local);
+                                                    ignore_MPI=ignore_MPI)
+    z_global, z_global_spectral = define_coordinate(make_global_input(z_local);
+                                                    ignore_MPI=ignore_MPI)
 
     return r_global, r_global_spectral, z_global, z_global_spectral
 end

From 56b09710e070fc309dde99eb9b1813040ebcefc6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 22 Feb 2024 16:09:12 +0000
Subject: [PATCH 121/428] Option to skip wpa^1 moment constraint for electrons
 at sheath entrance

---
 moment_kinetics/src/electron_kinetic_equation.jl | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 740dd7dcd..a1a13a33e 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -499,7 +499,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
 
     # pdf_adjustment_option determines the velocity-dependent pre-factor for the
     # corrections to the pdf needed to ensure moment constraints are satisfied
-    pdf_adjustment_option = "vpa4_gaussian"
+    #pdf_adjustment_option = "vpa4"
+    #pdf_adjustment_option = "vpa4_gaussian"
+    pdf_adjustment_option = "no1st_vpa2"
 
     cutoff_step_width = 0.1
 
@@ -693,6 +695,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             #normalisation_constant_C = 0.0
             @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "no1st_vpa2"
+            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
         else
             println("pdf_adjustment_option not recognised")
             stop()
@@ -921,6 +928,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
             #normalisation_constant_C = 0.0
             @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "no1st_vpa2"
+            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
         else
             println("pdf_adjustment_option not recognised")
             stop()

From fa8d27bc9760d32edd249d8f6c7020a79cfa055c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 22 Feb 2024 22:02:07 +0000
Subject: [PATCH 122/428] Correct shift/interpolation of
 g_e(wpa(vpa))->g_e(wpa(-vpa))

Was missing factor of 2.
---
 .../src/electron_kinetic_equation.jl          | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index a1a13a33e..56c02713c 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -518,9 +518,20 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
         #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
         #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+
+        # Want to construct the w-grid corresponding to -vpa.
+        #   wpa(vpa) = (vpa - upar)/vth
+        #   ⇒ vpa = vth*wpa(vpa) + upar
+        #   wpa(-vpa) = (-vpa - upar)/vth
+        #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+        #             = (-vth*wpa - 2*upar)/vth
+        #             = -wpa - 2*upar/vth
+        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
         # Need to reverse vpa.grid because the grid passed as the second argument of
         # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[1,ir] / vthe[1,ir]
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
+        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
         @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
@@ -736,7 +747,13 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
         #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
         # Need to reverse vpa.grid because the grid passed as the second argument of
         # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- upar[end,ir] / vthe[end,ir]
+
+        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
+        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
         @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).

From 0272dfe92647f6275cf4dfc00de1c94a9118c0db Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 22 Feb 2024 22:03:08 +0000
Subject: [PATCH 123/428] Use sharp cutoff for g_e, instead of smoothed (tanh)
 step

---
 .../src/electron_kinetic_equation.jl          | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 56c02713c..5caef372f 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -594,11 +594,13 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
                0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
 
         #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
-        wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
+        #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
         #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
-        @loop_vpa ivpa begin
-            reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
-        end
+        #@loop_vpa ivpa begin
+        #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
+        #end
+        reversed_pdf[ivpa_max+1:end] .= 0
+        reversed_pdf[ivpa_max] *= fraction_of_pdf
         #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
         #println("done first cutoff loop")
         # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
@@ -825,10 +827,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa,
                0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
 
         #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
-        wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
-        @loop_vpa ivpa begin
-            reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
-        end
+        #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
+        #@loop_vpa ivpa begin
+        #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
+        #end
+        reversed_pdf[1:ivpa_min-1] .= 0
+        reversed_pdf[ivpa_min] *= fraction_of_pdf
         #println("done second cutoff loop")
 
         # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)

From 0ce07233153d8ee1b416c262adfaff096bc9f322 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 23 Feb 2024 13:17:59 +0000
Subject: [PATCH 124/428] Move logLambda to reference_parameters, add electron
 values

---
 moment_kinetics/src/krook_collisions.jl     |  7 +------
 moment_kinetics/src/reference_parameters.jl | 23 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index a9223933e..ddda4ebaa 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -18,12 +18,7 @@ function setup_krook_collisions(reference_params)
     mref = reference_params.mref
     timeref = reference_params.timeref
     cref = reference_params.cref
-
-    Nref_per_cm3 = Nref * 1.0e-6
-
-    # Coulomb logarithm at reference parameters for same-species ion-ion collisions, using
-    # NRL formulary. Formula given for n in units of cm^-3 and T in units of eV.
-    logLambda_ii = 23.0 - log(sqrt(2.0*Nref_per_cm3) / Tref^1.5)
+    logLambda_ii = reference_params.logLambda_ii
 
     # Collision frequency, using \hat{\nu} from Appendix, p. 277 of Helander "Collisional
     # Transport in Magnetized Plasmas" (2002).
diff --git a/moment_kinetics/src/reference_parameters.jl b/moment_kinetics/src/reference_parameters.jl
index efdbb4ded..12d960401 100644
--- a/moment_kinetics/src/reference_parameters.jl
+++ b/moment_kinetics/src/reference_parameters.jl
@@ -28,6 +28,29 @@ function setup_reference_parameters(input_dict)
     reference_parameter_section["timeref"] = reference_parameter_section["Lref"] / reference_parameter_section["cref"]
     reference_parameter_section["Omegaref"] = proton_charge * reference_parameter_section["Bref"] / reference_parameter_section["mref"]
 
+    Nref_per_cm3 = reference_parameter_section["Nref"] * 1.0e-6
+    Tref = reference_parameter_section["Tref"]
+
+    # Coulomb logarithm at reference parameters for same-species, singly-charged ion-ion
+    # collisions, using NRL formulary. Formula given for n in units of cm^-3 and T in
+    # units of eV.
+    reference_parameter_section["logLambda_ii"] = 23.0 - log(sqrt(2.0*Nref_per_cm3) / Tref^1.5)
+
+    # Coulomb logarithm at reference parameters for electron-electron collisions, using
+    # NRL formulary. Formula given for n in units of cm^-3 and T in units of eV.
+    reference_parameter_section["logLambda_ee"] = 23.5 - log(sqrt(Nref_per_cm3) / Tref^1.25) - sqrt(1.0e-5 + (log(Tref) -2.0)^2 / 16.0)
+
+    # Coulomb logarithm at reference parameters for electron-ion collisions with
+    # singly-charged ions, using NRL formulary. Formula given for n in units of cm^-3 and
+    # T in units of eV.
+    # Note: assume reference temperature is the same for ions and electrons, so ignore
+    # case in NRL formulary where Te < Ti*me/mi.
+    if Tref < 10.0
+        reference_parameter_section["logLambda_ei"] = 23.0 - log(sqrt(Nref_per_cm3) / Tref^1.5)
+    else
+        reference_parameter_section["logLambda_ei"] = 24.0 - log(sqrt(Nref_per_cm3) / Tref)
+    end
+
     reference_params = Dict_to_NamedTuple(reference_parameter_section)
 
     return reference_params

From cf22a8839561566ca79e32071e0fcf311921d97c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 23 Feb 2024 13:27:31 +0000
Subject: [PATCH 125/428] Move more Krook collisions setup into
 setup_krook_collisions!()

This encapsulates the Krook collisions functionality better in the
krook_collisions module, rather than having some moment_kinetics_input.
---
 moment_kinetics/src/krook_collisions.jl      | 19 ++++++++++++++++---
 moment_kinetics/src/moment_kinetics_input.jl | 17 +++--------------
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index ddda4ebaa..31ab409ca 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -2,7 +2,7 @@
 """
 module krook_collisions
 
-export setup_krook_collisions, get_collision_frequency, krook_collisions!
+export setup_krook_collisions!, get_collision_frequency, krook_collisions!
 
 using ..constants: epsilon0, proton_charge
 using ..looping
@@ -12,7 +12,7 @@ Calculate normalized collision frequency at reference parameters for Coulomb col
 
 Currently valid only for hydrogenic ions (Z=1)
 """
-function setup_krook_collisions(reference_params)
+function setup_krook_collisions!(collisions, reference_params, scan_input)
     Nref = reference_params.Nref
     Tref = reference_params.Tref
     mref = reference_params.mref
@@ -26,7 +26,20 @@ function setup_krook_collisions(reference_params)
                    (4.0 * π * epsilon0^2 * mref^2 * cref^3) # s^-1
     nu_ii0 = nu_ii0_per_s * timeref
 
-    return nu_ii0
+    collisions.krook_collisions_option = get(scan_input, "krook_collisions_option", "none")
+    if collisions.krook_collisions_option == "reference_parameters"
+        collisions.krook_collision_frequency_prefactor = nu_ii0
+    elseif collisions.krook_collisions_option == "manual" # get the frequency from the input file
+        collisions.krook_collision_frequency_prefactor = get(scan_input, "nuii_krook", nu_ii0)
+    elseif collisions.krook_collisions_option == "none"
+        # By default, no krook collisions included
+        collisions.krook_collision_frequency_prefactor = -1.0
+    else
+        error("Invalid option "
+              * "krook_collisions_option=$(collisions.krook_collisions_option) passed")
+    end
+
+    return nothing
 end
 
 """
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 660409d53..a9bcef6bf 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -13,7 +13,7 @@ using ..communication
 using ..coordinates: define_coordinate
 using ..external_sources
 using ..file_io: io_has_parallel, input_option_error, open_ascii_output_file
-using ..krook_collisions: setup_krook_collisions
+using ..krook_collisions: setup_krook_collisions!
 using ..finite_differences: fd_check_option
 using ..input_structs
 using ..numerical_dissipation: setup_numerical_dissipation
@@ -177,19 +177,8 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     collisions.ionization_energy = get(scan_input, "ionization_energy", 0.0)
     collisions.constant_ionization_rate = get(scan_input, "constant_ionization_rate", false)
     collisions.nu_ei = get(scan_input, "nu_ei", 0.0)
-    collisions.krook_collisions_option = get(scan_input, "krook_collisions_option", "none")
-    nuii_krook_default = setup_krook_collisions(reference_params)
-    if collisions.krook_collisions_option == "reference_parameters"
-        collisions.krook_collision_frequency_prefactor = nuii_krook_default
-    elseif collisions.krook_collisions_option == "manual" # get the frequency from the input file
-        collisions.krook_collision_frequency_prefactor = get(scan_input, "nuii_krook", nuii_krook_default)
-    elseif collisions.krook_collisions_option == "none"
-        # By default, no krook collisions included
-        collisions.krook_collision_frequency_prefactor = -1.0
-    else
-        error("Invalid option "
-              * "krook_collisions_option=$(collisions.krook_collisions_option) passed")
-    end
+    # Set options for Krook collisions in the `collisions` struct
+    setup_krook_collisions!(collisions, reference_params, scan_input)
     # set the Fokker-Planck collision frequency
     collisions.nuii = get(scan_input, "nuii", 0.0)
     

From 8c90caa81d34a5df8f2d4d8c5b462aa4a9825310 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 23 Feb 2024 14:19:08 +0000
Subject: [PATCH 126/428] Create function for Krook collisions for electrons

---
 .../src/makie_post_processing.jl              |   1 -
 moment_kinetics/src/input_structs.jl          |   8 +-
 moment_kinetics/src/krook_collisions.jl       | 305 ++++++++++++++++--
 moment_kinetics/src/load_data.jl              |  16 +-
 moment_kinetics/src/moment_kinetics_input.jl  |   2 +
 5 files changed, 307 insertions(+), 25 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 1ff58364f..978d8e839 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -33,7 +33,6 @@ using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input,
                                      set_defaults_and_check_top_level!,
                                      set_defaults_and_check_section!, Dict_to_NamedTuple
-using moment_kinetics.krook_collisions: get_collision_frequency
 using moment_kinetics.looping: all_dimensions, ion_dimensions, neutral_dimensions
 using moment_kinetics.manufactured_solns: manufactured_solutions,
                                           manufactured_electric_fields
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 041c9e95c..081a7d744 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -331,8 +331,12 @@ mutable struct collisions_input
     constant_ionization_rate::Bool
     # electron-ion collision frequency
     nu_ei::mk_float
-    # Coulomb collision rate at the reference density and temperature
-    krook_collision_frequency_prefactor::mk_float
+    # Ion-ion Coulomb collision rate at the reference density and temperature
+    krook_collision_frequency_prefactor_ii::mk_float
+    # Electron-electron Coulomb collision rate at the reference density and temperature
+    krook_collision_frequency_prefactor_ee::mk_float
+    # Electron-ion Coulomb collision rate at the reference density and temperature
+    krook_collision_frequency_prefactor_ei::mk_float
     # Setting to switch between different options for Krook collision operator
     krook_collisions_option::String
     # ion-ion self collision frequency
diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index 31ab409ca..ca0c47037 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -2,9 +2,10 @@
 """
 module krook_collisions
 
-export setup_krook_collisions!, get_collision_frequency, krook_collisions!
+export setup_krook_collisions!, get_collision_frequency_ii, get_collision_frequency_ee,
+       get_collision_frequency_ei, krook_collisions!, electron_krook_collisions!
 
-using ..constants: epsilon0, proton_charge
+using ..constants
 using ..looping
 
 """
@@ -19,21 +20,69 @@ function setup_krook_collisions!(collisions, reference_params, scan_input)
     timeref = reference_params.timeref
     cref = reference_params.cref
     logLambda_ii = reference_params.logLambda_ii
+    logLambda_ee = reference_params.logLambda_ee
+    logLambda_ei = reference_params.logLambda_ei
 
-    # Collision frequency, using \hat{\nu} from Appendix, p. 277 of Helander "Collisional
-    # Transport in Magnetized Plasmas" (2002).
-    nu_ii0_per_s = Nref * proton_charge^4 * logLambda_ii  /
+    # Collision frequencies, using \hat{\nu} from Appendix, p. 277 of Helander
+    # "Collisional Transport in Magnetized Plasmas" (2002).
+    nu_ii0_per_s = Nref * proton_charge^4 * logLambda_ii /
                    (4.0 * π * epsilon0^2 * mref^2 * cref^3) # s^-1
     nu_ii0 = nu_ii0_per_s * timeref
 
+    # Note the electron thermal speed used in the code is normalised to cref, so we use
+    # cref in these two formulas rather than a reference electron thermal speed, so that
+    # when multiplied by the normalised electron thermal speed we get the correct
+    # normalised collision frequency.
+    nu_ee0_per_s = Nref * proton_charge^4 * logLambda_ee /
+                   (4.0 * π * epsilon0^2 * electron_mass^2 * cref^3) # s^-1
+    nu_ee0 = nu_ee0_per_s * timeref
+
+    nu_ei0_per_s = Nref * proton_charge^4 * logLambda_ei /
+                   (4.0 * π * epsilon0^2 * electron_mass^2 * cref^3) # s^-1
+    nu_ei0 = nu_ei0_per_s * timeref
+
     collisions.krook_collisions_option = get(scan_input, "krook_collisions_option", "none")
     if collisions.krook_collisions_option == "reference_parameters"
-        collisions.krook_collision_frequency_prefactor = nu_ii0
+        collisions.krook_collision_frequency_prefactor_ii = nu_ii0
+    elseif collisions.krook_collisions_option == "manual" # get the frequency from the input file
+        collisions.krook_collision_frequency_prefactor_ii = get(scan_input, "nuii_krook", nu_ii0)
+    elseif collisions.krook_collisions_option == "none"
+        # By default, no krook collisions included
+        collisions.krook_collision_frequency_prefactor_ii = -1.0
+    else
+        error("Invalid option "
+              * "krook_collisions_option=$(collisions.krook_collisions_option) passed")
+    end
+
+    if collisions.krook_collisions_option == "reference_parameters"
+        collisions.krook_collision_frequency_prefactor_ee = nu_ee0
+    elseif collisions.krook_collisions_option == "manual" # get the frequency from the input file
+        # If the "manual" option is used, the collision frequency is not multiplied by
+        # vthe^(-3), so need to correct it to be evaluated with the electron thermal speed
+        # at Tref for the default value.
+        vthe_ref = sqrt(2.0 * Tref / electron_mass)
+        collisions.krook_collision_frequency_prefactor_ee =
+            get(scan_input, "nuee_krook", nu_ee0 * (cref/vthe_ref)^3)
+    elseif collisions.krook_collisions_option == "none"
+        # By default, no krook collisions included
+        collisions.krook_collision_frequency_prefactor_ee = -1.0
+    else
+        error("Invalid option "
+              * "krook_collisions_option=$(collisions.krook_collisions_option) passed")
+    end
+
+    if collisions.krook_collisions_option == "reference_parameters"
+        collisions.krook_collision_frequency_prefactor_ei = nu_ei0
     elseif collisions.krook_collisions_option == "manual" # get the frequency from the input file
-        collisions.krook_collision_frequency_prefactor = get(scan_input, "nuii_krook", nu_ii0)
+        # If the "manual" option is used, the collision frequency is not multiplied by
+        # vthe^(-3), so need to correct it to be evaluated with the electron thermal speed
+        # at Tref for the default value.
+        vthe_ref = sqrt(2.0 * Tref / electron_mass)
+        collisions.krook_collision_frequency_prefactor_ei =
+            get(scan_input, "nuei_krook", nu_ei0 * (cref/vthe_ref)^3)
     elseif collisions.krook_collisions_option == "none"
         # By default, no krook collisions included
-        collisions.krook_collision_frequency_prefactor = -1.0
+        collisions.krook_collision_frequency_prefactor_ei = -1.0
     else
         error("Invalid option "
               * "krook_collisions_option=$(collisions.krook_collisions_option) passed")
@@ -43,24 +92,88 @@ function setup_krook_collisions!(collisions, reference_params, scan_input)
 end
 
 """
-    get_collision_frequency(collisions, n, vth)
+    get_collision_frequency_ii(collisions, n, vth)
 
-Calculate the collision frequency, depending on the settings/parameters in `collisions`,
-for the given density `n` and thermal speed `vth`.
+Calculate the ion-ion collision frequency, depending on the settings/parameters in
+`collisions`, for the given density `n` and thermal speed `vth`.
 
 `n` and `vth` may be scalars or arrays, but should have shapes that can be broadcasted
 together.
 """
-function get_collision_frequency(collisions, n, vth)
+function get_collision_frequency_ii(collisions, n, vth)
     if collisions.krook_collisions_option == "reference_parameters"
-        return @. collisions.krook_collision_frequency_prefactor * n * vth^(-3)
+        return @. collisions.krook_collision_frequency_prefactor_ii * n * vth^(-3)
     elseif collisions.krook_collisions_option == "manual"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
         # returning a scalar from this branch but an array from the "reference_parameters"
         # branch).
-        return @. collisions.krook_collision_frequency_prefactor + 0.0 * n
+        return @. collisions.krook_collision_frequency_prefactor_ii + 0.0 * n
+    elseif collisions.krook_collisions_option == "none"
+        # Include 0.0*n so that the result gets promoted to an array if n is an array,
+        # which hopefully means this function will have a fixed return type given the
+        # types of the arguments (we don't want to be 'type unstable' for array inputs by
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. 0.0 * n
+    else
+        error("Unrecognised option "
+              * "krook_collisions_option=$(collisions.krook_collisions_option)")
+    end
+end
+
+"""
+    get_collision_frequency_ee(collisions, n, vthe)
+
+Calculate the electron-electron collision frequency, depending on the settings/parameters
+in `collisions`, for the given density `n` and electron thermal speed `vthe`.
+
+`n` and `vthe` may be scalars or arrays, but should have shapes that can be broadcasted
+together.
+"""
+function get_collision_frequency_ee(collisions, n, vthe)
+    if collisions.krook_collisions_option == "reference_parameters"
+        return @. collisions.krook_collision_frequency_prefactor_ee * n * vthe^(-3)
+    elseif collisions.krook_collisions_option == "manual"
+        # Include 0.0*n so that the result gets promoted to an array if n is an array,
+        # which hopefully means this function will have a fixed return type given the
+        # types of the arguments (we don't want to be 'type unstable' for array inputs by
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. collisions.krook_collision_frequency_prefactor_ee + 0.0 * n
+    elseif collisions.krook_collisions_option == "none"
+        # Include 0.0*n so that the result gets promoted to an array if n is an array,
+        # which hopefully means this function will have a fixed return type given the
+        # types of the arguments (we don't want to be 'type unstable' for array inputs by
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. 0.0 * n
+    else
+        error("Unrecognised option "
+              * "krook_collisions_option=$(collisions.krook_collisions_option)")
+    end
+end
+
+"""
+    get_collision_frequency_ei(collisions, n, vthe)
+
+Calculate the electron-electron collision frequency, depending on the settings/parameters
+in `collisions`, for the given density `n` and electron thermal speed `vthe`.
+
+`n` and `vthe` may be scalars or arrays, but should have shapes that can be broadcasted
+together.
+"""
+function get_collision_frequency_ei(collisions, n, vthe)
+    if collisions.krook_collisions_option == "reference_parameters"
+        return @. collisions.krook_collision_frequency_prefactor_ei * n * vthe^(-3)
+    elseif collisions.krook_collisions_option == "manual"
+        # Include 0.0*n so that the result gets promoted to an array if n is an array,
+        # which hopefully means this function will have a fixed return type given the
+        # types of the arguments (we don't want to be 'type unstable' for array inputs by
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. collisions.krook_collision_frequency_prefactor_ei + 0.0 * n
     elseif collisions.krook_collisions_option == "none"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
@@ -94,7 +207,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            nu_ii = get_collision_frequency(collisions, n, vth)
+            nu_ii = get_collision_frequency_ii(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
                     (fvec_in.pdf[ivpa,ivperp,iz,ir,is]
@@ -107,7 +220,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            nu_ii = get_collision_frequency(collisions, n, vth)
+            nu_ii = get_collision_frequency_ii(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
                     (fvec_in.pdf[ivpa,ivperp,iz,ir,is]
@@ -120,7 +233,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            nu_ii = get_collision_frequency(collisions, n, vth)
+            nu_ii = get_collision_frequency_ii(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
                     (fvec_in.pdf[ivpa,ivperp,iz,ir,is]
@@ -134,7 +247,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
         @loop_s_r_z is ir iz begin
             n = fvec_in.density[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            nu_ii = get_collision_frequency(collisions, n, vth)
+            nu_ii = get_collision_frequency_ii(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
                 (fvec_in.pdf[ivpa,ivperp,iz,ir,is]
@@ -152,7 +265,7 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
             else
                 vth_prefactor = 1.0 / vth^3
             end
-            nu_ii = get_collision_frequency(collisions, n, vth)
+            nu_ii = get_collision_frequency_ii(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir,is] -= dt * nu_ii *
                     (fvec_in.pdf[ivpa,ivperp,iz,ir,is]
@@ -166,4 +279,158 @@ function krook_collisions!(pdf_out, fvec_in, moments, composition, collisions, v
     return nothing
 end
 
+"""
+Add Krook collision operator for electrons
+"""
+function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, collisions,
+                                    vperp, vpa, dt)
+    begin_r_z_region()
+
+    # For now, electrons are always fully moment-kinetic
+    evolve_density = true
+    evolve_upar = true
+    evolve_ppar = true
+
+    if vperp.n > 1 && (evolve_density || evolve_upar || evolve_ppar)
+        error("Krook collisions not implemented for 2V moment-kinetic cases yet")
+    end
+
+    # Note: do not need 1/sqrt(pi) for the 'Maxwellian' term because the pdf is already
+    # normalized by sqrt(pi) (see velocity_moments.integrate_over_vspace).
+    if evolve_ppar && evolve_upar
+        # Compared to evolve_upar version, grid is already normalized by vth, and multiply
+        # through by vth, remembering pdf is already multiplied by vth
+        @loop_r_z ir iz begin
+            n = fvec_in.electron_density[iz,ir]
+            vth = moments.electron.vth[iz,ir]
+            nu_ee = get_collision_frequency_ee(collisions, n, vth)
+            nu_ei = get_collision_frequency_ei(collisions, n, vth)
+
+            # e-i collisions push electrons towards a Maxwellian drifting at the ion
+            # parallel flow, so need a corresponding normalised parallel velocity
+            # coordinate.
+            # For now, assume there is only one ion species rather than bothering to
+            # calculate an average ion flow speed, or sum over ion species here.
+            @. vpa.scratch = vpa.grid + (fvec_in.upar[iz,ir,1] - fvec_in.electron_upar[iz,ir]) / vth
+
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= dt * (
+                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                             - exp(-vpa.grid[ivpa]^2 - vperp.grid[ivperp]^2))
+                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                               - exp(-vpa.scratch[ivpa]^2 - vperp.grid[ivperp]^2))
+                   )
+            end
+        end
+    elseif evolve_ppar
+        # Compared to full-f collision operater, multiply through by vth, remembering pdf
+        # is already multiplied by vth, and grid is already normalized by vth
+        @loop_r_z ir iz begin
+            n = fvec_in.electron_density[iz,ir]
+            vth = moments.electron.vth[iz,ir]
+            nu_ee = get_collision_frequency_ee(collisions, n, vth)
+            nu_ei = get_collision_frequency_ei(collisions, n, vth)
+
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= dt * (
+                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                             - exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir])/vth)^2
+                                   - (vperp.grid[ivperp]/vth)^2))
+                    # e-i collisions push electrons towards a Maxwellian drifting at the ion
+                    # parallel flow, so need a corresponding normalised parallel velocity
+                    # coordinate.
+                    # For now, assume there is only one ion species rather than bothering to
+                    # calculate an average ion flow speed, or sum over ion species here.
+                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                               - exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1])/vth)^2
+                                     - (vperp.grid[ivperp]/vth)^2))
+                   )
+            end
+        end
+    elseif evolve_upar
+        # Compared to evolve_density version, grid is already shifted by upar
+        @loop_r_z ir iz begin
+            n = fvec_in.electron_density[iz,ir]
+            vth = moments.electron.vth[iz,ir]
+            nu_ee = get_collision_frequency_ee(collisions, n, vth)
+            nu_ei = get_collision_frequency_ei(collisions, n, vth)
+
+            # e-i collisions push electrons towards a Maxwellian drifting at the ion
+            # parallel flow, so need a corresponding normalised parallel velocity
+            # coordinate.
+            # For now, assume there is only one ion species rather than bothering to
+            # calculate an average ion flow speed, or sum over ion species here.
+            @. vpa.scratch = vpa.grid + (fvec_in.upar[iz,ir,1] - fvec_in.electron_upar[iz,ir])
+
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= dt * (
+                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                             - 1.0 / vth * exp(-(vpa.grid[ivpa] / vth)^2
+                                               - (vperp.grid[ivperp] / vth)^2))
+                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                               - 1.0 / vth * exp(-(vpa.scratch[ivpa] / vth)^2
+                                                 - (vperp.grid[ivperp] / vth)^2))
+                   )
+            end
+        end
+    elseif evolve_density
+        # Compared to full-f collision operater, divide through by density, remembering
+        # that pdf is already normalized by density
+        @loop_r_z ir iz begin
+            n = fvec_in.electron_density[iz,ir]
+            vth = moments.electron.vth[iz,ir]
+            nu_ee = get_collision_frequency_ee(collisions, n, vth)
+            nu_ei = get_collision_frequency_ei(collisions, n, vth)
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= dt * (
+                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                             - 1.0 / vth
+                             * exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir]) / vth)^2
+                                   - (vperp.grid[ivperp]/vth)^2))
+                    # e-i collisions push electrons towards a Maxwellian drifting at the ion
+                    # parallel flow, so need a corresponding normalised parallel velocity
+                    # coordinate.
+                    # For now, assume there is only one ion species rather than bothering to
+                    # calculate an average ion flow speed, or sum over ion species here.
+                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                               - 1.0 / vth
+                               * exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1]) / vth)^2
+                                     - (vperp.grid[ivperp]/vth)^2))
+                   )
+            end
+        end
+    else
+        @loop_r_z ir iz begin
+            n = fvec_in.electron_density[iz,ir]
+            vth = moments.electron.vth[iz,ir]
+            if vperp.n == 1
+                vth_prefactor = 1.0 / vth
+            else
+                vth_prefactor = 1.0 / vth^3
+            end
+            nu_ee = get_collision_frequency_ee(collisions, n, vth)
+            nu_ei = get_collision_frequency_ei(collisions, n, vth)
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= dt * (
+                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                             - n * vth_prefactor
+                             * exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir])/vth)^2
+                                   - (vperp.grid[ivperp]/vth)^2))
+                    # e-i collisions push electrons towards a Maxwellian drifting at the ion
+                    # parallel flow, so need a corresponding normalised parallel velocity
+                    # coordinate.
+                    # For now, assume there is only one ion species rather than bothering to
+                    # calculate an average ion flow speed, or sum over ion species here.
+                    + nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                               - n * vth_prefactor
+                               * exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1])/vth)^2
+                                     - (vperp.grid[ivperp]/vth)^2))
+                   )
+            end
+        end
+    end
+
+    return nothing
+end
+
 end # krook_collisions
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 18c75b7ed..550dc2eb7 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -18,6 +18,7 @@ using ..coordinates: coordinate, define_coordinate
 using ..file_io: check_io_implementation, get_group, get_subgroup_keys, get_variable_keys
 using ..input_structs: advection_input, grid_input, hdf5, netcdf
 using ..interpolation: interpolate_to_grid_1d!
+using ..krook_collisions
 using ..looping
 using ..moment_kinetics_input: mk_input
 using ..type_definitions: mk_float, mk_int
@@ -29,7 +30,8 @@ using MPI
 const em_variables = ("phi", "Er", "Ez")
 const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "thermal_speed", "temperature", "parallel_heat_flux",
-                              "collision_frequency", "sound_speed", "mach_number")
+                              "collision_frequency_ii", "collision_frequency_ee",
+                              "collision_frequency_ei", "sound_speed", "mach_number")
 const electron_moment_variables = ("electron_density", "electron_parallel_flow",
                                    "electron_parallel_pressure", "electron_thermal_speed",
                                    "electron_temperature", "electron_parallel_heat_flux")
@@ -3690,10 +3692,18 @@ function get_variable(run_info, variable_name; kwargs...)
     if variable_name == "temperature"
         vth = postproc_load_variable(run_info, "thermal_speed"; kwargs...)
         variable = vth.^2
-    elseif variable_name == "collision_frequency"
+    elseif variable_name == "collision_frequency_ii"
         n = postproc_load_variable(run_info, "density"; kwargs...)
         vth = postproc_load_variable(run_info, "thermal_speed"; kwargs...)
-        variable = get_collision_frequency(run_info.collisions, n, vth)
+        variable = get_collision_frequency_ii(run_info.collisions, n, vth)
+    elseif variable_name == "collision_frequency_ee"
+        n = postproc_load_variable(run_info, "electron_density"; kwargs...)
+        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        variable = get_collision_frequency_ee(run_info.collisions, n, vth)
+    elseif variable_name == "collision_frequency_ei"
+        n = postproc_load_variable(run_info, "electron_density"; kwargs...)
+        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        variable = get_collision_frequency_ei(run_info.collisions, n, vth)
     elseif variable_name == "temperature_neutral"
         vth = postproc_load_variable(run_info, "thermal_speed_neutral"; kwargs...)
         variable = vth.^2
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index a9bcef6bf..71f01cad9 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1038,6 +1038,8 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     collisions = collisions_input(charge_exchange, charge_exchange_electron, ionization,
                                   ionization_electron, ionization_energy,
                                   constant_ionization_rate, nu_ei,
+                                  krook_collision_frequency_prefactor,
+                                  krook_collision_frequency_prefactor,
                                   krook_collision_frequency_prefactor, "none", nuii)
     Bzed = 1.0 # magnetic field component along z
     Bmag = 1.0 # magnetic field strength

From 2f753a9a3c636c056d674489155bbaaa561b0fab Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 23 Feb 2024 14:22:29 +0000
Subject: [PATCH 127/428] Remove use of `fvec_in` in
 electron_krook_collisions!()

The electron kinetic equation does not currently use the `scratch_pdf`
structs, so cannot use them in this function.
---
 moment_kinetics/src/krook_collisions.jl | 60 ++++++++++++-------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index ca0c47037..7df4c82ff 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -282,8 +282,8 @@ end
 """
 Add Krook collision operator for electrons
 """
-function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, collisions,
-                                    vperp, vpa, dt)
+function electron_krook_collisions!(pdf_out, pdf_in, dens_in, upar_in, upar_ion_in,
+                                    vth_in, collisions, vperp, vpa, dt)
     begin_r_z_region()
 
     # For now, electrons are always fully moment-kinetic
@@ -301,8 +301,8 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
         # Compared to evolve_upar version, grid is already normalized by vth, and multiply
         # through by vth, remembering pdf is already multiplied by vth
         @loop_r_z ir iz begin
-            n = fvec_in.electron_density[iz,ir]
-            vth = moments.electron.vth[iz,ir]
+            n = dens_in[iz,ir]
+            vth = vth_in[iz,ir]
             nu_ee = get_collision_frequency_ee(collisions, n, vth)
             nu_ei = get_collision_frequency_ei(collisions, n, vth)
 
@@ -311,13 +311,13 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
             # coordinate.
             # For now, assume there is only one ion species rather than bothering to
             # calculate an average ion flow speed, or sum over ion species here.
-            @. vpa.scratch = vpa.grid + (fvec_in.upar[iz,ir,1] - fvec_in.electron_upar[iz,ir]) / vth
+            @. vpa.scratch = vpa.grid + (upar_ion_in[iz,ir,1] - upar_in[iz,ir]) / vth
 
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= dt * (
-                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
                              - exp(-vpa.grid[ivpa]^2 - vperp.grid[ivperp]^2))
-                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    + nu_ei * (pdf_in[ivpa,ivperp,iz,ir]
                                - exp(-vpa.scratch[ivpa]^2 - vperp.grid[ivperp]^2))
                    )
             end
@@ -326,23 +326,23 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
         # Compared to full-f collision operater, multiply through by vth, remembering pdf
         # is already multiplied by vth, and grid is already normalized by vth
         @loop_r_z ir iz begin
-            n = fvec_in.electron_density[iz,ir]
-            vth = moments.electron.vth[iz,ir]
+            n = dens_in[iz,ir]
+            vth = vth_in[iz,ir]
             nu_ee = get_collision_frequency_ee(collisions, n, vth)
             nu_ei = get_collision_frequency_ei(collisions, n, vth)
 
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= dt * (
-                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
-                             - exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir])/vth)^2
+                    nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
+                             - exp(-((vpa.grid[ivpa] - upar_in[iz,ir])/vth)^2
                                    - (vperp.grid[ivperp]/vth)^2))
                     # e-i collisions push electrons towards a Maxwellian drifting at the ion
                     # parallel flow, so need a corresponding normalised parallel velocity
                     # coordinate.
                     # For now, assume there is only one ion species rather than bothering to
                     # calculate an average ion flow speed, or sum over ion species here.
-                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
-                               - exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1])/vth)^2
+                    + nu_ei * (pdf_in[ivpa,ivperp,iz,ir]
+                               - exp(-((vpa.grid[ivpa] - upar_ion_in[iz,ir,1])/vth)^2
                                      - (vperp.grid[ivperp]/vth)^2))
                    )
             end
@@ -350,8 +350,8 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
     elseif evolve_upar
         # Compared to evolve_density version, grid is already shifted by upar
         @loop_r_z ir iz begin
-            n = fvec_in.electron_density[iz,ir]
-            vth = moments.electron.vth[iz,ir]
+            n = dens_in[iz,ir]
+            vth = vth_in[iz,ir]
             nu_ee = get_collision_frequency_ee(collisions, n, vth)
             nu_ei = get_collision_frequency_ei(collisions, n, vth)
 
@@ -360,14 +360,14 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
             # coordinate.
             # For now, assume there is only one ion species rather than bothering to
             # calculate an average ion flow speed, or sum over ion species here.
-            @. vpa.scratch = vpa.grid + (fvec_in.upar[iz,ir,1] - fvec_in.electron_upar[iz,ir])
+            @. vpa.scratch = vpa.grid + (upar_ion_in[iz,ir,1] - upar_in[iz,ir])
 
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= dt * (
-                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
                              - 1.0 / vth * exp(-(vpa.grid[ivpa] / vth)^2
                                                - (vperp.grid[ivperp] / vth)^2))
-                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    + nu_ei * (pdf_in[ivpa,ivperp,iz,ir]
                                - 1.0 / vth * exp(-(vpa.scratch[ivpa] / vth)^2
                                                  - (vperp.grid[ivperp] / vth)^2))
                    )
@@ -377,32 +377,32 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
         # Compared to full-f collision operater, divide through by density, remembering
         # that pdf is already normalized by density
         @loop_r_z ir iz begin
-            n = fvec_in.electron_density[iz,ir]
-            vth = moments.electron.vth[iz,ir]
+            n = dens_in[iz,ir]
+            vth = vth_in[iz,ir]
             nu_ee = get_collision_frequency_ee(collisions, n, vth)
             nu_ei = get_collision_frequency_ei(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= dt * (
-                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
                              - 1.0 / vth
-                             * exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir]) / vth)^2
+                             * exp(-((vpa.grid[ivpa] - upar_in[iz,ir]) / vth)^2
                                    - (vperp.grid[ivperp]/vth)^2))
                     # e-i collisions push electrons towards a Maxwellian drifting at the ion
                     # parallel flow, so need a corresponding normalised parallel velocity
                     # coordinate.
                     # For now, assume there is only one ion species rather than bothering to
                     # calculate an average ion flow speed, or sum over ion species here.
-                    + nu_ei * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    + nu_ei * (pdf_in[ivpa,ivperp,iz,ir]
                                - 1.0 / vth
-                               * exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1]) / vth)^2
+                               * exp(-((vpa.grid[ivpa] - upar_ion_in[iz,ir,1]) / vth)^2
                                      - (vperp.grid[ivperp]/vth)^2))
                    )
             end
         end
     else
         @loop_r_z ir iz begin
-            n = fvec_in.electron_density[iz,ir]
-            vth = moments.electron.vth[iz,ir]
+            n = dens_in[iz,ir]
+            vth = vth_in[iz,ir]
             if vperp.n == 1
                 vth_prefactor = 1.0 / vth
             else
@@ -412,18 +412,18 @@ function electron_krook_collisions!(pdf_out, fvec_in, moments, composition, coll
             nu_ei = get_collision_frequency_ei(collisions, n, vth)
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= dt * (
-                    nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
                              - n * vth_prefactor
-                             * exp(-((vpa.grid[ivpa] - fvec_in.electron_upar[iz,ir])/vth)^2
+                             * exp(-((vpa.grid[ivpa] - upar_in[iz,ir])/vth)^2
                                    - (vperp.grid[ivperp]/vth)^2))
                     # e-i collisions push electrons towards a Maxwellian drifting at the ion
                     # parallel flow, so need a corresponding normalised parallel velocity
                     # coordinate.
                     # For now, assume there is only one ion species rather than bothering to
                     # calculate an average ion flow speed, or sum over ion species here.
-                    + nu_ee * (fvec_in.electron_pdf[ivpa,ivperp,iz,ir]
+                    + nu_ee * (pdf_in[ivpa,ivperp,iz,ir]
                                - n * vth_prefactor
-                               * exp(-((vpa.grid[ivpa] - fvec_in.upar[iz,ir,1])/vth)^2
+                               * exp(-((vpa.grid[ivpa] - upar_ion_in[iz,ir,1])/vth)^2
                                      - (vperp.grid[ivperp]/vth)^2))
                    )
             end

From 9d9fde7b1b3cac78e2cbd6d63d9462f67ba751b0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 23 Feb 2024 14:31:19 +0000
Subject: [PATCH 128/428] Implement option for Krook collisions on kinetic
 electrons

No Krook collision terms currently implemented in the moment equations.
This is OK if the ions and electrons are forced to have the same
parallel flow (by ambipolarity plus zero sheath currents assumed) as
then there is no friction. We can also neglect temperature equilibration
as small in sqrt(mass ratio).
---
 .../src/electron_kinetic_equation.jl          | 30 +++++++++++++------
 moment_kinetics/src/input_structs.jl          |  2 +-
 moment_kinetics/src/moment_kinetics.jl        |  2 +-
 moment_kinetics/src/time_advance.jl           | 14 ++++-----
 4 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 5caef372f..4a7020884 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -18,6 +18,7 @@ using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!
 using ..electron_vpa_advection: electron_vpa_advection!
 using ..file_io: write_initial_electron_state, finish_initial_electron_io
+using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
 using ..velocity_moments: integrate_over_vspace
 
@@ -122,6 +123,8 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     #ppar_old .= ppar
 
     upar = moments.electron.upar
+    println("WARNING: when time-evolving the ions, this should really be something like fvec_in.upar, so that it is defined at the correct time level...")
+    upar_ion = moments.ion.upar
 
     # create a (z,r) dimension dummy array for use in taking derivatives
     dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
@@ -185,10 +188,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     # initialise the electron pdf convergence flag to false
     electron_pdf_converged = false
     # calculate the residual of the electron kinetic equation for the initial guess of the electron pdf
-    dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vthe, ppar, 
+    dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, upar_ion,
                                         ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
-                                        z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                        num_diss_params, dt_max)
+                                        z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                        collisions, num_diss_params, dt_max)
     # Divide by wpa to relax CFL condition at large wpa - only looking for steady
     # state here, so does not matter that this makes time evolution incorrect.
     # Also increase the effective timestep for z-values far from the sheath boundary -
@@ -416,10 +419,10 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # TMP FOR TESTING
         #dqpar_dz .= 0.0
         # calculate the residual of the electron kinetic equation for the updated electron pdf
-        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vthe, ppar, ddens_dz, 
+        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, ddens_dz, upar_ion,
                                             dppar_dz, dqpar_dz, dvth_dz, 
-                                            z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                            num_diss_params, dt_max)
+                                            z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                            collisions, num_diss_params, dt_max)
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
@@ -1219,10 +1222,10 @@ INPUTS:
 OUTPUT:
     residual = updated residual of the electron kinetic equation
 """
-function electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, vth, ppar, 
+function electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vth, ppar, upar_ion,
                                              ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
-                                             z, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                             num_diss_params, dt_electron)
+                                             z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                             collisions, num_diss_params, dt_electron)
 
     # initialise the residual to zero                                             
     begin_r_vperp_vpa_region()
@@ -1278,6 +1281,15 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     # end
     # stop()
     #dt_max = min(dt_max_zadv, dt_max_vadv)
+
+    if collisions.krook_collision_frequency_prefactor_ee > 0.0
+        # Add a Krook collision operator
+        # Set dt=-1 as we update the residual here rather than adding an update to
+        # 'fvec_out'.
+        electron_krook_collisions!(residual, pdf, dens, upar, upar_ion, vth,
+                                   collisions, vperp, vpa, -1.0)
+    end
+
     dt_max = dt_electron
     #println("dt_max: ", dt_max, " dt_max_zadv: ", dt_max_zadv, " dt_max_vadv: ", dt_max_vadv)
     return dt_max
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 081a7d744..c54dbf10f 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -60,7 +60,7 @@ mutable struct advance_info
     ionization_collisions::Bool
     ionization_collisions_1V::Bool
     ionization_source::Bool
-    krook_collisions::Bool
+    krook_collisions_ii::Bool
     explicit_weakform_fp_collisions::Bool
     external_source::Bool
     numerical_dissipation::Bool
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 248b0038f..aa2b20f07 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -53,10 +53,10 @@ include("electron_vpa_advection.jl")
 include("neutral_r_advection.jl")
 include("neutral_z_advection.jl")
 include("neutral_vz_advection.jl")
-include("electron_kinetic_equation.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("krook_collisions.jl")
+include("electron_kinetic_equation.jl")
 include("continuity.jl")
 include("energy_equation.jl")
 include("force_balance.jl")
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 107f27009..43a98a4cc 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -560,7 +560,7 @@ function setup_advance_flags(moments, composition, t_input, collisions,
     advance_ionization = false
     advance_ionization_1V = false
     advance_ionization_source = false
-    advance_krook_collisions = false
+    advance_krook_collisions_ii = false
     advance_external_source = false
     advance_numerical_dissipation = false
     advance_sources = false
@@ -637,8 +637,8 @@ function setup_advance_flags(moments, composition, t_input, collisions,
         if collisions.ionization > 0.0 && collisions.constant_ionization_rate
             advance_ionization_source = true
         end
-        if collisions.krook_collision_frequency_prefactor > 0.0
-            advance_krook_collisions = true
+        if collisions.krook_collision_frequency_prefactor_ii > 0.0
+            advance_krook_collisions_ii = true
         end
         advance_external_source = external_source_settings.ion.active
         advance_neutral_external_source = external_source_settings.neutral.active
@@ -695,7 +695,7 @@ function setup_advance_flags(moments, composition, t_input, collisions,
                         advance_neutral_z_advection, advance_neutral_r_advection,
                         advance_neutral_vz_advection, advance_cx, advance_cx_1V,
                         advance_ionization, advance_ionization_1V,
-                        advance_ionization_source, advance_krook_collisions,
+                        advance_ionization_source, advance_krook_collisions_ii,
                         explicit_weakform_fp_collisions,
                         advance_external_source, advance_numerical_dissipation,
                         advance_sources, advance_continuity, advance_force_balance,
@@ -1148,12 +1148,12 @@ function time_advance_split_operators!(pdf, scratch, t, t_input, vpa, z,
             end
         end
         if collisions.krook_collision_frequency_prefactor  > 0.0
-            advance.krook_collisions = true
+            advance.krook_collisions_ii = true
             time_advance_no_splitting!(pdf, scratch, t, t_input, z, vpa,
                 z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
                 z_SL, vpa_SL, composition, collisions, sources, num_diss_params,
                 advance, istep)
-            advance.krook_collisions = false
+            advance.krook_collisions_ii = false
         end
         # and add the source terms associated with redefining g = pdf/density or pdf*vth/density
         # to the kinetic equation
@@ -1846,7 +1846,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
 
     # Add Krook collision operator for ions
-    if advance.krook_collisions
+    if advance.krook_collisions_ii
         krook_collisions!(fvec_out.pdf, fvec_in, moments, composition, collisions,
                           vperp, vpa, dt)
     end

From 8a9e9a6827a40a8867752fa169544a0a054401a5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 24 Feb 2024 21:07:12 +0000
Subject: [PATCH 129/428] Use get_variable() in VariableCache

Allows more makie_post_processing functions to work with variables
constructed by get_variable(), e.g. "collision_frequency_ii", etc.
---
 .../makie_post_processing/src/makie_post_processing.jl        | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 978d8e839..06eb32965 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -806,8 +806,8 @@ function VariableCache(run_info, variable_name::String, t_chunk_size::mk_int;
     tinds_chunk = 1:t_chunk_size
     dim_slices = (is=is, iz=iz, ir=ir, ivperp=ivperp, ivpa=ivpa, ivzeta=ivzeta, ivr=ivr,
                   ivz=ivz)
-    data_chunk = postproc_load_variable(run_info, variable_name;
-                                        it=tinds_range_global[tinds_chunk], dim_slices...)
+    data_chunk = get_variable(run_info, variable_name; it=tinds_range_global[tinds_chunk],
+                              dim_slices...)
 
     return VariableCache(run_info, variable_name, t_chunk_size,
                          n_tinds, tinds_range_global, Ref(tinds_chunk),

From ca98977b67cd031058cbdb1471124dcbd849da43 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 24 Feb 2024 21:08:45 +0000
Subject: [PATCH 130/428] Support selection by `it` for the time variable in
 title of animations

---
 .../src/makie_post_processing.jl                   | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 06eb32965..d36f25506 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1735,9 +1735,11 @@ for dim ∈ one_dimension_combinations_no_t
                              all(isapprox.(ri.time, run_info[1].time))
                              for ri ∈ run_info[2:end])
                          # All times are the same
-                         title = lift(i->string("t = ", run_info[1].time[i]), frame_index)
+                         time = select_slice(run_info[1].time, :t; input=input, it=it)
+                         title = lift(i->string("t = ", time[i]), frame_index)
                      else
-                         title = lift(i->join((string("t", irun, " = ", ri.time[i])
+                         time = select_slice(ri.time, :t; input=input, it=it)
+                         title = lift(i->join((string("t", irun, " = ", time[i])
                                                for (irun,ri) ∈ enumerate(run_info)), "; "),
                                       frame_index)
                      end
@@ -1817,7 +1819,8 @@ for dim ∈ one_dimension_combinations_no_t
                      ind = frame_index
                  end
                  if ax === nothing
-                     title = lift(i->string("t = ", run_info.time[i]), ind)
+                     time = select_slice(run_info.time, :t; input=input, it=it)
+                     title = lift(i->string("t = ", time[i]), ind)
                      fig, ax = get_1d_ax(; xlabel="$($dim_str)",
                                          ylabel=get_variable_symbol(var_name),
                                          yscale=yscale, title=title, axis_args...)
@@ -1962,10 +1965,12 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
 
                      if length(run_info) > 1
                          title = get_variable_symbol(var_name)
-                         subtitles = (lift(i->string(ri.run_name, "\nt = ", ri.time[i]),
+                         time = select_slice(ri.time, :t; input=input, it=it)
+                         subtitles = (lift(i->string(ri.run_name, "\nt = ", time[i]),
                                            frame_index)
                                       for ri ∈ run_info)
                      else
+                         time = select_slice(run_info[1].time, :t; input=input, it=it)
                          title = lift(i->string(get_variable_symbol(var_name), "\nt = ",
                                                 run_info[1].time[i]),
                                       frame_index)
@@ -2040,6 +2045,7 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
                      colormap = input.colormap
                  end
                  if title === nothing && ax == nothing
+                     time = select_slice(run_info.time, :t; input=input, it=it)
                      title = lift(i->string(get_variable_symbol(var_name), "\nt = ",
                                             run_info.time[i]),
                                   ind)

From a73c3dd1089d399770c22ab9c0da785960512700 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 25 Feb 2024 15:06:39 +0000
Subject: [PATCH 131/428] Support electron pdf in
 enforce_vperp_boundary_condition!()

---
 moment_kinetics/src/initial_conditions.jl | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index acb2fb7af..9f7736c92 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -2507,20 +2507,29 @@ end
 """
 enforce zero boundary condition at vperp -> infinity
 """
-function enforce_vperp_boundary_condition!(f, bc, vperp, vperp_spectral)
+function enforce_vperp_boundary_condition! end
+
+function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vperp, vperp_spectral)
+    @loop_s is begin
+        @views enforce_vperp_boundary_condition!(f[:,:,:,:,is], bc, vperp, vperp_spectral)
+    end
+    return nothing
+end
+
+function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral)
     if bc == "zero"
         nvperp = vperp.n
         ngrid = vperp.ngrid
         # set zero boundary condition
-        @loop_s_r_z_vpa is ir iz ivpa begin
-            f[ivpa,nvperp,iz,ir,is] = 0.0
+        @loop_r_z_vpa ir iz ivpa begin
+            f[ivpa,nvperp,iz,ir] = 0.0
         end
         # set regularity condition d F / d vperp = 0 at vperp = 0
         if vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral"
             D0 = vperp_spectral.radau.D0
-            @loop_s_r_z_vpa is ir iz ivpa begin
+            @loop_r_z_vpa ir iz ivpa begin
                 # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
-                f[ivpa,1,iz,ir,is] = -sum(D0[2:ngrid].*f[ivpa,2:ngrid,iz,ir,is])/D0[1]
+                f[ivpa,1,iz,ir] = -sum(D0[2:ngrid].*f[ivpa,2:ngrid,iz,ir])/D0[1]
             end
         else
             println("vperp.bc=\"$bc\" not supported by discretization "

From 972285d8894de45ed9b6787d0f775f73c65074c9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 25 Feb 2024 15:16:13 +0000
Subject: [PATCH 132/428] Move boundary conditions functions to separate module

Makes sense to separate initial conditions and boundary conditions
functions.

Also avoids cyclic dependency between `initial_conditions` and
`electron_kinetic_equation` modules.
---
 moment_kinetics/src/analysis.jl            |    2 +-
 moment_kinetics/src/boundary_conditions.jl | 1000 ++++++++++++++++++++
 moment_kinetics/src/initial_conditions.jl  |  986 +------------------
 moment_kinetics/src/moment_kinetics.jl     |    3 +-
 moment_kinetics/src/time_advance.jl        |    4 +-
 5 files changed, 1006 insertions(+), 989 deletions(-)
 create mode 100644 moment_kinetics/src/boundary_conditions.jl

diff --git a/moment_kinetics/src/analysis.jl b/moment_kinetics/src/analysis.jl
index 0f43a7177..31d52a49d 100644
--- a/moment_kinetics/src/analysis.jl
+++ b/moment_kinetics/src/analysis.jl
@@ -10,7 +10,7 @@ using ..array_allocation: allocate_float
 using ..calculus: integral
 using ..communication
 using ..coordinates: coordinate
-using ..initial_conditions: vpagrid_to_dzdt
+using ..boundary_conditions: vpagrid_to_dzdt
 using ..interpolation: interpolate_to_grid_1d
 using ..load_data: open_readonly_output_file, get_nranks, load_pdf_data, load_rank_data
 using ..load_data: load_distributed_ion_pdf_slice
diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
new file mode 100644
index 000000000..d98cd22eb
--- /dev/null
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -0,0 +1,1000 @@
+"""
+Functions for applying boundary conditions
+"""
+module boundary_conditions
+
+export enforce_boundary_conditions!
+export enforce_neutral_boundary_conditions!
+
+using SpecialFunctions: erfc
+
+using ..calculus: reconcile_element_boundaries_MPI!
+using ..coordinates: coordinate
+using ..interpolation: interpolate_to_grid_1d!
+using ..looping
+using ..moment_kinetics_structs: scratch_pdf
+using ..type_definitions: mk_float, mk_int
+using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace,
+                          integrate_over_positive_vz, integrate_over_negative_vz
+
+"""
+enforce boundary conditions in vpa and z on the evolved pdf;
+also enforce boundary conditions in z on all separately evolved velocity space moments of the pdf
+"""
+function enforce_boundary_conditions!(f, f_r_bc, density, upar, ppar, moments, vpa_bc,
+        z_bc, r_bc, vpa, vperp, z, r, vpa_spectral, vperp_spectral, vpa_adv, z_adv, r_adv,
+        composition, scratch_dummy, r_diffusion, vpa_diffusion)
+    if vpa.n > 1
+        begin_s_r_z_vperp_region()
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            # enforce the vpa BC
+            # use that adv.speed independent of vpa 
+            @views enforce_v_boundary_condition_local!(f[:,ivperp,iz,ir,is], vpa_bc,
+                             vpa_adv[is].speed[:,ivperp,iz,ir], vpa_diffusion,
+                             vpa, vpa_spectral)
+        end
+    end
+    if vperp.n > 1
+        begin_s_r_z_vpa_region()
+        @views enforce_vperp_boundary_condition!(f, vperp.bc, vperp, vperp_spectral)
+    end
+    if z.n > 1
+        begin_s_r_vperp_vpa_region()
+        # enforce the z BC on the evolved velocity space moments of the pdf
+        @views enforce_z_boundary_condition_moments!(density, moments, z_bc)
+        @views enforce_z_boundary_condition!(f, density, upar, ppar, moments, z_bc, z_adv, z,
+                                             vperp, vpa, composition,
+                                             scratch_dummy.buffer_vpavperprs_1, scratch_dummy.buffer_vpavperprs_2,
+                                             scratch_dummy.buffer_vpavperprs_3, scratch_dummy.buffer_vpavperprs_4)
+                                              
+    end
+    if r.n > 1
+        begin_s_z_vperp_vpa_region()
+        @views enforce_r_boundary_condition!(f, f_r_bc, r_bc, r_adv, vpa, vperp, z, r, composition,
+            scratch_dummy.buffer_vpavperpzs_1, scratch_dummy.buffer_vpavperpzs_2,
+            scratch_dummy.buffer_vpavperpzs_3, scratch_dummy.buffer_vpavperpzs_4,
+            r_diffusion)
+    end
+end
+function enforce_boundary_conditions!(fvec_out::scratch_pdf, moments, f_r_bc, vpa_bc,
+        z_bc, r_bc, vpa, vperp, z, r, vpa_spectral, vperp_spectral, vpa_adv, z_adv, r_adv, composition, scratch_dummy,
+        r_diffusion, vpa_diffusion)
+    enforce_boundary_conditions!(fvec_out.pdf, f_r_bc, fvec_out.density, fvec_out.upar,
+        fvec_out.ppar, moments, vpa_bc, z_bc, r_bc, vpa, vperp, z, r, 
+        vpa_spectral, vperp_spectral, vpa_adv, z_adv,
+        r_adv, composition, scratch_dummy, r_diffusion, vpa_diffusion)
+end
+
+"""
+enforce boundary conditions on f in r
+"""
+function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc::String,
+        adv, vpa, vperp, z, r, composition, end1::AbstractArray{mk_float,4},
+        end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4},
+        buffer2::AbstractArray{mk_float,4}, r_diffusion::Bool)
+
+    nr = r.n
+
+    if r.nelement_global > r.nelement_local
+        # reconcile internal element boundaries across processes
+        # & enforce periodicity and external boundaries if needed
+        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
+            end1[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,1,is]
+            end2[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,nr,is]
+        end
+        @views reconcile_element_boundaries_MPI!(f,
+            end1, end2,	buffer1, buffer2, r)
+    end
+
+    # 'periodic' BC enforces periodicity by taking the average of the boundary points
+    # enforce the condition if r is local
+    if bc == "periodic" && r.nelement_global == r.nelement_local
+        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
+            f[ivpa,ivperp,iz,1,is] = 0.5*(f[ivpa,ivperp,iz,nr,is]+f[ivpa,ivperp,iz,1,is])
+            f[ivpa,ivperp,iz,nr,is] = f[ivpa,ivperp,iz,1,is]
+        end
+    end
+    if bc == "Dirichlet"
+        zero = 1.0e-10
+        # use the old distribution to force the new distribution to have
+        # consistant-in-time values at the boundary
+        # with bc = "Dirichlet" and r_diffusion = false
+        # impose bc for incoming parts of velocity space only (Hyperbolic PDE)
+        # with bc = "Dirichlet" and r_diffusion = true
+        # impose bc on both sides of the domain to accomodate a diffusion operator d^2 / d r^2
+        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
+            ir = 1 # r = -L/2 -- check that the point is on lowest rank
+            if r.irank == 0 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] > zero)
+                f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,1,is]
+            end
+            ir = r.n # r = L/2 -- check that the point is on highest rank
+            if r.irank == r.nrank - 1 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] < -zero)
+                f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,end,is]
+            end
+        end
+    end
+end
+
+"""
+enforce boundary conditions on ion particle f in z
+"""
+function enforce_z_boundary_condition!(pdf, density, upar, ppar, moments, bc::String, adv,
+                                       z, vperp, vpa, composition, end1::AbstractArray{mk_float,4},
+                                       end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4},
+                                       buffer2::AbstractArray{mk_float,4})
+    # this block ensures periodic BC can be supported with distributed memory MPI
+    if z.nelement_global > z.nelement_local
+        # reconcile internal element boundaries across processes
+        # & enforce periodicity and external boundaries if needed
+        nz = z.n
+        @loop_s_r_vperp_vpa is ir ivperp ivpa begin
+            end1[ivpa,ivperp,ir,is] = pdf[ivpa,ivperp,1,ir,is]
+            end2[ivpa,ivperp,ir,is] = pdf[ivpa,ivperp,nz,ir,is]
+        end
+        # check on periodic bc happens inside this call below
+        @views reconcile_element_boundaries_MPI!(pdf,
+            end1, end2,	buffer1, buffer2, z)
+    end
+    # define a zero that accounts for finite precision
+    zero = 1.0e-14
+    # 'constant' BC is time-independent f at upwind boundary
+    # and constant f beyond boundary 
+    if bc == "constant"
+        begin_s_r_vperp_vpa_region()
+        density_offset = 1.0
+        vwidth = 1.0
+        if z.irank == 0
+            @loop_s_r_vperp_vpa is ir ivperp ivpa begin
+                if adv[is].speed[ivpa,1,ir] > 0.0
+                    pdf[ivpa,ivperp,1,ir,is] = density_offset * exp(-(vpa.grid[ivpa]^2 + vperp.grid[ivperp]^2)/vwidth^2) / sqrt(pi)
+                end
+            end
+        end
+        if z.irank == z.nrank - 1
+            @loop_s_r_vperp_vpa is ir ivperp ivpa begin
+                if adv[is].speed[ivpa,end,ir] > 0.0
+                    pdf[ivpa,ivperp,end,ir,is] = density_offset * exp(-(vpa.grid[ivpa]^2 + vperp.grid[ivperp]^2)/vwidth^2) / sqrt(pi)
+                end
+            end
+        end
+    # 'periodic' BC enforces periodicity by taking the average of the boundary points
+    elseif bc == "periodic" && z.nelement_global == z.nelement_local
+        begin_s_r_vperp_vpa_region()
+        @loop_s_r_vperp_vpa is ir ivperp ivpa begin
+            pdf[ivpa,ivperp,1,ir,is] = 0.5*(pdf[ivpa,ivperp,z.n,ir,is]+pdf[ivpa,ivperp,1,ir,is])
+            pdf[ivpa,ivperp,z.n,ir,is] = pdf[ivpa,ivperp,1,ir,is]
+        end
+    # 'wall' BC enforces wall boundary conditions
+    elseif bc == "wall"
+        # Need integrals over vpa at wall boundaries in z, so cannot parallelize over z
+        # or vpa.
+        begin_s_r_region()
+        @loop_s is begin
+            # zero incoming BC for ions, as they recombine at the wall
+            if moments.evolve_upar
+                @loop_r ir begin
+                    @views enforce_zero_incoming_bc!(
+                        pdf[:,:,:,ir,is], z, vpa, density[:,ir,is], upar[:,ir,is],
+                        ppar[:,ir,is], moments.evolve_upar, moments.evolve_ppar, zero)
+                end
+            else
+                @loop_r ir begin
+                    @views enforce_zero_incoming_bc!(pdf[:,:,:,ir,is],
+                                                     adv[is].speed[:,:,:,ir], z, zero)
+                end
+            end
+        end
+    end
+end
+
+"""
+enforce boundary conditions on neutral particle distribution function
+"""
+function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
+        boundary_distributions, density_neutral, uz_neutral, pz_neutral, moments,
+        density_ion, upar_ion, Er, vzeta_spectral, vr_spectral, vz_spectral, r_adv, z_adv,
+        vzeta_adv, vr_adv, vz_adv, r, z, vzeta, vr, vz, composition, geometry,
+        scratch_dummy, r_diffusion, vz_diffusion)
+
+    # without acceleration of neutrals bc on vz vr vzeta should not be required as no
+    # advection or diffusion in these coordinates
+
+    if vzeta_adv !== nothing && vzeta.n_global > 1 && vzeta.bc != "none"
+        begin_sn_r_z_vr_vz_region()
+        @loop_sn_r_z_vr_vz isn ir iz ivr ivz begin
+            # enforce the vz BC
+            @views enforce_v_boundary_condition_local!(f_neutral[ivz,ivr,:,iz,ir,isn],
+                                                       vzeta.bc,
+                                                       vzeta_adv[isn].speed[ivz,ivr,:,iz,ir],
+                                                       false, vzeta, vzeta_spectral)
+        end
+    end
+    if vr_adv !== nothing && vr.n_global > 1 && vr.bc != "none"
+        begin_sn_r_z_vzeta_vz_region()
+        @loop_sn_r_z_vzeta_vz isn ir iz ivzeta ivz begin
+            # enforce the vz BC
+            @views enforce_v_boundary_condition_local!(f_neutral[ivz,:,ivzeta,iz,ir,isn],
+                                                       vr.bc,
+                                                       vr_adv[isn].speed[ivz,:,ivzeta,iz,ir],
+                                                       false, vr, vr_spectral)
+        end
+    end
+    if vz_adv !== nothing && vz.n_global > 1 && vz.bc != "none"
+        begin_sn_r_z_vzeta_vr_region()
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            # enforce the vz BC
+            @views enforce_v_boundary_condition_local!(f_neutral[:,ivr,ivzeta,iz,ir,isn],
+                                                       vz.bc,
+                                                       vz_adv[isn].speed[:,ivr,ivzeta,iz,ir],
+                                                       vz_diffusion, vz, vz_spectral)
+        end
+    end
+    # f_initial contains the initial condition for enforcing a fixed-boundary-value condition
+    if z.n > 1
+        begin_sn_r_vzeta_vr_vz_region()
+        @views enforce_neutral_z_boundary_condition!(f_neutral, density_neutral, uz_neutral,
+            pz_neutral, moments, density_ion, upar_ion, Er, boundary_distributions,
+            z_adv, z, vzeta, vr, vz, composition, geometry, 
+            scratch_dummy.buffer_vzvrvzetarsn_1, scratch_dummy.buffer_vzvrvzetarsn_2,
+            scratch_dummy.buffer_vzvrvzetarsn_3, scratch_dummy.buffer_vzvrvzetarsn_4)
+    end
+    if r.n > 1
+        begin_sn_z_vzeta_vr_vz_region()
+        @views enforce_neutral_r_boundary_condition!(f_neutral, boundary_distributions.pdf_rboundary_neutral,
+                                    r_adv, vz, vr, vzeta, z, r, composition,
+                                    scratch_dummy.buffer_vzvrvzetazsn_1, scratch_dummy.buffer_vzvrvzetazsn_2,
+                                    scratch_dummy.buffer_vzvrvzetazsn_3, scratch_dummy.buffer_vzvrvzetazsn_4,
+                                    r_diffusion)
+    end
+end
+
+function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6},
+        f_r_bc::AbstractArray{mk_float,6}, adv, vz, vr, vzeta, z, r, composition,
+        end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5},
+        buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5},
+        r_diffusion) #f_initial,
+
+    bc = r.bc
+    nr = r.n
+
+    if r.nelement_global > r.nelement_local
+        # reconcile internal element boundaries across processes
+        # & enforce periodicity and external boundaries if needed
+        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
+            end1[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,1,isn]
+            end2[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,nr,isn]
+        end
+        @views reconcile_element_boundaries_MPI!(f,
+            end1, end2,	buffer1, buffer2, r)
+    end
+    # 'periodic' BC enforces periodicity by taking the average of the boundary points
+    # local case only when no communication required
+    if bc == "periodic" && r.nelement_global == r.nelement_local
+        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
+            f[ivz,ivr,ivzeta,iz,1,isn] = 0.5*(f[ivz,ivr,ivzeta,iz,1,isn]+f[ivz,ivr,ivzeta,iz,nr,isn])
+            f[ivz,ivr,ivzeta,iz,nr,isn] = f[ivz,ivr,ivzeta,iz,1,isn]
+        end
+    end
+    # Dirichlet boundary condition for external endpoints
+    if bc == "Dirichlet"
+        zero = 1.0e-10
+        # use the old distribution to force the new distribution to have
+        # consistant-in-time values at the boundary
+        # impose bc for incoming parts of velocity space only (Hyperbolic PDE)
+        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
+            ir = 1 # r = -L/2
+            # incoming particles and on lowest rank
+            if r.irank == 0 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] > zero)
+                f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,1,isn]
+            end
+            ir = nr # r = L/2
+            # incoming particles and on highest rank
+            if r.irank == r.nrank - 1 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] < -zero)
+                f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,end,isn]
+            end
+        end
+    end
+end
+
+"""
+enforce boundary conditions on neutral particle f in z
+"""
+function enforce_neutral_z_boundary_condition!(pdf, density, uz, pz, moments, density_ion,
+                                               upar_ion, Er, boundary_distributions, adv,
+                                               z, vzeta, vr, vz, composition, geometry,
+                                               end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5},
+                                               buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5})
+    
+
+    if z.nelement_global > z.nelement_local
+        # reconcile internal element boundaries across processes
+        # & enforce periodicity and external boundaries if needed
+        nz = z.n
+        @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+            end1[ivz,ivr,ivzeta,ir,isn] = pdf[ivz,ivr,ivzeta,1,ir,isn]
+            end2[ivz,ivr,ivzeta,ir,isn] = pdf[ivz,ivr,ivzeta,nz,ir,isn]
+        end
+        # check on periodic bc occurs within this call below
+        @views reconcile_element_boundaries_MPI!(pdf,
+            end1, end2,	buffer1, buffer2, z)
+    end
+
+    zero = 1.0e-14
+    # 'constant' BC is time-independent f at upwind boundary
+    # and constant f beyond boundary
+    if z.bc == "constant"
+        begin_sn_r_vzeta_vr_vz_region()
+        density_offset = 1.0
+        vwidth = 1.0
+        if z.irank == 0
+            @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+                if adv[isn].speed[ivz,ivr,ivzeta,1,ir] > 0.0
+                    pdf[ivz,ivr,ivzeta,1,ir,is] = density_offset *
+                        exp(-(vzeta.grid[ivzeta]^2 + vr.grid[ivr] + vz.grid[ivz])/vwidth^2) /
+                        sqrt(pi)
+                end
+            end
+        end
+        if z.irank == z.nrank - 1
+            @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+                if adv[isn].speed[ivz,ivr,ivzeta,end,ir] > 0.0
+                    pdf[ivz,ivr,ivzeta,end,ir,is] = density_offset *
+                        exp(-(vzeta.grid[ivzeta]^2 + vr.grid[ivr] + vz.grid[ivz])/vwidth^2) /
+                        sqrt(pi)
+                end
+            end
+        end
+    # 'periodic' BC enforces periodicity by taking the average of the boundary points
+    elseif z.bc == "periodic" && z.nelement_global == z.nelement_local
+        begin_sn_r_vzeta_vr_vz_region()
+        @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
+            pdf[ivz,ivr,ivzeta,1,ir,isn] = 0.5*(pdf[ivz,ivr,ivzeta,1,ir,isn] +
+                                                pdf[ivz,ivr,ivzeta,end,ir,isn])
+            pdf[ivz,ivr,ivzeta,end,ir,isn] = pdf[ivz,ivr,ivzeta,1,ir,isn]
+        end
+    # 'wall' BC enforces wall boundary conditions
+    elseif z.bc == "wall"
+        # Need integrals over vpa at wall boundaries in z, so cannot parallelize over z
+        # or vpa.
+        begin_sn_r_region()
+        @loop_sn isn begin
+            # BC for neutrals
+            @loop_r ir begin
+                # define vtfac to avoid repeated computation below
+                vtfac = sqrt(composition.T_wall * composition.mn_over_mi)
+                # Assume for now that the ion species index corresponding to this neutral
+                # species is the same as the neutral species index.
+                # Note, have already calculated moments of ion distribution function(s),
+                # so can use the moments here to get the flux
+                if z.irank == 0
+                    ion_flux_0 = -density_ion[1,ir,isn] * (upar_ion[1,ir,isn]*geometry.bzed - 0.5*geometry.rhostar*Er[1,ir])
+                else
+                    ion_flux_0 = NaN
+                end
+                if z.irank == z.nrank - 1
+                    ion_flux_L = density_ion[end,ir,isn] * (upar_ion[end,ir,isn]*geometry.bzed - 0.5*geometry.rhostar*Er[end,ir])
+                else
+                    ion_flux_L = NaN
+                end
+                # enforce boundary condition on the neutral pdf that all ions and neutrals
+                # that leave the domain re-enter as neutrals
+                @views enforce_neutral_wall_bc!(
+                    pdf[:,:,:,:,ir,isn], z, vzeta, vr, vz, pz[:,ir,isn], uz[:,ir,isn],
+                    density[:,ir,isn], ion_flux_0, ion_flux_L, boundary_distributions,
+                    vtfac, composition.recycling_fraction, moments.evolve_ppar,
+                    moments.evolve_upar, moments.evolve_density, zero)
+            end
+        end
+    end
+end
+
+"""
+enforce a zero incoming BC in z for given species pdf at each radial location
+"""
+function enforce_zero_incoming_bc!(pdf, speed, z, zero)
+    nvpa = size(pdf,1)
+    # no parallel BC should be enforced for dz/dt = 0
+    # note that the parallel velocity coordinate vpa may be dz/dt or
+    # some version of the peculiar velocity (dz/dt - upar),
+    # so use advection speed below instead of vpa
+    if z.irank == 0
+        @loop_vperp_vpa ivperp ivpa begin
+            # for left boundary in zed (z = -Lz/2), want
+            # f(z=-Lz/2, v_parallel > 0) = 0
+            if speed[1,ivpa,ivperp] > zero
+                pdf[ivpa,ivperp,1] = 0.0
+            end
+        end
+    end
+    if z.irank == z.nrank - 1
+        @loop_vperp_vpa ivperp ivpa begin
+            # for right boundary in zed (z = Lz/2), want
+            # f(z=Lz/2, v_parallel < 0) = 0
+            if speed[end,ivpa,ivperp] < -zero
+                pdf[ivpa,ivperp,end] = 0.0
+            end
+        end
+    end
+end
+function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density, upar,
+                                   ppar, evolve_upar, evolve_ppar, zero)
+    if z.irank != 0 && z.irank != z.nrank - 1
+        # No z-boundary in this block
+        return nothing
+    end
+    nvpa, nvperp, nz = size(pdf)
+    # no parallel BC should be enforced for dz/dt = 0
+    # note that the parallel velocity coordinate vpa may be dz/dt or
+    # some version of the peculiar velocity (dz/dt - upar),
+    # so use advection speed below instead of vpa
+
+    # absolute velocity at left boundary
+    if z.irank == 0
+        @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[1]/density[1])),
+                                         upar[1], evolve_ppar, evolve_upar)
+        @loop_vpa ivpa begin
+            # for left boundary in zed (z = -Lz/2), want
+            # f(z=-Lz/2, v_parallel > 0) = 0
+            if vpa.scratch[ivpa] > zero
+                pdf[ivpa,:,1] .= 0.0
+            end
+        end
+    end
+    # absolute velocity at right boundary
+    if z.irank == z.nrank - 1
+        @. vpa.scratch2 = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[end]/density[end])),
+                                          upar[end], evolve_ppar, evolve_upar)
+        @loop_vpa ivpa begin
+            # for right boundary in zed (z = Lz/2), want
+            # f(z=Lz/2, v_parallel < 0) = 0
+            if vpa.scratch2[ivpa] < -zero
+                pdf[ivpa,:,end] .= 0.0
+            end
+        end
+    end
+
+    # Special constraint-forcing code that tries to keep the modifications smooth at
+    # v_parallel=0.
+    if z.irank == 0 && z.irank == z.nrank - 1
+        # Both z-boundaries in this block
+        z_range = (1,nz)
+    elseif z.irank == 0
+        z_range = (1,)
+    elseif z.irank == z.nrank - 1
+        z_range = (nz,)
+    else
+        error("No boundary in this block, should have returned already")
+    end
+    for iz ∈ z_range
+        # moment-kinetic approach only implemented for 1V case so far
+        @boundscheck size(pdf,2) == 1
+
+        f = @view pdf[:,1,iz]
+        if evolve_ppar && evolve_upar
+            I0 = integrate_over_vspace(f, vpa.wgts)
+            I1 = integrate_over_vspace(f, vpa.grid, vpa.wgts)
+            I2 = integrate_over_vspace(f, vpa.grid, 2, vpa.wgts)
+
+            # Store v_parallel with upar shift removed in vpa.scratch
+            vth = sqrt(2.0*ppar[iz]/density[iz])
+            @. vpa.scratch = vpa.grid + upar[iz]/vth
+            # Introduce factor to ensure corrections go smoothly to zero near
+            # v_parallel=0
+            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
+            J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
+            J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
+            J3 = integrate_over_vspace(vpa.scratch2, vpa.grid, 3, vpa.wgts)
+            J4 = integrate_over_vspace(vpa.scratch2, vpa.grid, 4, vpa.wgts)
+
+            A = (J3^2 - J2*J4 + 0.5*(J2^2 - J1*J3)) /
+                (I0*(J3^2 - J2*J4) + I1*(J1*J4 - J2*J3) + I2*(J2^2 - J1*J3))
+            B = (0.5*J3 + A*(I1*J4 - I2*J3)) / (J3^2 - J2*J4)
+            C = (0.5 - A*I2 -B*J3) / J4
+
+            @. f = A*f + B*vpa.grid*vpa.scratch2 + C*vpa.grid*vpa.grid*vpa.scratch2
+        elseif evolve_upar
+            I0 = integrate_over_vspace(f, vpa.wgts)
+            I1 = integrate_over_vspace(f, vpa.grid, vpa.wgts)
+
+            # Store v_parallel with upar shift removed in vpa.scratch
+            @. vpa.scratch = vpa.grid + upar[iz]
+            # Introduce factor to ensure corrections go smoothly to zero near
+            # v_parallel=0
+            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
+            J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
+            J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
+
+            A = 1.0 / (I0 - I1*J1/J2)
+            B = -A*I1/J2
+
+            @. f = A*f + B*vpa.grid*vpa.scratch2
+        elseif evolve_density
+            I0 = integrate_over_vspace(f, vpa.wgts)
+            @. f = f / I0
+        end
+    end
+end
+
+"""
+Set up an initial condition that tries to be smoothly compatible with the sheath
+boundary condition for ions, by setting f(±(v_parallel-u0)<0) where u0=0 at the sheath
+boundaries and for z<0 increases linearly to u0=vpa.L at z=0, while for z>0 increases
+from u0=-vpa.L at z=0 to zero at the z=z.L/2 sheath.
+
+To be applied to 'full-f' distribution function on v_parallel grid (not w_parallel
+grid).
+"""
+function enforce_initial_tapered_zero_incoming!(pdf, z::coordinate, vpa::coordinate)
+    nvpa = size(pdf,1)
+    zero = 1.0e-14
+    # no parallel BC should be enforced for dz/dt = 0
+    # note that the parallel velocity coordinate vpa may be dz/dt or
+    # some version of the peculiar velocity (dz/dt - upar),
+    # so use advection speed below instead of vpa
+
+    for iz ∈ 1:z.n
+        u0 = (2.0*z.grid[iz]/z.L - sign(z.grid[iz])) * vpa.L / 2.0
+        if z.grid[iz] < -zero
+            for ivpa ∈ 1:nvpa
+                if vpa.grid[ivpa] > u0 + zero
+                    pdf[ivpa,iz] = 0.0
+                end
+            end
+        elseif z.grid[iz] > zero
+            for ivpa ∈ 1:nvpa
+                if vpa.grid[ivpa] < u0 - zero
+                    pdf[ivpa,iz] = 0.0
+                end
+            end
+        end
+    end
+end
+
+"""
+enforce the wall boundary condition on neutrals;
+i.e., the incoming flux of neutrals equals the sum of the ion/neutral outgoing fluxes
+"""
+function enforce_neutral_wall_bc!(pdf, z, vzeta, vr, vz, pz, uz, density, wall_flux_0,
+                                  wall_flux_L, boundary_distributions, vtfac,
+                                  recycling_fraction, evolve_ppar, evolve_upar,
+                                  evolve_density, zero)
+
+    # Reduce the ion flux by `recycling_fraction` to account for ions absorbed by the
+    # wall.
+    wall_flux_0 *= recycling_fraction
+    wall_flux_L *= recycling_fraction
+
+    if !evolve_density && !evolve_upar
+        knudsen_cosine = boundary_distributions.knudsen
+
+        if z.irank == 0
+            ## treat z = -Lz/2 boundary ##
+
+            # add the neutral species's contribution to the combined ion/neutral particle
+            # flux out of the domain at z=-Lz/2
+            @views wall_flux_0 += integrate_over_negative_vz(abs.(vz.grid) .* pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+
+            # for left boundary in zed (z = -Lz/2), want
+            # f_n(z=-Lz/2, v_parallel > 0) = Γ_0 * f_KW(v_parallel)
+            @loop_vz ivz begin
+                if vz.grid[ivz] >= -zero
+                    @views @. pdf[ivz,:,:,1] = wall_flux_0 * knudsen_cosine[ivz,:,:]
+                end
+            end
+        end
+
+        if z.irank == z.nrank - 1
+            ## treat the right boundary at z = Lz/2 ##
+
+            # add the neutral species's contribution to the combined ion/neutral particle
+            # flux out of the domain at z=-Lz/2
+            @views wall_flux_L += integrate_over_positive_vz(abs.(vz.grid) .* pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+
+            # for right boundary in zed (z = Lz/2), want
+            # f_n(z=Lz/2, v_parallel < 0) = Γ_Lz * f_KW(v_parallel)
+            @loop_vz ivz begin
+                if vz.grid[ivz] <= zero
+                    @views @. pdf[ivz,:,:,end] = wall_flux_L * knudsen_cosine[ivz,:,:]
+                end
+            end
+        end
+    elseif !evolve_upar
+        # Evolving density case
+        knudsen_cosine = boundary_distributions.knudsen
+
+        if z.irank == 0
+            ## treat z = -Lz/2 boundary ##
+
+            # Note the numerical integrol of knudsen_cosine was forced to be 1 (to machine
+            # precision) when it was initialised.
+            @views pdf_integral_0 = integrate_over_negative_vz(pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            @views pdf_integral_1 = integrate_over_negative_vz(vz.grid .* pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_0 = integrate_over_positive_vz(knudsen_cosine, vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_1 = 1.0 # This is enforced in initialization
+
+            # Calculate normalisation factors N_in for the incoming and N_out for the
+            # Knudsen parts of the distirbution so that ∫dvpa F = 1 and ∫dvpa vpa F = uz
+            # Note wall_flux_0 is the ion flux into the wall (reduced by the recycling
+            # fraction), and the neutral flux should be out of the wall (i.e. uz>0) so
+            # n*uz = |n*uz| = wall_flux_0
+            # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
+            #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = uz
+            uz = wall_flux_0 / density[1]
+            N_in = (1 - uz * knudsen_integral_0 / knudsen_integral_1) /
+                   (pdf_integral_0
+                    - pdf_integral_1 / knudsen_integral_1 * knudsen_integral_0)
+            N_out = (uz - N_in * pdf_integral_1) / knudsen_integral_1
+
+            @loop_vz ivz begin
+                if vz.grid[ivz] >= -zero
+                    @views @. pdf[ivz,:,:,1] = N_out * knudsen_cosine[ivz,:,:]
+                else
+                    @views @. pdf[ivz,:,:,1] *= N_in
+                end
+            end
+        end
+
+        if z.irank == z.nrank - 1
+            ## treat the right boundary at z = Lz/2 ##
+
+            # Note the numerical integrol of knudsen_cosine was forced to be 1 (to machine
+            # precision) when it was initialised.
+            @views pdf_integral_0 = integrate_over_positive_vz(pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            @views pdf_integral_1 = integrate_over_positive_vz(vz.grid .* pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_0 = integrate_over_negative_vz(knudsen_cosine, vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_1 = -1.0 # This is enforced in initialization
+
+            # Calculate normalisation factors N_in for the incoming and N_out for the
+            # Knudsen parts of the distirbution so that ∫dvpa F = 1 and ∫dvpa vpa F = uz
+            # Note wall_flux_L is the ion flux into the wall (reduced by the recycling
+            # fraction), and the neutral flux should be out of the wall (i.e. uz<0) so
+            # -n*uz = |n*uz| = wall_flux_L
+            # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
+            #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = uz
+            uz = -wall_flux_L / density[end]
+            N_in = (1 - uz * knudsen_integral_0 / knudsen_integral_1) /
+                   (pdf_integral_0
+                    - pdf_integral_1 / knudsen_integral_1 * knudsen_integral_0)
+            N_out = (uz - N_in * pdf_integral_1) / knudsen_integral_1
+
+            @loop_vz ivz begin
+                if vz.grid[ivz] <= zero
+                    @views @. pdf[ivz,:,:,end] = N_out * knudsen_cosine[ivz,:,:]
+                else
+                    @views @. pdf[ivz,:,:,end] *= N_in
+                end
+            end
+        end
+    else
+        if z.irank == 0
+            ## treat z = -Lz/2 boundary ##
+            # populate vz.scratch2 array with dz/dt values at z = -Lz/2
+            if evolve_ppar
+                vth = sqrt(2.0*pz[1]/density[1])
+            else
+                vth = nothing
+            end
+            @. vz.scratch2 = vpagrid_to_dzdt(vz.grid, vth, uz[1], evolve_ppar, evolve_upar)
+
+            # First apply boundary condition that total neutral outflux is equal to ion
+            # influx to uz
+            uz[1] = wall_flux_0 / density[1]
+            #would setting density work better??
+            #density[1] = - wall_flux_0 / uz[1]
+
+            # Create normalised Knudsen cosine distribution, to use for positive v_parallel
+            # at z = -Lz/2
+            # Note this only makes sense for the 1V case with vr.n=vzeta.n=1
+            @. vz.scratch = (3.0*pi/vtfac^3)*abs(vz.scratch2)*erfc(abs(vz.scratch2)/vtfac)
+
+            # The v_parallel>0 part of the pdf is replaced by the Knudsen cosine
+            # distribution. To ensure the constraints ∫dwpa wpa^m F = 0 are satisfied when
+            # necessary, calculate a normalisation factor for the Knudsen distribution (in
+            # vz.scratch) and correction terms for the incoming pdf similar to
+            # enforce_moment_constraints!().
+            #
+            # Note that it seems to be important that this boundary condition not be
+            # modified by the moment constraints, as that could cause numerical instability.
+            # By ensuring that the constraints are satisfied already here,
+            # enforce_moment_constraints!() will not change the pdf at the boundary. For
+            # ions this is not an issue, because points set to 0 by the bc are not modified
+            # from 0 by enforce_moment_constraints!().
+            knudsen_integral_0 = integrate_over_positive_vz(vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_1 = integrate_over_positive_vz(vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+
+            @views pdf_integral_0 = integrate_over_negative_vz(pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            @views pdf_integral_1 = integrate_over_negative_vz(vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            if !evolve_ppar
+                # Calculate normalisation factors N_in for the incoming and N_out for the
+                # Knudsen parts of the distirbution so that ∫dwpa F = 1 and ∫dwpa wpa F = 0
+                # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
+                #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = 0
+                N_in = 1.0 / (pdf_integral_0 - pdf_integral_1/knudsen_integral_1*knudsen_integral_0)
+                N_out = -N_in * pdf_integral_1 / knudsen_integral_1
+
+                zero_vz_ind = 0
+                for ivz ∈ 1:vz.n
+                    if vz.scratch2[ivz] <= -zero
+                        pdf[ivz,:,:,1] .= N_in*pdf[ivz,:,:,1]
+                    else
+                        zero_vz_ind = ivz
+                        if abs(vz.scratch2[ivz]) < zero
+                            # v_z = 0 point, half contribution from original pdf and half
+                            # from Knudsen cosine distribution, to be consistent with weights
+                            # used in
+                            # integrate_over_positive_vz()/integrate_over_negative_vz().
+                            @. pdf[ivz,:,:,1] = 0.5*(N_in*pdf[ivz,:,:,1] + N_out*vz.scratch[ivz])
+                        else
+                            pdf[ivz,:,:,1] .= N_out*vz.scratch[ivz]
+                        end
+                        break
+                    end
+                end
+                for ivz ∈ zero_vz_ind+1:vz.n
+                    pdf[ivz,:,:,1] .= N_out*vz.scratch[ivz]
+                end
+            else
+                knudsen_integral_2 = integrate_over_positive_vz(vz.grid .* vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                @views pdf_integral_2 = integrate_over_negative_vz(vz.grid .* vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                @views pdf_integral_3 = integrate_over_negative_vz(vz.grid .* vz.grid .* vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                # Calculate normalisation factor N_out for the Knudsen part of the
+                # distirbution and normalisation factor N_in and correction term C*wpa*F_in
+                # for the incoming distribution so that ∫dwpa F = 1, ∫dwpa wpa F = 0, and
+                # ∫dwpa wpa^2 F = 1/2
+                # ⇒ N_in*pdf_integral_0 + C*pdf_integral_1 + N_out*knudsen_integral_0 = 1
+                #   N_in*pdf_integral_1 + C*pdf_integral_2 + N_out*knudsen_integral_1 = 0
+                #   N_in*pdf_integral_2 + C*pdf_integral_3 + N_out*knudsen_integral_2 = 1/2
+                N_in = (0.5*knudsen_integral_0*pdf_integral_2 +
+                        knudsen_integral_1*(pdf_integral_3 - 0.5*pdf_integral_1) -
+                        knudsen_integral_2*pdf_integral_2) /
+                       (knudsen_integral_0*(pdf_integral_2^2 - pdf_integral_1*pdf_integral_3) +
+                        knudsen_integral_1*(pdf_integral_0*pdf_integral_3 - pdf_integral_1*pdf_integral_2) +
+                        knudsen_integral_2*(pdf_integral_1^2 - pdf_integral_0*pdf_integral_2))
+                N_out = -(N_in*(pdf_integral_1*pdf_integral_3 - pdf_integral_2^2) + 0.5*pdf_integral_2) /
+                         (knudsen_integral_1*pdf_integral_3 - knudsen_integral_2*pdf_integral_2)
+                C = (0.5 - N_out*knudsen_integral_2 - N_in*pdf_integral_2)/pdf_integral_3
+
+                zero_vz_ind = 0
+                for ivz ∈ 1:vz.n
+                    if vz.scratch2[ivz] <= -zero
+                        @views @. pdf[ivz,:,:,1] = N_in*pdf[ivz,:,:,1] + C*vz.grid[ivz]*pdf[ivz,:,:,1]
+                    else
+                        zero_vz_ind = ivz
+                        if abs(vz.scratch2[ivz]) < zero
+                            # v_parallel = 0 point, half contribution from original pdf and half
+                            # from Knudsen cosine distribution, to be consistent with weights
+                            # used in
+                            # integrate_over_positive_vz()/integrate_over_negative_vz().
+                            @views @. pdf[ivz,:,:,1] = 0.5*(N_in*pdf[ivz,:,:,1] +
+                                                            C*vz.grid[ivz]*pdf[ivz,:,:,1] +
+                                                            N_out*vz.scratch[ivz])
+                        else
+                            @. pdf[ivz,:,:,1] = N_out*vz.scratch[ivz]
+                        end
+                        break
+                    end
+                end
+                for ivz ∈ zero_vz_ind+1:vz.n
+                    @. pdf[ivz,:,:,1] = N_out*vz.scratch[ivz]
+                end
+            end
+        end
+
+        if z.irank == z.nrank - 1
+            ## treat the right boundary at z = Lz/2 ##
+            # populate vz.scratch2 array with dz/dt values at z = Lz/2
+            if evolve_ppar
+                vth = sqrt(2.0*pz[end]/density[end])
+            else
+                vth = nothing
+            end
+            @. vz.scratch2 = vpagrid_to_dzdt(vz.grid, vth, uz[end], evolve_ppar, evolve_upar)
+
+            # First apply boundary condition that total neutral outflux is equal to ion
+            # influx to uz
+            uz[end] = - wall_flux_L / density[end]
+            #would setting density work better??
+            #density[end] = - wall_flux_L / upar[end]
+
+            # obtain the Knudsen cosine distribution at z = Lz/2
+            # the z-dependence is only introduced if the peculiiar velocity is used as vz
+            # Note this only makes sense for the 1V case with vr.n=vzeta.n=1
+            @. vz.scratch = (3.0*pi/vtfac^3)*abs(vz.scratch2)*erfc(abs(vz.scratch2)/vtfac)
+
+            # The v_parallel<0 part of the pdf is replaced by the Knudsen cosine
+            # distribution. To ensure the constraint ∫dwpa wpa F = 0 is satisfied, multiply
+            # the Knudsen distribution (in vz.scratch) by a normalisation factor given by
+            # the integral (over negative v_parallel) of the outgoing Knudsen distribution
+            # and (over positive v_parallel) of the incoming pdf.
+            knudsen_integral_0 = integrate_over_negative_vz(vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            knudsen_integral_1 = integrate_over_negative_vz(vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+
+            @views pdf_integral_0 = integrate_over_positive_vz(pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+            @views pdf_integral_1 = integrate_over_positive_vz(vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+
+            if !evolve_ppar
+                # Calculate normalisation factors N_in for the incoming and N_out for the
+                # Knudsen parts of the distirbution so that ∫dwpa F = 1 and ∫dwpa wpa F = 0
+                # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
+                #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = 0
+                N_in = 1.0 / (pdf_integral_0 - pdf_integral_1/knudsen_integral_1*knudsen_integral_0)
+                N_out = -N_in * pdf_integral_1 / knudsen_integral_1
+
+                zero_vz_ind = 0
+                for ivz ∈ vz.n:-1:1
+                    if vz.scratch2[ivz] >= zero
+                        @views @. pdf[ivz,:,:,end] = N_in*pdf[ivz,:,:,end]
+                    else
+                        zero_vz_ind = ivz
+                        if abs(vz.scratch2[ivz]) < zero
+                            # v_parallel = 0 point, half contribution from original pdf and half
+                            # from Knudsen cosine distribution, to be consistent with weights
+                            # used in
+                            # integrate_over_positive_vz()/integrate_over_negative_vz().
+                            @views @. pdf[ivz,:,:,end] = 0.5*(N_in*pdf[ivz,:,:,end] + N_out*vz.scratch[ivz])
+                        else
+                            @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
+                        end
+                        break
+                    end
+                end
+                for ivz ∈ 1:zero_vz_ind-1
+                    @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
+                end
+            else
+                knudsen_integral_2 = integrate_over_negative_vz(vz.grid .* vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                @views pdf_integral_2 = integrate_over_positive_vz(vz.grid .* vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                @views pdf_integral_3 = integrate_over_positive_vz(vz.grid .* vz.grid .* vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
+                # Calculate normalisation factor N_out for the Knudsen part of the
+                # distirbution and normalisation factor N_in and correction term C*wpa*F_in
+                # for the incoming distribution so that ∫dwpa F = 1, ∫dwpa wpa F = 0, and
+                # ∫dwpa wpa^2 F = 1/2
+                # ⇒ N_in*pdf_integral_0 + C*pdf_integral_1 + N_out*knudsen_integral_0 = 1
+                #   N_in*pdf_integral_1 + C*pdf_integral_2 + N_out*knudsen_integral_1 = 0
+                #   N_in*pdf_integral_2 + C*pdf_integral_3 + N_out*knudsen_integral_2 = 1/2
+                N_in = (0.5*knudsen_integral_0*pdf_integral_2 +
+                        knudsen_integral_1*(pdf_integral_3 - 0.5*pdf_integral_1) -
+                        knudsen_integral_2*pdf_integral_2) /
+                       (knudsen_integral_0*(pdf_integral_2^2 - pdf_integral_1*pdf_integral_3) +
+                        knudsen_integral_1*(pdf_integral_0*pdf_integral_3 - pdf_integral_1*pdf_integral_2) +
+                        knudsen_integral_2*(pdf_integral_1^2 - pdf_integral_0*pdf_integral_2))
+                N_out = -(N_in*(pdf_integral_1*pdf_integral_3 - pdf_integral_2^2) + 0.5*pdf_integral_2) /
+                         (knudsen_integral_1*pdf_integral_3 - knudsen_integral_2*pdf_integral_2)
+                C = (0.5 - N_out*knudsen_integral_2 - N_in*pdf_integral_2)/pdf_integral_3
+
+                zero_vz_ind = 0
+                for ivz ∈ vz.n:-1:1
+                    if vz.scratch2[ivz] >= zero
+                        @views @. pdf[ivz,:,:,end] = N_in*pdf[ivz,:,:,end] + C*vz.grid[ivz]*pdf[ivz,:,:,end]
+                    else
+                        zero_vz_ind = ivz
+                        if abs(vz.scratch2[ivz]) < zero
+                            # v_parallel = 0 point, half contribution from original pdf and half
+                            # from Knudsen cosine distribution, to be consistent with weights
+                            # used in
+                            # integrate_over_positive_vz()/integrate_over_negative_vz().
+                            @views @. pdf[ivz,:,:,end] = 0.5*(N_in*pdf[ivz,:,:,end] +
+                                                              C*vz.grid[ivz]*pdf[ivz,:,:,end] +
+                                                              N_out*vz.scratch[ivz])
+                        else
+                            @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
+                        end
+                        break
+                    end
+                end
+                for ivz ∈ 1:zero_vz_ind-1
+                    @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
+                end
+            end
+        end
+    end
+end
+
+"""
+create an array of dz/dt values corresponding to the given vpagrid values
+"""
+function vpagrid_to_dzdt(vpagrid, vth, upar, evolve_ppar, evolve_upar)
+    if evolve_ppar
+        if evolve_upar
+            return vpagrid .* vth .+ upar
+        else
+            return vpagrid .* vth
+        end
+    elseif evolve_upar
+        return vpagrid .+ upar
+    else
+        return vpagrid
+    end
+end
+
+"""
+enforce the z boundary condition on the evolved velocity space moments of f
+"""
+function enforce_z_boundary_condition_moments!(density, moments, bc::String)
+    ## TODO: parallelise
+    #begin_serial_region()
+    #@serial_region begin
+    #    # enforce z boundary condition on density if it is evolved separately from f
+    #	if moments.evolve_density
+    #        # TODO: extend to 'periodic' BC case, as this requires further code modifications to be consistent
+    #        # with finite difference derivatives (should be fine for Chebyshev)
+    #        if bc == "wall"
+    #            @loop_s_r is ir begin
+    #                density[1,ir,is] = 0.5*(density[1,ir,is] + density[end,ir,is])
+    #                density[end,ir,is] = density[1,ir,is]
+    #        	end
+    #        end
+    #    end
+    #end
+end
+
+"""
+"""
+function enforce_v_boundary_condition_local!(f, bc, speed, v_diffusion, v, v_spectral)
+    if bc == "zero"
+        if v_diffusion || speed[1] > 0.0
+            # 'upwind' boundary
+            f[1] = 0.0
+        end
+        if v_diffusion || speed[end] < 0.0
+            # 'upwind' boundary
+            f[end] = 0.0
+        end
+    elseif bc == "both_zero"
+        f[1] = 0.0
+        f[end] = 0.0
+    elseif bc == "zero_gradient"
+        D0 = v_spectral.lobatto.Dmat[1,:]
+        # adjust F(vpa = -L/2) so that d F / d vpa = 0 at vpa = -L/2
+        f[1] = -sum(D0[2:v.ngrid].*f[2:v.ngrid])/D0[1]
+
+        D0 = v_spectral.lobatto.Dmat[end,:]
+        # adjust F(vpa = L/2) so that d F / d vpa = 0 at vpa = L/2
+        f[end] = -sum(D0[1:ngrid-1].*f[end-v.ngrid+1:end-1])/D0[v.ngrid]
+    elseif bc == "periodic"
+        f[1] = 0.5*(f[1]+f[end])
+        f[end] = f[1]
+    else
+        error("Unsupported boundary condition option '$bc' for $(v.name)")
+    end
+end
+
+"""
+enforce zero boundary condition at vperp -> infinity
+"""
+function enforce_vperp_boundary_condition! end
+
+function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vperp, vperp_spectral)
+    @loop_s is begin
+        @views enforce_vperp_boundary_condition!(f[:,:,:,:,is], bc, vperp, vperp_spectral)
+    end
+    return nothing
+end
+
+function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral)
+    if bc == "zero"
+        nvperp = vperp.n
+        ngrid = vperp.ngrid
+        # set zero boundary condition
+        @loop_r_z_vpa ir iz ivpa begin
+            f[ivpa,nvperp,iz,ir] = 0.0
+        end
+        # set regularity condition d F / d vperp = 0 at vperp = 0
+        if vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral"
+            D0 = vperp_spectral.radau.D0
+            @loop_r_z_vpa ir iz ivpa begin
+                # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
+                f[ivpa,1,iz,ir] = -sum(D0[2:ngrid].*f[ivpa,2:ngrid,iz,ir])/D0[1]
+            end
+        else
+            println("vperp.bc=\"$bc\" not supported by discretization "
+                    * "$(vperp.discretization)")
+        end
+    elseif bc == "none"
+        # Do nothing
+    else
+        error("Unsupported boundary condition option '$bc' for vperp")
+    end
+end
+
+end # boundary_conditions
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 9f7736c92..74dec52d7 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -4,8 +4,6 @@ module initial_conditions
 
 export allocate_pdf_and_moments
 export init_pdf_and_moments!
-export enforce_boundary_conditions!
-export enforce_neutral_boundary_conditions!
 
 # functional testing 
 export create_boundary_distributions
@@ -17,9 +15,8 @@ using SpecialFunctions: erfc
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_shared_float
 using ..bgk: init_bgk_pdf!
+using ..boundary_conditions: vpagrid_to_dzdt
 using ..communication
-using ..calculus: reconcile_element_boundaries_MPI!
-using ..coordinates: coordinate
 using ..external_sources
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
@@ -30,7 +27,6 @@ using ..load_data: reload_electron_data!
 using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, pdf_struct, moments_struct,
                                  boundary_distributions_struct
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
-using ..velocity_moments: integrate_over_positive_vpa, integrate_over_negative_vpa
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
 using ..velocity_moments: create_moments_ion, create_moments_electron, create_moments_neutral
 using ..velocity_moments: update_qpar!
@@ -1463,896 +1459,6 @@ function init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta
     return nothing
 end
 
-"""
-enforce boundary conditions in vpa and z on the evolved pdf;
-also enforce boundary conditions in z on all separately evolved velocity space moments of the pdf
-"""
-function enforce_boundary_conditions!(f, f_r_bc, density, upar, ppar, moments, vpa_bc,
-        z_bc, r_bc, vpa, vperp, z, r, vpa_spectral, vperp_spectral, vpa_adv, z_adv, r_adv,
-        composition, scratch_dummy, r_diffusion, vpa_diffusion)
-    if vpa.n > 1
-        begin_s_r_z_vperp_region()
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            # enforce the vpa BC
-            # use that adv.speed independent of vpa 
-            @views enforce_v_boundary_condition_local!(f[:,ivperp,iz,ir,is], vpa_bc,
-                             vpa_adv[is].speed[:,ivperp,iz,ir], vpa_diffusion,
-                             vpa, vpa_spectral)
-        end
-    end
-    if vperp.n > 1
-        begin_s_r_z_vpa_region()
-        @views enforce_vperp_boundary_condition!(f, vperp.bc, vperp, vperp_spectral)
-    end
-    if z.n > 1
-        begin_s_r_vperp_vpa_region()
-        # enforce the z BC on the evolved velocity space moments of the pdf
-        @views enforce_z_boundary_condition_moments!(density, moments, z_bc)
-        @views enforce_z_boundary_condition!(f, density, upar, ppar, moments, z_bc, z_adv, z,
-                                             vperp, vpa, composition,
-                                             scratch_dummy.buffer_vpavperprs_1, scratch_dummy.buffer_vpavperprs_2,
-                                             scratch_dummy.buffer_vpavperprs_3, scratch_dummy.buffer_vpavperprs_4)
-                                              
-    end
-    if r.n > 1
-        begin_s_z_vperp_vpa_region()
-        @views enforce_r_boundary_condition!(f, f_r_bc, r_bc, r_adv, vpa, vperp, z, r, composition,
-            scratch_dummy.buffer_vpavperpzs_1, scratch_dummy.buffer_vpavperpzs_2,
-            scratch_dummy.buffer_vpavperpzs_3, scratch_dummy.buffer_vpavperpzs_4,
-            r_diffusion)
-    end
-end
-function enforce_boundary_conditions!(fvec_out::scratch_pdf, moments, f_r_bc, vpa_bc,
-        z_bc, r_bc, vpa, vperp, z, r, vpa_spectral, vperp_spectral, vpa_adv, z_adv, r_adv, composition, scratch_dummy,
-        r_diffusion, vpa_diffusion)
-    enforce_boundary_conditions!(fvec_out.pdf, f_r_bc, fvec_out.density, fvec_out.upar,
-        fvec_out.ppar, moments, vpa_bc, z_bc, r_bc, vpa, vperp, z, r, 
-        vpa_spectral, vperp_spectral, vpa_adv, z_adv,
-        r_adv, composition, scratch_dummy, r_diffusion, vpa_diffusion)
-end
-
-"""
-enforce boundary conditions on f in r
-"""
-function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc::String,
-        adv, vpa, vperp, z, r, composition, end1::AbstractArray{mk_float,4},
-        end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4},
-        buffer2::AbstractArray{mk_float,4}, r_diffusion::Bool)
-
-    nr = r.n
-
-    if r.nelement_global > r.nelement_local
-        # reconcile internal element boundaries across processes
-        # & enforce periodicity and external boundaries if needed
-        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
-            end1[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,1,is]
-            end2[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,nr,is]
-        end
-        @views reconcile_element_boundaries_MPI!(f,
-            end1, end2,	buffer1, buffer2, r)
-    end
-
-    # 'periodic' BC enforces periodicity by taking the average of the boundary points
-    # enforce the condition if r is local
-    if bc == "periodic" && r.nelement_global == r.nelement_local
-        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
-            f[ivpa,ivperp,iz,1,is] = 0.5*(f[ivpa,ivperp,iz,nr,is]+f[ivpa,ivperp,iz,1,is])
-            f[ivpa,ivperp,iz,nr,is] = f[ivpa,ivperp,iz,1,is]
-        end
-    end
-    if bc == "Dirichlet"
-        zero = 1.0e-10
-        # use the old distribution to force the new distribution to have
-        # consistant-in-time values at the boundary
-        # with bc = "Dirichlet" and r_diffusion = false
-        # impose bc for incoming parts of velocity space only (Hyperbolic PDE)
-        # with bc = "Dirichlet" and r_diffusion = true
-        # impose bc on both sides of the domain to accomodate a diffusion operator d^2 / d r^2
-        @loop_s_z_vperp_vpa is iz ivperp ivpa begin
-            ir = 1 # r = -L/2 -- check that the point is on lowest rank
-            if r.irank == 0 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] > zero)
-                f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,1,is]
-            end
-            ir = r.n # r = L/2 -- check that the point is on highest rank
-            if r.irank == r.nrank - 1 && (r_diffusion || adv[is].speed[ir,ivpa,ivperp,iz] < -zero)
-                f[ivpa,ivperp,iz,ir,is] = f_r_bc[ivpa,ivperp,iz,end,is]
-            end
-        end
-    end
-end
-
-"""
-enforce boundary conditions on ion particle f in z
-"""
-function enforce_z_boundary_condition!(pdf, density, upar, ppar, moments, bc::String, adv,
-                                       z, vperp, vpa, composition, end1::AbstractArray{mk_float,4},
-                                       end2::AbstractArray{mk_float,4}, buffer1::AbstractArray{mk_float,4},
-                                       buffer2::AbstractArray{mk_float,4})
-    # this block ensures periodic BC can be supported with distributed memory MPI
-    if z.nelement_global > z.nelement_local
-        # reconcile internal element boundaries across processes
-        # & enforce periodicity and external boundaries if needed
-        nz = z.n
-        @loop_s_r_vperp_vpa is ir ivperp ivpa begin
-            end1[ivpa,ivperp,ir,is] = pdf[ivpa,ivperp,1,ir,is]
-            end2[ivpa,ivperp,ir,is] = pdf[ivpa,ivperp,nz,ir,is]
-        end
-        # check on periodic bc happens inside this call below
-        @views reconcile_element_boundaries_MPI!(pdf,
-            end1, end2,	buffer1, buffer2, z)
-    end
-    # define a zero that accounts for finite precision
-    zero = 1.0e-14
-    # 'constant' BC is time-independent f at upwind boundary
-    # and constant f beyond boundary 
-    if bc == "constant"
-        begin_s_r_vperp_vpa_region()
-        density_offset = 1.0
-        vwidth = 1.0
-        if z.irank == 0
-            @loop_s_r_vperp_vpa is ir ivperp ivpa begin
-                if adv[is].speed[ivpa,1,ir] > 0.0
-                    pdf[ivpa,ivperp,1,ir,is] = density_offset * exp(-(vpa.grid[ivpa]^2 + vperp.grid[ivperp]^2)/vwidth^2) / sqrt(pi)
-                end
-            end
-        end
-        if z.irank == z.nrank - 1
-            @loop_s_r_vperp_vpa is ir ivperp ivpa begin
-                if adv[is].speed[ivpa,end,ir] > 0.0
-                    pdf[ivpa,ivperp,end,ir,is] = density_offset * exp(-(vpa.grid[ivpa]^2 + vperp.grid[ivperp]^2)/vwidth^2) / sqrt(pi)
-                end
-            end
-        end
-    # 'periodic' BC enforces periodicity by taking the average of the boundary points
-    elseif bc == "periodic" && z.nelement_global == z.nelement_local
-        begin_s_r_vperp_vpa_region()
-        @loop_s_r_vperp_vpa is ir ivperp ivpa begin
-            pdf[ivpa,ivperp,1,ir,is] = 0.5*(pdf[ivpa,ivperp,z.n,ir,is]+pdf[ivpa,ivperp,1,ir,is])
-            pdf[ivpa,ivperp,z.n,ir,is] = pdf[ivpa,ivperp,1,ir,is]
-        end
-    # 'wall' BC enforces wall boundary conditions
-    elseif bc == "wall"
-        # Need integrals over vpa at wall boundaries in z, so cannot parallelize over z
-        # or vpa.
-        begin_s_r_region()
-        @loop_s is begin
-            # zero incoming BC for ions, as they recombine at the wall
-            if moments.evolve_upar
-                @loop_r ir begin
-                    @views enforce_zero_incoming_bc!(
-                        pdf[:,:,:,ir,is], z, vpa, density[:,ir,is], upar[:,ir,is],
-                        ppar[:,ir,is], moments.evolve_upar, moments.evolve_ppar, zero)
-                end
-            else
-                @loop_r ir begin
-                    @views enforce_zero_incoming_bc!(pdf[:,:,:,ir,is],
-                                                     adv[is].speed[:,:,:,ir], z, zero)
-                end
-            end
-        end
-    end
-end
-
-"""
-enforce boundary conditions on neutral particle distribution function
-"""
-function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
-        boundary_distributions, density_neutral, uz_neutral, pz_neutral, moments,
-        density_ion, upar_ion, Er, vzeta_spectral, vr_spectral, vz_spectral, r_adv, z_adv,
-        vzeta_adv, vr_adv, vz_adv, r, z, vzeta, vr, vz, composition, geometry,
-        scratch_dummy, r_diffusion, vz_diffusion)
-
-    # without acceleration of neutrals bc on vz vr vzeta should not be required as no
-    # advection or diffusion in these coordinates
-
-    if vzeta_adv !== nothing && vzeta.n_global > 1 && vzeta.bc != "none"
-        begin_sn_r_z_vr_vz_region()
-        @loop_sn_r_z_vr_vz isn ir iz ivr ivz begin
-            # enforce the vz BC
-            @views enforce_v_boundary_condition_local!(f_neutral[ivz,ivr,:,iz,ir,isn],
-                                                       vzeta.bc,
-                                                       vzeta_adv[isn].speed[ivz,ivr,:,iz,ir],
-                                                       false, vzeta, vzeta_spectral)
-        end
-    end
-    if vr_adv !== nothing && vr.n_global > 1 && vr.bc != "none"
-        begin_sn_r_z_vzeta_vz_region()
-        @loop_sn_r_z_vzeta_vz isn ir iz ivzeta ivz begin
-            # enforce the vz BC
-            @views enforce_v_boundary_condition_local!(f_neutral[ivz,:,ivzeta,iz,ir,isn],
-                                                       vr.bc,
-                                                       vr_adv[isn].speed[ivz,:,ivzeta,iz,ir],
-                                                       false, vr, vr_spectral)
-        end
-    end
-    if vz_adv !== nothing && vz.n_global > 1 && vz.bc != "none"
-        begin_sn_r_z_vzeta_vr_region()
-        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
-            # enforce the vz BC
-            @views enforce_v_boundary_condition_local!(f_neutral[:,ivr,ivzeta,iz,ir,isn],
-                                                       vz.bc,
-                                                       vz_adv[isn].speed[:,ivr,ivzeta,iz,ir],
-                                                       vz_diffusion, vz, vz_spectral)
-        end
-    end
-    # f_initial contains the initial condition for enforcing a fixed-boundary-value condition
-    if z.n > 1
-        begin_sn_r_vzeta_vr_vz_region()
-        @views enforce_neutral_z_boundary_condition!(f_neutral, density_neutral, uz_neutral,
-            pz_neutral, moments, density_ion, upar_ion, Er, boundary_distributions,
-            z_adv, z, vzeta, vr, vz, composition, geometry, 
-            scratch_dummy.buffer_vzvrvzetarsn_1, scratch_dummy.buffer_vzvrvzetarsn_2,
-            scratch_dummy.buffer_vzvrvzetarsn_3, scratch_dummy.buffer_vzvrvzetarsn_4)
-    end
-    if r.n > 1
-        begin_sn_z_vzeta_vr_vz_region()
-        @views enforce_neutral_r_boundary_condition!(f_neutral, boundary_distributions.pdf_rboundary_neutral,
-                                    r_adv, vz, vr, vzeta, z, r, composition,
-                                    scratch_dummy.buffer_vzvrvzetazsn_1, scratch_dummy.buffer_vzvrvzetazsn_2,
-                                    scratch_dummy.buffer_vzvrvzetazsn_3, scratch_dummy.buffer_vzvrvzetazsn_4,
-                                    r_diffusion)
-    end
-end
-
-function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6},
-        f_r_bc::AbstractArray{mk_float,6}, adv, vz, vr, vzeta, z, r, composition,
-        end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5},
-        buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5},
-        r_diffusion) #f_initial,
-
-    bc = r.bc
-    nr = r.n
-
-    if r.nelement_global > r.nelement_local
-        # reconcile internal element boundaries across processes
-        # & enforce periodicity and external boundaries if needed
-        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
-            end1[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,1,isn]
-            end2[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,nr,isn]
-        end
-        @views reconcile_element_boundaries_MPI!(f,
-            end1, end2,	buffer1, buffer2, r)
-    end
-    # 'periodic' BC enforces periodicity by taking the average of the boundary points
-    # local case only when no communication required
-    if bc == "periodic" && r.nelement_global == r.nelement_local
-        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
-            f[ivz,ivr,ivzeta,iz,1,isn] = 0.5*(f[ivz,ivr,ivzeta,iz,1,isn]+f[ivz,ivr,ivzeta,iz,nr,isn])
-            f[ivz,ivr,ivzeta,iz,nr,isn] = f[ivz,ivr,ivzeta,iz,1,isn]
-        end
-    end
-    # Dirichlet boundary condition for external endpoints
-    if bc == "Dirichlet"
-        zero = 1.0e-10
-        # use the old distribution to force the new distribution to have
-        # consistant-in-time values at the boundary
-        # impose bc for incoming parts of velocity space only (Hyperbolic PDE)
-        @loop_sn_z_vzeta_vr_vz isn iz ivzeta ivr ivz begin
-            ir = 1 # r = -L/2
-            # incoming particles and on lowest rank
-            if r.irank == 0 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] > zero)
-                f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,1,isn]
-            end
-            ir = nr # r = L/2
-            # incoming particles and on highest rank
-            if r.irank == r.nrank - 1 && (r_diffusion || adv[isn].speed[ir,ivz,ivr,ivzeta,iz] < -zero)
-                f[ivz,ivr,ivzeta,iz,ir,isn] = f_r_bc[ivz,ivr,ivzeta,iz,end,isn]
-            end
-        end
-    end
-end
-
-"""
-enforce boundary conditions on neutral particle f in z
-"""
-function enforce_neutral_z_boundary_condition!(pdf, density, uz, pz, moments, density_ion,
-                                               upar_ion, Er, boundary_distributions, adv,
-                                               z, vzeta, vr, vz, composition, geometry,
-                                               end1::AbstractArray{mk_float,5}, end2::AbstractArray{mk_float,5},
-                                               buffer1::AbstractArray{mk_float,5}, buffer2::AbstractArray{mk_float,5})
-    
-
-    if z.nelement_global > z.nelement_local
-        # reconcile internal element boundaries across processes
-        # & enforce periodicity and external boundaries if needed
-        nz = z.n
-        @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
-            end1[ivz,ivr,ivzeta,ir,isn] = pdf[ivz,ivr,ivzeta,1,ir,isn]
-            end2[ivz,ivr,ivzeta,ir,isn] = pdf[ivz,ivr,ivzeta,nz,ir,isn]
-        end
-        # check on periodic bc occurs within this call below
-        @views reconcile_element_boundaries_MPI!(pdf,
-            end1, end2,	buffer1, buffer2, z)
-    end
-
-    zero = 1.0e-14
-    # 'constant' BC is time-independent f at upwind boundary
-    # and constant f beyond boundary
-    if z.bc == "constant"
-        begin_sn_r_vzeta_vr_vz_region()
-        density_offset = 1.0
-        vwidth = 1.0
-        if z.irank == 0
-            @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
-                if adv[isn].speed[ivz,ivr,ivzeta,1,ir] > 0.0
-                    pdf[ivz,ivr,ivzeta,1,ir,is] = density_offset *
-                        exp(-(vzeta.grid[ivzeta]^2 + vr.grid[ivr] + vz.grid[ivz])/vwidth^2) /
-                        sqrt(pi)
-                end
-            end
-        end
-        if z.irank == z.nrank - 1
-            @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
-                if adv[isn].speed[ivz,ivr,ivzeta,end,ir] > 0.0
-                    pdf[ivz,ivr,ivzeta,end,ir,is] = density_offset *
-                        exp(-(vzeta.grid[ivzeta]^2 + vr.grid[ivr] + vz.grid[ivz])/vwidth^2) /
-                        sqrt(pi)
-                end
-            end
-        end
-    # 'periodic' BC enforces periodicity by taking the average of the boundary points
-    elseif z.bc == "periodic" && z.nelement_global == z.nelement_local
-        begin_sn_r_vzeta_vr_vz_region()
-        @loop_sn_r_vzeta_vr_vz isn ir ivzeta ivr ivz begin
-            pdf[ivz,ivr,ivzeta,1,ir,isn] = 0.5*(pdf[ivz,ivr,ivzeta,1,ir,isn] +
-                                                pdf[ivz,ivr,ivzeta,end,ir,isn])
-            pdf[ivz,ivr,ivzeta,end,ir,isn] = pdf[ivz,ivr,ivzeta,1,ir,isn]
-        end
-    # 'wall' BC enforces wall boundary conditions
-    elseif z.bc == "wall"
-        # Need integrals over vpa at wall boundaries in z, so cannot parallelize over z
-        # or vpa.
-        begin_sn_r_region()
-        @loop_sn isn begin
-            # BC for neutrals
-            @loop_r ir begin
-                # define vtfac to avoid repeated computation below
-                vtfac = sqrt(composition.T_wall * composition.mn_over_mi)
-                # Assume for now that the ion species index corresponding to this neutral
-                # species is the same as the neutral species index.
-                # Note, have already calculated moments of ion distribution function(s),
-                # so can use the moments here to get the flux
-                if z.irank == 0
-                    ion_flux_0 = -density_ion[1,ir,isn] * (upar_ion[1,ir,isn]*geometry.bzed - 0.5*geometry.rhostar*Er[1,ir])
-                else
-                    ion_flux_0 = NaN
-                end
-                if z.irank == z.nrank - 1
-                    ion_flux_L = density_ion[end,ir,isn] * (upar_ion[end,ir,isn]*geometry.bzed - 0.5*geometry.rhostar*Er[end,ir])
-                else
-                    ion_flux_L = NaN
-                end
-                # enforce boundary condition on the neutral pdf that all ions and neutrals
-                # that leave the domain re-enter as neutrals
-                @views enforce_neutral_wall_bc!(
-                    pdf[:,:,:,:,ir,isn], z, vzeta, vr, vz, pz[:,ir,isn], uz[:,ir,isn],
-                    density[:,ir,isn], ion_flux_0, ion_flux_L, boundary_distributions,
-                    vtfac, composition.recycling_fraction, moments.evolve_ppar,
-                    moments.evolve_upar, moments.evolve_density, zero)
-            end
-        end
-    end
-end
-
-"""
-enforce a zero incoming BC in z for given species pdf at each radial location
-"""
-function enforce_zero_incoming_bc!(pdf, speed, z, zero)
-    nvpa = size(pdf,1)
-    # no parallel BC should be enforced for dz/dt = 0
-    # note that the parallel velocity coordinate vpa may be dz/dt or
-    # some version of the peculiar velocity (dz/dt - upar),
-    # so use advection speed below instead of vpa
-    if z.irank == 0
-        @loop_vperp_vpa ivperp ivpa begin
-            # for left boundary in zed (z = -Lz/2), want
-            # f(z=-Lz/2, v_parallel > 0) = 0
-            if speed[1,ivpa,ivperp] > zero
-                pdf[ivpa,ivperp,1] = 0.0
-            end
-        end
-    end
-    if z.irank == z.nrank - 1
-        @loop_vperp_vpa ivperp ivpa begin
-            # for right boundary in zed (z = Lz/2), want
-            # f(z=Lz/2, v_parallel < 0) = 0
-            if speed[end,ivpa,ivperp] < -zero
-                pdf[ivpa,ivperp,end] = 0.0
-            end
-        end
-    end
-end
-function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density, upar,
-                                   ppar, evolve_upar, evolve_ppar, zero)
-    if z.irank != 0 && z.irank != z.nrank - 1
-        # No z-boundary in this block
-        return nothing
-    end
-    nvpa, nvperp, nz = size(pdf)
-    # no parallel BC should be enforced for dz/dt = 0
-    # note that the parallel velocity coordinate vpa may be dz/dt or
-    # some version of the peculiar velocity (dz/dt - upar),
-    # so use advection speed below instead of vpa
-
-    # absolute velocity at left boundary
-    if z.irank == 0
-        @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[1]/density[1])),
-                                         upar[1], evolve_ppar, evolve_upar)
-        @loop_vpa ivpa begin
-            # for left boundary in zed (z = -Lz/2), want
-            # f(z=-Lz/2, v_parallel > 0) = 0
-            if vpa.scratch[ivpa] > zero
-                pdf[ivpa,:,1] .= 0.0
-            end
-        end
-    end
-    # absolute velocity at right boundary
-    if z.irank == z.nrank - 1
-        @. vpa.scratch2 = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[end]/density[end])),
-                                          upar[end], evolve_ppar, evolve_upar)
-        @loop_vpa ivpa begin
-            # for right boundary in zed (z = Lz/2), want
-            # f(z=Lz/2, v_parallel < 0) = 0
-            if vpa.scratch2[ivpa] < -zero
-                pdf[ivpa,:,end] .= 0.0
-            end
-        end
-    end
-
-    # Special constraint-forcing code that tries to keep the modifications smooth at
-    # v_parallel=0.
-    if z.irank == 0 && z.irank == z.nrank - 1
-        # Both z-boundaries in this block
-        z_range = (1,nz)
-    elseif z.irank == 0
-        z_range = (1,)
-    elseif z.irank == z.nrank - 1
-        z_range = (nz,)
-    else
-        error("No boundary in this block, should have returned already")
-    end
-    for iz ∈ z_range
-        # moment-kinetic approach only implemented for 1V case so far
-        @boundscheck size(pdf,2) == 1
-
-        f = @view pdf[:,1,iz]
-        if evolve_ppar && evolve_upar
-            I0 = integrate_over_vspace(f, vpa.wgts)
-            I1 = integrate_over_vspace(f, vpa.grid, vpa.wgts)
-            I2 = integrate_over_vspace(f, vpa.grid, 2, vpa.wgts)
-
-            # Store v_parallel with upar shift removed in vpa.scratch
-            vth = sqrt(2.0*ppar[iz]/density[iz])
-            @. vpa.scratch = vpa.grid + upar[iz]/vth
-            # Introduce factor to ensure corrections go smoothly to zero near
-            # v_parallel=0
-            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
-            J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
-            J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
-            J3 = integrate_over_vspace(vpa.scratch2, vpa.grid, 3, vpa.wgts)
-            J4 = integrate_over_vspace(vpa.scratch2, vpa.grid, 4, vpa.wgts)
-
-            A = (J3^2 - J2*J4 + 0.5*(J2^2 - J1*J3)) /
-                (I0*(J3^2 - J2*J4) + I1*(J1*J4 - J2*J3) + I2*(J2^2 - J1*J3))
-            B = (0.5*J3 + A*(I1*J4 - I2*J3)) / (J3^2 - J2*J4)
-            C = (0.5 - A*I2 -B*J3) / J4
-
-            @. f = A*f + B*vpa.grid*vpa.scratch2 + C*vpa.grid*vpa.grid*vpa.scratch2
-        elseif evolve_upar
-            I0 = integrate_over_vspace(f, vpa.wgts)
-            I1 = integrate_over_vspace(f, vpa.grid, vpa.wgts)
-
-            # Store v_parallel with upar shift removed in vpa.scratch
-            @. vpa.scratch = vpa.grid + upar[iz]
-            # Introduce factor to ensure corrections go smoothly to zero near
-            # v_parallel=0
-            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
-            J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
-            J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
-
-            A = 1.0 / (I0 - I1*J1/J2)
-            B = -A*I1/J2
-
-            @. f = A*f + B*vpa.grid*vpa.scratch2
-        elseif evolve_density
-            I0 = integrate_over_vspace(f, vpa.wgts)
-            @. f = f / I0
-        end
-    end
-end
-
-"""
-Set up an initial condition that tries to be smoothly compatible with the sheath
-boundary condition for ions, by setting f(±(v_parallel-u0)<0) where u0=0 at the sheath
-boundaries and for z<0 increases linearly to u0=vpa.L at z=0, while for z>0 increases
-from u0=-vpa.L at z=0 to zero at the z=z.L/2 sheath.
-
-To be applied to 'full-f' distribution function on v_parallel grid (not w_parallel
-grid).
-"""
-function enforce_initial_tapered_zero_incoming!(pdf, z::coordinate, vpa::coordinate)
-    nvpa = size(pdf,1)
-    zero = 1.0e-14
-    # no parallel BC should be enforced for dz/dt = 0
-    # note that the parallel velocity coordinate vpa may be dz/dt or
-    # some version of the peculiar velocity (dz/dt - upar),
-    # so use advection speed below instead of vpa
-
-    for iz ∈ 1:z.n
-        u0 = (2.0*z.grid[iz]/z.L - sign(z.grid[iz])) * vpa.L / 2.0
-        if z.grid[iz] < -zero
-            for ivpa ∈ 1:nvpa
-                if vpa.grid[ivpa] > u0 + zero
-                    pdf[ivpa,iz] = 0.0
-                end
-            end
-        elseif z.grid[iz] > zero
-            for ivpa ∈ 1:nvpa
-                if vpa.grid[ivpa] < u0 - zero
-                    pdf[ivpa,iz] = 0.0
-                end
-            end
-        end
-    end
-end
-
-"""
-enforce the wall boundary condition on neutrals;
-i.e., the incoming flux of neutrals equals the sum of the ion/neutral outgoing fluxes
-"""
-function enforce_neutral_wall_bc!(pdf, z, vzeta, vr, vz, pz, uz, density, wall_flux_0,
-                                  wall_flux_L, boundary_distributions, vtfac,
-                                  recycling_fraction, evolve_ppar, evolve_upar,
-                                  evolve_density, zero)
-
-    # Reduce the ion flux by `recycling_fraction` to account for ions absorbed by the
-    # wall.
-    wall_flux_0 *= recycling_fraction
-    wall_flux_L *= recycling_fraction
-
-    if !evolve_density && !evolve_upar
-        knudsen_cosine = boundary_distributions.knudsen
-
-        if z.irank == 0
-            ## treat z = -Lz/2 boundary ##
-
-            # add the neutral species's contribution to the combined ion/neutral particle
-            # flux out of the domain at z=-Lz/2
-            @views wall_flux_0 += integrate_over_negative_vz(abs.(vz.grid) .* pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-
-            # for left boundary in zed (z = -Lz/2), want
-            # f_n(z=-Lz/2, v_parallel > 0) = Γ_0 * f_KW(v_parallel)
-            @loop_vz ivz begin
-                if vz.grid[ivz] >= -zero
-                    @views @. pdf[ivz,:,:,1] = wall_flux_0 * knudsen_cosine[ivz,:,:]
-                end
-            end
-        end
-
-        if z.irank == z.nrank - 1
-            ## treat the right boundary at z = Lz/2 ##
-
-            # add the neutral species's contribution to the combined ion/neutral particle
-            # flux out of the domain at z=-Lz/2
-            @views wall_flux_L += integrate_over_positive_vz(abs.(vz.grid) .* pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-
-            # for right boundary in zed (z = Lz/2), want
-            # f_n(z=Lz/2, v_parallel < 0) = Γ_Lz * f_KW(v_parallel)
-            @loop_vz ivz begin
-                if vz.grid[ivz] <= zero
-                    @views @. pdf[ivz,:,:,end] = wall_flux_L * knudsen_cosine[ivz,:,:]
-                end
-            end
-        end
-    elseif !evolve_upar
-        # Evolving density case
-        knudsen_cosine = boundary_distributions.knudsen
-
-        if z.irank == 0
-            ## treat z = -Lz/2 boundary ##
-
-            # Note the numerical integrol of knudsen_cosine was forced to be 1 (to machine
-            # precision) when it was initialised.
-            @views pdf_integral_0 = integrate_over_negative_vz(pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            @views pdf_integral_1 = integrate_over_negative_vz(vz.grid .* pdf[:,:,:,1], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_0 = integrate_over_positive_vz(knudsen_cosine, vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_1 = 1.0 # This is enforced in initialization
-
-            # Calculate normalisation factors N_in for the incoming and N_out for the
-            # Knudsen parts of the distirbution so that ∫dvpa F = 1 and ∫dvpa vpa F = uz
-            # Note wall_flux_0 is the ion flux into the wall (reduced by the recycling
-            # fraction), and the neutral flux should be out of the wall (i.e. uz>0) so
-            # n*uz = |n*uz| = wall_flux_0
-            # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
-            #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = uz
-            uz = wall_flux_0 / density[1]
-            N_in = (1 - uz * knudsen_integral_0 / knudsen_integral_1) /
-                   (pdf_integral_0
-                    - pdf_integral_1 / knudsen_integral_1 * knudsen_integral_0)
-            N_out = (uz - N_in * pdf_integral_1) / knudsen_integral_1
-
-            @loop_vz ivz begin
-                if vz.grid[ivz] >= -zero
-                    @views @. pdf[ivz,:,:,1] = N_out * knudsen_cosine[ivz,:,:]
-                else
-                    @views @. pdf[ivz,:,:,1] *= N_in
-                end
-            end
-        end
-
-        if z.irank == z.nrank - 1
-            ## treat the right boundary at z = Lz/2 ##
-
-            # Note the numerical integrol of knudsen_cosine was forced to be 1 (to machine
-            # precision) when it was initialised.
-            @views pdf_integral_0 = integrate_over_positive_vz(pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            @views pdf_integral_1 = integrate_over_positive_vz(vz.grid .* pdf[:,:,:,end], vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_0 = integrate_over_negative_vz(knudsen_cosine, vz.grid, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_1 = -1.0 # This is enforced in initialization
-
-            # Calculate normalisation factors N_in for the incoming and N_out for the
-            # Knudsen parts of the distirbution so that ∫dvpa F = 1 and ∫dvpa vpa F = uz
-            # Note wall_flux_L is the ion flux into the wall (reduced by the recycling
-            # fraction), and the neutral flux should be out of the wall (i.e. uz<0) so
-            # -n*uz = |n*uz| = wall_flux_L
-            # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
-            #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = uz
-            uz = -wall_flux_L / density[end]
-            N_in = (1 - uz * knudsen_integral_0 / knudsen_integral_1) /
-                   (pdf_integral_0
-                    - pdf_integral_1 / knudsen_integral_1 * knudsen_integral_0)
-            N_out = (uz - N_in * pdf_integral_1) / knudsen_integral_1
-
-            @loop_vz ivz begin
-                if vz.grid[ivz] <= zero
-                    @views @. pdf[ivz,:,:,end] = N_out * knudsen_cosine[ivz,:,:]
-                else
-                    @views @. pdf[ivz,:,:,end] *= N_in
-                end
-            end
-        end
-    else
-        if z.irank == 0
-            ## treat z = -Lz/2 boundary ##
-            # populate vz.scratch2 array with dz/dt values at z = -Lz/2
-            if evolve_ppar
-                vth = sqrt(2.0*pz[1]/density[1])
-            else
-                vth = nothing
-            end
-            @. vz.scratch2 = vpagrid_to_dzdt(vz.grid, vth, uz[1], evolve_ppar, evolve_upar)
-
-            # First apply boundary condition that total neutral outflux is equal to ion
-            # influx to uz
-            uz[1] = wall_flux_0 / density[1]
-            #would setting density work better??
-            #density[1] = - wall_flux_0 / uz[1]
-
-            # Create normalised Knudsen cosine distribution, to use for positive v_parallel
-            # at z = -Lz/2
-            # Note this only makes sense for the 1V case with vr.n=vzeta.n=1
-            @. vz.scratch = (3.0*pi/vtfac^3)*abs(vz.scratch2)*erfc(abs(vz.scratch2)/vtfac)
-
-            # The v_parallel>0 part of the pdf is replaced by the Knudsen cosine
-            # distribution. To ensure the constraints ∫dwpa wpa^m F = 0 are satisfied when
-            # necessary, calculate a normalisation factor for the Knudsen distribution (in
-            # vz.scratch) and correction terms for the incoming pdf similar to
-            # enforce_moment_constraints!().
-            #
-            # Note that it seems to be important that this boundary condition not be
-            # modified by the moment constraints, as that could cause numerical instability.
-            # By ensuring that the constraints are satisfied already here,
-            # enforce_moment_constraints!() will not change the pdf at the boundary. For
-            # ions this is not an issue, because points set to 0 by the bc are not modified
-            # from 0 by enforce_moment_constraints!().
-            knudsen_integral_0 = integrate_over_positive_vz(vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_1 = integrate_over_positive_vz(vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-
-            @views pdf_integral_0 = integrate_over_negative_vz(pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            @views pdf_integral_1 = integrate_over_negative_vz(vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            if !evolve_ppar
-                # Calculate normalisation factors N_in for the incoming and N_out for the
-                # Knudsen parts of the distirbution so that ∫dwpa F = 1 and ∫dwpa wpa F = 0
-                # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
-                #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = 0
-                N_in = 1.0 / (pdf_integral_0 - pdf_integral_1/knudsen_integral_1*knudsen_integral_0)
-                N_out = -N_in * pdf_integral_1 / knudsen_integral_1
-
-                zero_vz_ind = 0
-                for ivz ∈ 1:vz.n
-                    if vz.scratch2[ivz] <= -zero
-                        pdf[ivz,:,:,1] .= N_in*pdf[ivz,:,:,1]
-                    else
-                        zero_vz_ind = ivz
-                        if abs(vz.scratch2[ivz]) < zero
-                            # v_z = 0 point, half contribution from original pdf and half
-                            # from Knudsen cosine distribution, to be consistent with weights
-                            # used in
-                            # integrate_over_positive_vz()/integrate_over_negative_vz().
-                            @. pdf[ivz,:,:,1] = 0.5*(N_in*pdf[ivz,:,:,1] + N_out*vz.scratch[ivz])
-                        else
-                            pdf[ivz,:,:,1] .= N_out*vz.scratch[ivz]
-                        end
-                        break
-                    end
-                end
-                for ivz ∈ zero_vz_ind+1:vz.n
-                    pdf[ivz,:,:,1] .= N_out*vz.scratch[ivz]
-                end
-            else
-                knudsen_integral_2 = integrate_over_positive_vz(vz.grid .* vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                @views pdf_integral_2 = integrate_over_negative_vz(vz.grid .* vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                @views pdf_integral_3 = integrate_over_negative_vz(vz.grid .* vz.grid .* vz.grid .* pdf[:,:,:,1], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                # Calculate normalisation factor N_out for the Knudsen part of the
-                # distirbution and normalisation factor N_in and correction term C*wpa*F_in
-                # for the incoming distribution so that ∫dwpa F = 1, ∫dwpa wpa F = 0, and
-                # ∫dwpa wpa^2 F = 1/2
-                # ⇒ N_in*pdf_integral_0 + C*pdf_integral_1 + N_out*knudsen_integral_0 = 1
-                #   N_in*pdf_integral_1 + C*pdf_integral_2 + N_out*knudsen_integral_1 = 0
-                #   N_in*pdf_integral_2 + C*pdf_integral_3 + N_out*knudsen_integral_2 = 1/2
-                N_in = (0.5*knudsen_integral_0*pdf_integral_2 +
-                        knudsen_integral_1*(pdf_integral_3 - 0.5*pdf_integral_1) -
-                        knudsen_integral_2*pdf_integral_2) /
-                       (knudsen_integral_0*(pdf_integral_2^2 - pdf_integral_1*pdf_integral_3) +
-                        knudsen_integral_1*(pdf_integral_0*pdf_integral_3 - pdf_integral_1*pdf_integral_2) +
-                        knudsen_integral_2*(pdf_integral_1^2 - pdf_integral_0*pdf_integral_2))
-                N_out = -(N_in*(pdf_integral_1*pdf_integral_3 - pdf_integral_2^2) + 0.5*pdf_integral_2) /
-                         (knudsen_integral_1*pdf_integral_3 - knudsen_integral_2*pdf_integral_2)
-                C = (0.5 - N_out*knudsen_integral_2 - N_in*pdf_integral_2)/pdf_integral_3
-
-                zero_vz_ind = 0
-                for ivz ∈ 1:vz.n
-                    if vz.scratch2[ivz] <= -zero
-                        @views @. pdf[ivz,:,:,1] = N_in*pdf[ivz,:,:,1] + C*vz.grid[ivz]*pdf[ivz,:,:,1]
-                    else
-                        zero_vz_ind = ivz
-                        if abs(vz.scratch2[ivz]) < zero
-                            # v_parallel = 0 point, half contribution from original pdf and half
-                            # from Knudsen cosine distribution, to be consistent with weights
-                            # used in
-                            # integrate_over_positive_vz()/integrate_over_negative_vz().
-                            @views @. pdf[ivz,:,:,1] = 0.5*(N_in*pdf[ivz,:,:,1] +
-                                                            C*vz.grid[ivz]*pdf[ivz,:,:,1] +
-                                                            N_out*vz.scratch[ivz])
-                        else
-                            @. pdf[ivz,:,:,1] = N_out*vz.scratch[ivz]
-                        end
-                        break
-                    end
-                end
-                for ivz ∈ zero_vz_ind+1:vz.n
-                    @. pdf[ivz,:,:,1] = N_out*vz.scratch[ivz]
-                end
-            end
-        end
-
-        if z.irank == z.nrank - 1
-            ## treat the right boundary at z = Lz/2 ##
-            # populate vz.scratch2 array with dz/dt values at z = Lz/2
-            if evolve_ppar
-                vth = sqrt(2.0*pz[end]/density[end])
-            else
-                vth = nothing
-            end
-            @. vz.scratch2 = vpagrid_to_dzdt(vz.grid, vth, uz[end], evolve_ppar, evolve_upar)
-
-            # First apply boundary condition that total neutral outflux is equal to ion
-            # influx to uz
-            uz[end] = - wall_flux_L / density[end]
-            #would setting density work better??
-            #density[end] = - wall_flux_L / upar[end]
-
-            # obtain the Knudsen cosine distribution at z = Lz/2
-            # the z-dependence is only introduced if the peculiiar velocity is used as vz
-            # Note this only makes sense for the 1V case with vr.n=vzeta.n=1
-            @. vz.scratch = (3.0*pi/vtfac^3)*abs(vz.scratch2)*erfc(abs(vz.scratch2)/vtfac)
-
-            # The v_parallel<0 part of the pdf is replaced by the Knudsen cosine
-            # distribution. To ensure the constraint ∫dwpa wpa F = 0 is satisfied, multiply
-            # the Knudsen distribution (in vz.scratch) by a normalisation factor given by
-            # the integral (over negative v_parallel) of the outgoing Knudsen distribution
-            # and (over positive v_parallel) of the incoming pdf.
-            knudsen_integral_0 = integrate_over_negative_vz(vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            knudsen_integral_1 = integrate_over_negative_vz(vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-
-            @views pdf_integral_0 = integrate_over_positive_vz(pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-            @views pdf_integral_1 = integrate_over_positive_vz(vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-
-            if !evolve_ppar
-                # Calculate normalisation factors N_in for the incoming and N_out for the
-                # Knudsen parts of the distirbution so that ∫dwpa F = 1 and ∫dwpa wpa F = 0
-                # ⇒ N_in*pdf_integral_0 + N_out*knudsen_integral_0 = 1
-                #   N_in*pdf_integral_1 + N_out*knudsen_integral_1 = 0
-                N_in = 1.0 / (pdf_integral_0 - pdf_integral_1/knudsen_integral_1*knudsen_integral_0)
-                N_out = -N_in * pdf_integral_1 / knudsen_integral_1
-
-                zero_vz_ind = 0
-                for ivz ∈ vz.n:-1:1
-                    if vz.scratch2[ivz] >= zero
-                        @views @. pdf[ivz,:,:,end] = N_in*pdf[ivz,:,:,end]
-                    else
-                        zero_vz_ind = ivz
-                        if abs(vz.scratch2[ivz]) < zero
-                            # v_parallel = 0 point, half contribution from original pdf and half
-                            # from Knudsen cosine distribution, to be consistent with weights
-                            # used in
-                            # integrate_over_positive_vz()/integrate_over_negative_vz().
-                            @views @. pdf[ivz,:,:,end] = 0.5*(N_in*pdf[ivz,:,:,end] + N_out*vz.scratch[ivz])
-                        else
-                            @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
-                        end
-                        break
-                    end
-                end
-                for ivz ∈ 1:zero_vz_ind-1
-                    @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
-                end
-            else
-                knudsen_integral_2 = integrate_over_negative_vz(vz.grid .* vz.grid .* vz.scratch, vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                @views pdf_integral_2 = integrate_over_positive_vz(vz.grid .* vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                @views pdf_integral_3 = integrate_over_positive_vz(vz.grid .* vz.grid .* vz.grid .* pdf[:,:,:,end], vz.scratch2, vz.wgts, vz.scratch3, vr.grid, vr.wgts, vzeta.grid, vzeta.wgts)
-                # Calculate normalisation factor N_out for the Knudsen part of the
-                # distirbution and normalisation factor N_in and correction term C*wpa*F_in
-                # for the incoming distribution so that ∫dwpa F = 1, ∫dwpa wpa F = 0, and
-                # ∫dwpa wpa^2 F = 1/2
-                # ⇒ N_in*pdf_integral_0 + C*pdf_integral_1 + N_out*knudsen_integral_0 = 1
-                #   N_in*pdf_integral_1 + C*pdf_integral_2 + N_out*knudsen_integral_1 = 0
-                #   N_in*pdf_integral_2 + C*pdf_integral_3 + N_out*knudsen_integral_2 = 1/2
-                N_in = (0.5*knudsen_integral_0*pdf_integral_2 +
-                        knudsen_integral_1*(pdf_integral_3 - 0.5*pdf_integral_1) -
-                        knudsen_integral_2*pdf_integral_2) /
-                       (knudsen_integral_0*(pdf_integral_2^2 - pdf_integral_1*pdf_integral_3) +
-                        knudsen_integral_1*(pdf_integral_0*pdf_integral_3 - pdf_integral_1*pdf_integral_2) +
-                        knudsen_integral_2*(pdf_integral_1^2 - pdf_integral_0*pdf_integral_2))
-                N_out = -(N_in*(pdf_integral_1*pdf_integral_3 - pdf_integral_2^2) + 0.5*pdf_integral_2) /
-                         (knudsen_integral_1*pdf_integral_3 - knudsen_integral_2*pdf_integral_2)
-                C = (0.5 - N_out*knudsen_integral_2 - N_in*pdf_integral_2)/pdf_integral_3
-
-                zero_vz_ind = 0
-                for ivz ∈ vz.n:-1:1
-                    if vz.scratch2[ivz] >= zero
-                        @views @. pdf[ivz,:,:,end] = N_in*pdf[ivz,:,:,end] + C*vz.grid[ivz]*pdf[ivz,:,:,end]
-                    else
-                        zero_vz_ind = ivz
-                        if abs(vz.scratch2[ivz]) < zero
-                            # v_parallel = 0 point, half contribution from original pdf and half
-                            # from Knudsen cosine distribution, to be consistent with weights
-                            # used in
-                            # integrate_over_positive_vz()/integrate_over_negative_vz().
-                            @views @. pdf[ivz,:,:,end] = 0.5*(N_in*pdf[ivz,:,:,end] +
-                                                              C*vz.grid[ivz]*pdf[ivz,:,:,end] +
-                                                              N_out*vz.scratch[ivz])
-                        else
-                            @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
-                        end
-                        break
-                    end
-                end
-                for ivz ∈ 1:zero_vz_ind-1
-                    @. pdf[ivz,:,:,end] = N_out*vz.scratch[ivz]
-                end
-            end
-        end
-    end
-end
-
-"""
-create an array of dz/dt values corresponding to the given vpagrid values
-"""
-function vpagrid_to_dzdt(vpagrid, vth, upar, evolve_ppar, evolve_upar)
-    if evolve_ppar
-        if evolve_upar
-            return vpagrid .* vth .+ upar
-        else
-            return vpagrid .* vth
-        end
-    elseif evolve_upar
-        return vpagrid .+ upar
-    else
-        return vpagrid
-    end
-end
-
 """
 Take the full ion distribution function, calculate the moments, then
 normalise and shift to the moment-kinetic grid.
@@ -2452,94 +1558,4 @@ function convert_full_f_neutral_to_normalised!(f, density, uz, pz, vth, vzeta, v
     return nothing
 end
 
-"""
-enforce the z boundary condition on the evolved velocity space moments of f
-"""
-function enforce_z_boundary_condition_moments!(density, moments, bc::String)
-    ## TODO: parallelise
-    #begin_serial_region()
-    #@serial_region begin
-    #    # enforce z boundary condition on density if it is evolved separately from f
-    #	if moments.evolve_density
-    #        # TODO: extend to 'periodic' BC case, as this requires further code modifications to be consistent
-    #        # with finite difference derivatives (should be fine for Chebyshev)
-    #        if bc == "wall"
-    #            @loop_s_r is ir begin
-    #                density[1,ir,is] = 0.5*(density[1,ir,is] + density[end,ir,is])
-    #                density[end,ir,is] = density[1,ir,is]
-    #        	end
-    #        end
-    #    end
-    #end
-end
-
-"""
-"""
-function enforce_v_boundary_condition_local!(f, bc, speed, v_diffusion, v, v_spectral)
-    if bc == "zero"
-        if v_diffusion || speed[1] > 0.0
-            # 'upwind' boundary
-            f[1] = 0.0
-        end
-        if v_diffusion || speed[end] < 0.0
-            # 'upwind' boundary
-            f[end] = 0.0
-        end
-    elseif bc == "both_zero"
-        f[1] = 0.0
-        f[end] = 0.0
-    elseif bc == "zero_gradient"
-        D0 = v_spectral.lobatto.Dmat[1,:]
-        # adjust F(vpa = -L/2) so that d F / d vpa = 0 at vpa = -L/2
-        f[1] = -sum(D0[2:v.ngrid].*f[2:v.ngrid])/D0[1]
-
-        D0 = v_spectral.lobatto.Dmat[end,:]
-        # adjust F(vpa = L/2) so that d F / d vpa = 0 at vpa = L/2
-        f[end] = -sum(D0[1:ngrid-1].*f[end-v.ngrid+1:end-1])/D0[v.ngrid]
-    elseif bc == "periodic"
-        f[1] = 0.5*(f[1]+f[end])
-        f[end] = f[1]
-    else
-        error("Unsupported boundary condition option '$bc' for $(v.name)")
-    end
-end
-
-"""
-enforce zero boundary condition at vperp -> infinity
-"""
-function enforce_vperp_boundary_condition! end
-
-function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vperp, vperp_spectral)
-    @loop_s is begin
-        @views enforce_vperp_boundary_condition!(f[:,:,:,:,is], bc, vperp, vperp_spectral)
-    end
-    return nothing
-end
-
-function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral)
-    if bc == "zero"
-        nvperp = vperp.n
-        ngrid = vperp.ngrid
-        # set zero boundary condition
-        @loop_r_z_vpa ir iz ivpa begin
-            f[ivpa,nvperp,iz,ir] = 0.0
-        end
-        # set regularity condition d F / d vperp = 0 at vperp = 0
-        if vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral"
-            D0 = vperp_spectral.radau.D0
-            @loop_r_z_vpa ir iz ivpa begin
-                # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
-                f[ivpa,1,iz,ir] = -sum(D0[2:ngrid].*f[ivpa,2:ngrid,iz,ir])/D0[1]
-            end
-        else
-            println("vperp.bc=\"$bc\" not supported by discretization "
-                    * "$(vperp.discretization)")
-        end
-    elseif bc == "none"
-        # Do nothing
-    else
-        error("Unsupported boundary condition option '$bc' for vperp")
-    end
-end
-
 end
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index aa2b20f07..1a47f9fcd 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -53,6 +53,7 @@ include("electron_vpa_advection.jl")
 include("neutral_r_advection.jl")
 include("neutral_z_advection.jl")
 include("neutral_vz_advection.jl")
+include("boundary_conditions.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("krook_collisions.jl")
@@ -84,7 +85,7 @@ using .debugging
 using .external_sources
 using .input_structs
 using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!,
-                           enforce_boundary_conditions!, initialize_scratch_arrays!
+                           initialize_scratch_arrays!
 using .load_data: reload_evolving_fields!
 using .looping
 using .moment_constraints: hard_force_moment_constraints!
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 43a98a4cc..ccd0b1d31 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -27,8 +27,8 @@ using ..velocity_moments: calculate_moment_derivatives!, calculate_moment_deriva
 using ..velocity_moments: calculate_electron_moment_derivatives!
 using ..velocity_moments: update_chodura!
 using ..velocity_grid_transforms: vzvrvzeta_to_vpavperp!, vpavperp_to_vzvrvzeta!
-using ..initial_conditions: enforce_boundary_conditions!
-using ..initial_conditions: enforce_neutral_boundary_conditions!
+using ..boundary_conditions: enforce_boundary_conditions!
+using ..boundary_conditions: enforce_neutral_boundary_conditions!
 using ..input_structs: advance_info, time_input
 using ..moment_constraints: hard_force_moment_constraints!,
                             hard_force_moment_constraints_neutral!

From a5d7c98f41a8157fd2fd84d35eec3edf6974f705 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 25 Feb 2024 15:41:49 +0000
Subject: [PATCH 133/428] Enforce velocity-space boundary conditions on
 electron pdf

---
 .../wall+sheath-bc_kinetic.toml               |  4 +-
 .../src/electron_kinetic_equation.jl          | 34 +++++++++++--
 moment_kinetics/src/initial_conditions.jl     | 51 ++++++++++++-------
 moment_kinetics/src/moment_kinetics.jl        |  4 +-
 4 files changed, 67 insertions(+), 26 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index faad6f68c..bdbdf5ae6 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -67,13 +67,13 @@ vpa_ngrid = 17
 #vpa_nelement = 10
 vpa_nelement = 20
 vpa_L = 12.0 #8.0
-vpa_bc = "periodic"
+vpa_bc = "zero"
 vpa_discretization = "chebyshev_pseudospectral"
 #vpa_discretization = "gausslegendre_pseudospectral"
 vz_ngrid = 17
 vz_nelement = 10
 vz_L = 8.0
-vz_bc = "periodic"
+vz_bc = "zero"
 vz_discretization = "chebyshev_pseudospectral"
 
 [output]
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 4a7020884..915b50e01 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -8,6 +8,8 @@ export get_electron_critical_velocities
 
 using ..looping
 using ..derivatives: derivative_z!
+using ..boundary_conditions: enforce_v_boundary_condition_local!,
+                             enforce_vperp_boundary_condition!
 using ..calculus: derivative!, second_derivative!, integral
 using ..communication
 using ..interpolation: interpolate_to_grid_1d!
@@ -50,7 +52,7 @@ OUTPUT:
 """
 function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
         phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
-        vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
+        vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
         num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing,
         initial_time=0.0)
 
@@ -62,7 +64,7 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_u
     if solution_method == "artificial_time_derivative"
         return update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
             moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
-            r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
+            r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
             num_diss_params, max_electron_pdf_iterations;
             io_initial_electron=io_initial_electron, initial_time=initial_time)
     elseif solution_method == "shooting_method"
@@ -109,7 +111,7 @@ OUTPUT:
 """
 function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
     moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition,
-    r, z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
+    r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
     num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0)
 
     begin_r_z_region()
@@ -334,7 +336,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
 
         # enforce the boundary condition(s) on the electron pdf
-        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, composition.me_over_mi)
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
+                                                    vperp_spectral, vpa_spectral,
+                                                    vpa_advect,
+                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                    composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
 
@@ -489,7 +495,25 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     return time, output_counter
 end
 
-function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
+function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
+                                                     vperp_spectral, vpa_spectral,
+                                                     vpa_adv, vpa_diffusion, me_over_mi)
+    # Enforce velocity-space boundary conditions
+    if vpa.n > 1
+        begin_r_z_vperp_region()
+        @loop_r_z_vperp ir iz ivperp begin
+            # enforce the vpa BC
+            # use that adv.speed independent of vpa
+            @views enforce_v_boundary_condition_local!(pdf[:,ivperp,iz,ir], vpa.bc,
+                                                       vpa_adv[1].speed[:,ivperp,iz,ir],
+                                                       vpa_diffusion, vpa, vpa_spectral)
+        end
+    end
+    if vperp.n > 1
+        begin_r_z_vpa_region()
+        @views enforce_vperp_boundary_condition!(pdf, vperp.bc, vperp, vperp_spectral)
+    end
+
     # first enforce the boundary condition at z_min.
     # this involves forcing the pdf to be zero for electron travelling faster than the max speed
     # they could attain by accelerating in the electric field between the wall and the simulation boundary;
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 74dec52d7..4944d124e 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -175,10 +175,11 @@ creates the normalised pdfs and the velocity-space moments and populates them
 with a self-consistent initial condition
 """
 function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry, composition, r, z,
-                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, vz_spectral,
-                               species, collisions, external_source_settings, manufactured_solns_input,
-                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
-                               io_input, input_dict)
+                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                               vperp_spectral, vpa_spectral, vz_spectral, species,
+                               collisions, external_source_settings,
+                               manufactured_solns_input, scratch_dummy, scratch, t_input,
+                               num_diss_params, advection_structs, io_input, input_dict)
     if manufactured_solns_input.use_for_init
         init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z,
                                              r, composition.n_ion_species,
@@ -310,9 +311,16 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
         scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
     if composition.electron_physics == kinetic_electrons
         # Initialise the array for the electron pdf
+        begin_serial_region()
+        speed = @view scratch_dummy.buffer_vpavperpzrs_1[:,:,:,:,1]
+        @serial_region begin
+            speed .= 0.0
+        end
         init_electron_pdf_over_density_and_boundary_phi!(
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
-            moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
+            moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
+            [(speed=speed,)], num_diss_params,
+            composition.me_over_mi)
     end
     # calculate the electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
@@ -339,7 +347,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
 
     # initialize the electron pdf that satisfies the electron kinetic equation
     initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
-                             vzeta, vr, vz, z_spectral, vpa_spectral,
+                             vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
@@ -467,10 +475,10 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
 end
 
 function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
-                                  vz, z_spectral, vpa_spectral, z_advect, vpa_advect,
-                                  scratch_dummy, collisions, composition, geometry,
-                                  external_source_settings, num_diss_params, dt, io_input,
-                                  input_dict)
+                                  vz, z_spectral, vperp_spectral, vpa_spectral, z_advect,
+                                  vpa_advect, scratch_dummy, collisions, composition,
+                                  geometry, external_source_settings, num_diss_params, dt,
+                                  io_input, input_dict)
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
@@ -546,9 +554,10 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                         moments.electron.dppar_dz,
                                         moments.electron.dqpar_dz,
                                         moments.electron.dvth_dz, r, z, vperp, vpa,
-                                        z_spectral, vpa_spectral, z_advect, vpa_advect,
-                                        scratch_dummy, dt, collisions, composition,
-                                        num_diss_params, max_electron_pdf_iterations;
+                                        z_spectral, vperp_spectral, vpa_spectral,
+                                        z_advect, vpa_advect, scratch_dummy, dt,
+                                        collisions, composition, num_diss_params,
+                                        max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
                                         initial_time=code_time)
 
@@ -1222,9 +1231,10 @@ NB: as the electron pdf is obtained via a time-independent equation,
 this 'initital' value for the electron will just be the first guess in an iterative solution
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
-        vpa, vperp, vpa_spectral, me_over_mi)
+        vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, num_diss_params, me_over_mi)
 
     if z.bc == "wall"
+        begin_r_region()
         @loop_r ir begin
             # Initialise an unshifted Maxwellian as a first step
             @loop_z iz begin
@@ -1233,9 +1243,16 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
                     @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
                 end
             end
-            # Apply the sheath boundary condition to get cut-off boundary distribution
-            # functions and boundary values of phi
-            enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vpa, vpa_spectral, me_over_mi)
+        end
+        # Apply the sheath boundary condition to get cut-off boundary distribution
+        # functions and boundary values of phi
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
+                                                    vperp_spectral, vpa_spectral,
+                                                    vpa_advect,
+                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                    me_over_mi)
+        begin_r_region()
+        @loop_r ir begin
             # get critical velocities beyond which electrons are lost to the wall
             #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
             #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 1a47f9fcd..d87917cee 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -276,8 +276,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         # each of which may be evolved separately depending on input choices.
         init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geometry,
                               composition, r, z, vperp, vpa, vzeta, vr, vz,
-                              z_spectral, r_spectral, vpa_spectral, vz_spectral, species,
-                              collisions, external_source_settings,
+                              z_spectral, r_spectral, vperp_spectral, vpa_spectral,
+                              vz_spectral, species, collisions, external_source_settings,
                               manufactured_solns_input, scratch_dummy, scratch, t_input,
                               num_diss_params, advection_structs, io_input, input_dict)
         # initialize time variable

From b0e851f17430736e85241e9a960ec54c2f0fdd6f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 25 Feb 2024 21:56:33 +0000
Subject: [PATCH 134/428] Move 'speedup hacks' to a function

---
 .../src/electron_kinetic_equation.jl          | 55 +++++++++----------
 1 file changed, 27 insertions(+), 28 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 915b50e01..1743b7659 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -194,19 +194,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                                         ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
                                         z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                         collisions, num_diss_params, dt_max)
-    # Divide by wpa to relax CFL condition at large wpa - only looking for steady
-    # state here, so does not matter that this makes time evolution incorrect.
-    # Also increase the effective timestep for z-values far from the sheath boundary -
-    # these have a less-limited timestep so letting them evolve faster speeds up
-    # convergence to the steady state.
-    Lz = z.L
-    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-        zval = z.grid[iz]
-        znorm = 2.0*zval/Lz
-        residual[ivpa,ivperp,iz,ir] *=
-            (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
-            sqrt(1.0 + vpa.grid[ivpa]^2)
-    end
+    speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
     if n_blocks[] == 1
         text_output_suffix = ""
     else
@@ -302,7 +290,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             # but without the wpa dependence.
             @loop_r_z ir iz begin
                 zval = z.grid[iz]
-                znorm = 2.0*zval/Lz
+                znorm = 2.0*zval/z.L
                 ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
                               (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
                               (1.0 + z_speedup_fac*(1.0 - znorm^2))
@@ -433,20 +421,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
         average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
 
-        # Divide by wpa to relax CFL condition at large wpa - only looking for steady
-        # state here, so does not matter that this makes time evolution incorrect.
-        # Also increase the effective timestep for z-values far from the sheath boundary -
-        # these have a less-limited timestep so letting them evolve faster speeds up
-        # convergence to the steady state.
-        begin_r_z_vperp_region()
-        Lz = z.L
-        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-            zval = z.grid[iz]
-            znorm = 2.0*zval/Lz
-            residual[ivpa,ivperp,iz,ir] *=
-                (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
-                sqrt(1.0 + vpa.grid[ivpa]^2)
-        end
+        speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
         if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
             break
         end
@@ -495,6 +470,30 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     return time, output_counter
 end
 
+function speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
+    # Divide by wpa to relax CFL condition at large wpa - only looking for steady
+    # state here, so does not matter that this makes time evolution incorrect.
+    # Also increase the effective timestep for z-values far from the sheath boundary -
+    # these have a less-limited timestep so letting them evolve faster speeds up
+    # convergence to the steady state.
+
+    # Actually modify so that large wpa does go faster (to allow some phase mixing - maybe
+    # this makes things more stable?), but not by so much.
+    #vpa_fudge_factor = 1.0
+    #vpa_fudge_factor = 0.8
+    vpa_fudge_factor = 0.0
+
+    Lz = z.L
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        zval = z.grid[iz]
+        znorm = 2.0*zval/Lz
+        residual[ivpa,ivperp,iz,ir] *=
+            (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
+            sqrt(1.0 + vpa_fudge_factor * vpa.grid[ivpa]^2)
+    end
+    return nothing
+end
+
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, vpa_diffusion, me_over_mi)

From 024f419de740e47a75f652f092ec7b72924fb2eb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 25 Feb 2024 21:57:26 +0000
Subject: [PATCH 135/428] Include upar in electron_z_advection

This should be a negligibly small contribution, but might be more
consistent with the electron boundary condition near vpa=0?
---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 moment_kinetics/src/electron_z_advection.jl      | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 1743b7659..703bcf7b6 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1256,7 +1256,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
         residual[ivpa,ivperp,iz,ir] = 0.0
     end
     # calculate the contribution to the residual from the z advection term
-    electron_z_advection!(residual, pdf, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
+    electron_z_advection!(residual, pdf, upar, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
     #dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
     #single_term .= residual
     #max_term .= abs.(residual)
diff --git a/moment_kinetics/src/electron_z_advection.jl b/moment_kinetics/src/electron_z_advection.jl
index f987ce028..10f1e8b2a 100644
--- a/moment_kinetics/src/electron_z_advection.jl
+++ b/moment_kinetics/src/electron_z_advection.jl
@@ -14,13 +14,13 @@ using ..calculus: second_derivative!, derivative!
 """
 calculate the z-advection term for the electron kinetic equation = wpa * vthe * df/dz
 """
-function electron_z_advection!(advection_term, pdf, vth, advect, z, vpa, spectral, scratch_dummy)
+function electron_z_advection!(advection_term, pdf, upar, vth, advect, z, vpa, spectral, scratch_dummy)
     # create a pointer to a scratch_dummy array to store the z-derivative of the electron pdf
     dpdf_dz = scratch_dummy.buffer_vpavperpzr_1
     d2pdf_dz2 = scratch_dummy.buffer_vpavperpzr_2
     begin_r_vperp_vpa_region()
     # get the updated speed along the z direction using the current pdf
-    @views update_electron_speed_z!(advect[1], vth, vpa)
+    @views update_electron_speed_z!(advect[1], upar, vth, vpa)
     # update adv_fac -- note that there is no factor of dt here because
     # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
     @loop_r_vperp_vpa ir ivperp ivpa begin
@@ -46,10 +46,11 @@ end
 """
 calculate the electron advection speed in the z-direction at each grid point
 """
-function update_electron_speed_z!(advect, vth, vpa)
+function update_electron_speed_z!(advect, upar, vth, vpa)
     # the electron advection speed in z is v_par = w_par * v_the
     @loop_r_vperp_vpa ir ivperp ivpa begin
-        @. @views advect.speed[:,ivpa,ivperp,ir] = vpa[ivpa] * vth[:,ir]
+        #@. @views advect.speed[:,ivpa,ivperp,ir] = vpa[ivpa] * vth[:,ir]
+        @. @views advect.speed[:,ivpa,ivperp,ir] = vpa[ivpa] * vth[:,ir] + upar[:,ir]
     end
     return nothing
 end

From 6aa103abdb54939b06626d1e4016fe394866b87c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 26 Feb 2024 12:27:44 +0000
Subject: [PATCH 136/428] Restore `using Glob` in top-level moment_kinetics
 module

This is required by the debug check test scripts.
---
 moment_kinetics/src/moment_kinetics.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index d87917cee..b12fb13a9 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -73,6 +73,7 @@ include("time_advance.jl")
 
 using TimerOutputs
 using Dates
+using Glob
 using Primes
 
 using .file_io: setup_file_io, finish_file_io

From abfd6fddbb33fd8f95e15c5fa03c6d723e3136ba Mon Sep 17 00:00:00 2001
From: Michael Barnes <michael.barnes@physics.ox.ac.uk>
Date: Mon, 26 Feb 2024 12:29:30 +0000
Subject: [PATCH 137/428] added ability to initialise a kinetic electron
 simulation from the end of a simulation with Boltzmann electrons; also added
 in division by vthe in the definition of the average_residual; other
 (possibly temporary) changes to probe behaviour of electron pdf at wall
 boundaries.

---
 .../src/electron_kinetic_equation.jl          | 578 +++++++++++++++++-
 moment_kinetics/src/initial_conditions.jl     | 247 ++++++--
 moment_kinetics/src/load_data.jl              |   6 +-
 moment_kinetics/src/moment_kinetics.jl        |  18 +-
 4 files changed, 768 insertions(+), 81 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 5caef372f..f9c906dfa 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -160,11 +160,11 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
 
     #dt_electron = dt * sqrt(composition.me_over_mi)
     #dt_max = 3.0e-8
-    dt_max = 1.0e-8
+    dt_max = 2.0e-8
     #dt_max = 2.5e-9
 
     #dt_energy = 1.0e-7
-    dt_energy = 1.0e-8
+    dt_energy = 2.0e-8
     #dt_energy = 2.5e-9
     #n_ppar_subcycles = 1000
     #n_ppar_subcycles = 200
@@ -290,18 +290,19 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                 end
             end
 
-            #dt_energy = dt_electron
-            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
-
-            # Apply same 'speed up' hack to ppar that we do to the distribution function,
-            # but without the wpa dependence.
-            @loop_r_z ir iz begin
-                zval = z.grid[iz]
-                znorm = 2.0*zval/Lz
-                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
-                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
-                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
-            end
+            # TMP FOR TESTING -- MAB
+            # #dt_energy = dt_electron
+            # electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
+
+            # # Apply same 'speed up' hack to ppar that we do to the distribution function,
+            # # but without the wpa dependence.
+            # @loop_r_z ir iz begin
+            #     zval = z.grid[iz]
+            #     znorm = 2.0*zval/Lz
+            #     ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
+            #                   (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
+            #                   (1.0 + z_speedup_fac*(1.0 - znorm^2))
+            # end
             begin_r_z_region()
             @loop_r_z ir iz begin
                 fvec.electron_ppar[iz,ir] = ppar[iz,ir]
@@ -331,7 +332,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
 
         # enforce the boundary condition(s) on the electron pdf
-        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, composition.me_over_mi)
+        enforce_boundary_condition_on_electron_pdf_noshift!(pdf, phi, vthe, upar, vpa, vpa_spectral, composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
 
@@ -392,6 +393,15 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         if (mod(iteration,output_interval) == 0)
             begin_serial_region()
             @serial_region begin
+                if (mod(iteration,100*output_interval) == 0)
+                    @loop_vpa ivpa begin
+                        @loop_z iz begin
+                            println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                        end
+                        println(io_pdf,"")
+                    end
+                    println(io_pdf,"")
+                end
                 @loop_z iz begin
                     println(io_upar, "z: ", z.grid[iz], " upar: ", upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
                     println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
@@ -486,9 +496,539 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     return time, output_counter
 end
 
+function enforce_boundary_condition_on_electron_pdf_noshift!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
+    # first enforce the boundary condition at z_min.
+    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
+    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
+    # for electrons with positive velocities less than this critical value, they must have the same
+    # pdf as electrons with negative velocities of the same magnitude.
+    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
+    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
+
+    begin_r_region()
+
+    # pdf_adjustment_option determines the velocity-dependent pre-factor for the
+    # corrections to the pdf needed to ensure moment constraints are satisfied
+    #pdf_adjustment_option = "vpa4"
+    #pdf_adjustment_option = "vpa4_gaussian"
+    pdf_adjustment_option = "no1st_vpa2"
+
+    cutoff_step_width = 0.1
+
+    # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
+    #wpa_values = vpa.scratch
+    # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
+    #interpolated_pdf = vpa.scratch2
+    reversed_pdf = vpa.scratch
+
+    # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
+    #ivpa_zero = (vpa.n+1)÷2
+
+    @loop_r ir begin
+        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+        #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
+        #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+
+        # Want to construct the w-grid corresponding to -vpa.
+        #   wpa(vpa) = (vpa - upar)/vth
+        #   ⇒ vpa = vth*wpa(vpa) + upar
+        #   wpa(-vpa) = (-vpa - upar)/vth
+        #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+        #             = (-vth*wpa - 2*upar)/vth
+        #             = -wpa - 2*upar/vth
+        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+        # reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+        # #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
+        # # interpolate the pdf onto this grid
+        # #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+        # @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+        # reverse!(reversed_pdf)
+        reversed_pdf .= pdf[:,1,1,ir]
+        # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
+        #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
+        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+        #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
+        # interpolate back onto the original wpa grid
+        #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+        # construct wpa * pdf
+        #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
+        # calculate the first moment of the normalised pdf
+        #first_vspace_moment = 0.0
+        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
+        ivpa = 0
+        ivpa_max = vpa.n + 1
+        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+        # if the first moment is positive, then the cutoff speed needs to be reduced
+        upar_over_vth = upar[1,ir] / vthe[1,ir]
+        #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
+        #println("$first_vspace_moment, u/vth=$upar_over_vth")
+        vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid
+        upar_integral = 0.0
+        #while first_vspace_moment > upar_over_vth # > 0.0
+        #    # zero out the pdf at the current cutoff velocity
+        #    pdf[ivpa_max,1,1,ir] = 0.0
+        #    # update wpa * pdf
+        #    vpa.scratch3[ivpa_max] = 0.0
+        #    # calculate the updated first moment of the normalised pdf
+        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
+        #    if first_vspace_moment > upar_over_vth #0.0
+        #        ivpa_max -= 1
+        #    end
+        #end
+        upar0 = upar[1,ir]
+        #println("before pdf left ", pdf[:,1,1,ir])
+        while upar_integral > upar0 && ivpa_max > 1
+            ivpa += 1
+            ivpa_max -= 1
+            # zero out the reversed pdf at the current cutoff velocity
+            #reversed_pdf[ivpa_max] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
+            #println("left ", ivpa, " ", upar_integral, " ", upar0)
+        end
+        integral_excess = upar_integral - upar0
+        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
+        #println("fraction_of_pdf=", fraction_of_pdf)
+
+        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+        # integral from the point at ivpa is required to make upar_integral<upar0) the
+        # cut-off velocity is half way between ivpa and ivpa+1, while when
+        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+        # ivpa is required to make upar_integral<upar0) the cut-off is half way between
+        # ivpa-1 and ivpa.
+        vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
+               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
+
+        #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
+        #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
+        #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
+        #@loop_vpa ivpa begin
+        #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
+        #end
+        reversed_pdf[ivpa_max+1:end] .= 0
+        reversed_pdf[ivpa_max] *= fraction_of_pdf
+        #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
+        #println("done first cutoff loop")
+        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+        #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+        phi[1,ir] = me_over_mi * vmax^2
+        iv0 = findfirst(x -> x>0.0, vpa_unnorm)
+        pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
+        #println("check reversed change ", reversed_pdf[iv0:end])
+        #println("reversed_pdf ", reversed_pdf)
+        #println("after pdf left ", pdf[:,1,1,ir])
+        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+        # of the modified pdf are 1, 0 and 1/2 respectively
+        # will need vpa / vthe = wpa + upar/vthe
+        @. vpa.scratch2 = vpa.grid
+        # first need to calculate int dwpa pdf = zeroth_moment
+        zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
+        # calculate int dwpa wpa^2 * pdf = wpa2_moment
+        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
+        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * pdf = vpa2_moment
+        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
+        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
+        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
+        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        
+        if pdf_adjustment_option == "absvpa"
+            # calculate int dwpa |vpa| * pdf = absvpa_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
+            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            # calculate the 'B' normalisation constant
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+            # calculate the 'A' normalisation constant
+            normalisation_constant_A = (1 + normalisation_constant_B * 
+                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+            # calculate the 'C' normalisation constant
+            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                + vpa.scratch2^2 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4"
+            # calculate int dwpa vpa^4 * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                + vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4_gaussian"
+            afac = 0.1
+            bfac = 0.2
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #                           (vpa2_wpa2_moment
+            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+            #normalisation_constant_A = (1 - normalisation_constant_B 
+            #                                * vpa2_moment) / zeroth_moment
+            #normalisation_constant_C = 0.0
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "no1st_vpa2"
+            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
+            # TMP FOR TESTING -- MAB
+            #@. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
+        else
+            println("pdf_adjustment_option not recognised")
+            stop()
+        end
+
+        # smooth the pdf at the boundary
+        #for ivpa ∈ 2:ivpa_max-1
+        #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
+        #end
+    end
+
+    # next enforce the boundary condition at z_max.
+    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
+    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
+    # for electrons with negative velocities less than this critical value, they must have the same
+    # pdf as electrons with positive velocities of the same magnitude.
+    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
+    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
+    
+    # io_pdf_stages = open("pdf_stages.txt", "w")
+    # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+    # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid^2
+    # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+    # @loop_vpa ivpa begin
+    #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+    #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 0)
+    # end
+    # println(io_pdf_stages,"")
+
+    @loop_r ir begin
+        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+        #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+
+        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+        # Need to reverse vpa.grid because the grid passed as the second argument of
+        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+        # reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+        # #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
+        # # interpolate the pdf onto this grid
+        # #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+        # @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+        # # reverse!(reversed_pdf)
+        reversed_pdf .= pdf[:,1,end,ir]
+        # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
+        #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
+        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+        #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
+        # interpolate back onto the original wpa grid
+        #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
+        # end
+        # println(io_pdf_stages,"")
+
+        # construct wpa * pdf
+        #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
+        # calculate the first moment of the normalised pdf
+        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
+        ivpa = vpa.n+1
+        ivpa_min = 0
+        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+        # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
+        upar_over_vth = upar[end,ir] / vthe[end,ir]
+        #println("$first_vspace_moment, u/vth=$upar_over_vth")
+        vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid
+        upar_integral = 0.0
+        #while first_vspace_moment < upar_over_vth # < 0.0
+        #    # zero out the pdf at the current cutoff velocity
+        #    pdf[ivpa_min,1,end,ir] = 0.0
+        #    # update wpa * pdf
+        #    vpa.scratch3[ivpa_min] = 0.0
+        #    # calculate the updated first moment of the normalised pdf
+        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        #    if first_vspace_moment < upar_over_vth
+        #        ivpa_min += 1
+        #    end
+        #end
+        upar_end = upar[end,ir]
+        #println("before pdf ", pdf[:,1,end,ir])
+        while upar_integral < upar_end && ivpa > 1
+            ivpa -= 1
+            ivpa_min += 1
+            # zero out the reversed pdf at the current cutoff velocity
+            #reversed_pdf[ivpa_min] = 0.0
+            # calculate the updated first moment of the normalised pdf
+            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
+            #println("right ", ivpa, " ", upar_integral, " ", upar_end)
+        end
+        integral_excess = upar_integral - upar_end
+        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
+        #println("B fraction_of_pdf=", fraction_of_pdf)
+
+        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+        # integral from the point at ivpa is required to make upar_integral>upar_end) the
+        # cut-off velocity is half way between ivpa-1 and ivpa, while when
+        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+        # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
+        # ivpa and ivpa+1.
+        vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
+               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
+
+        #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
+        #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
+        #@loop_vpa ivpa begin
+        #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
+        #end
+        reversed_pdf[1:ivpa_min-1] .= 0
+        reversed_pdf[ivpa_min] *= fraction_of_pdf
+        #println("done second cutoff loop")
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
+        # end
+        # println(io_pdf_stages,"")
+
+        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+        #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
+        phi[end,ir] = me_over_mi * vmin^2
+        iv0 = findlast(x -> x<0.0, vpa_unnorm)
+        pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
+        #println("after pdf ", pdf[:,1,end,ir])
+        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+        # of the modified pdf are 1, 0 and 1/2 respectively
+        # will need vpa / vthe = wpa + upar/vthe
+        @. vpa.scratch2 = vpa.grid
+        # first need to calculate int dwpa pdf = zeroth_moment
+        zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
+        # calculate int dwpa wpa^2 * pdf = wpa2_moment
+        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
+        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * pdf = vpa2_moment
+        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
+        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
+        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
+        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        
+        if pdf_adjustment_option == "absvpa"
+            # calculate int dwpa |vpa| * pdf = absvpa_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
+            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            # calculate the 'B' normalisation constant
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+            # calculate the 'A' normalisation constant
+            normalisation_constant_A = (1 + normalisation_constant_B * 
+                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+            # calculate the 'C' normalisation constant
+            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                + vpa.scratch2^2 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4"
+            # calculate int dwpa vpa^4 * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                + vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "vpa4_gaussian"
+            afac = 0.1
+            bfac = 0.2
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
+            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+            @. vpa.scratch3 *= vpa.grid
+            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+            normalisation_constant_A = (1 + normalisation_constant_B 
+                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+            #                           (vpa2_wpa2_moment
+            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+            #normalisation_constant_A = (1 - normalisation_constant_B 
+            #                                * vpa2_moment) / zeroth_moment
+            #normalisation_constant_C = 0.0
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+        elseif pdf_adjustment_option == "no1st_vpa2"
+            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
+            # TMP FOR TESTING -- MAB
+            #@. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
+        else
+            println("pdf_adjustment_option not recognised")
+            stop()
+        end
+
+        # smooth the pdf at the boundary
+        #for ivpa ∈ ivpa_min+1:vpa.n-1
+        #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
+        #end
+
+        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @. vpa.scratch3 *= vpa.grid
+        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+        # @loop_vpa ivpa begin
+        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
+        # end
+        # println(io_pdf_stages,"")
+    end
+
+    # # initialise the electron pdf for positive vpa to be the mirror reflection about vpa = 0
+    # ivpa_zero = (vpa.n+1)÷2
+    # ivpa_max = vpa.n
+    # @. pdf[ivpa_zero+1:end,:,1,:] = pdf[ivpa_zero-1:-1:1,:,1,:]
+    # # calculate the zeroth v-space moment of the normalised electron pdf:
+    # # if unity to within specified tolerance, then the boundary condition is satisfied;
+    # # otherwise, modify the cutoff velocity and repeat
+    # @loop_r ir begin
+    #     unity = 2.0
+    #     while unity > 1.0
+    #         unity = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
+    #         # if unity > 1.0, then the cutoff velocity is too high so reduce it
+    #         if unity > 1.0
+    #             pdf[ivpa_max,1,1,ir] = 0.0
+    #             ivpa_max -= 1
+    #         end
+    #     end
+    #     phi[1,ir] = vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+    # end
+    # # repeat the above procedure for the boundary at z_max
+    # @. pdf[ivpa_zero-1:-1:1,:,end,:] = pdf[ivpa_zero+1:end,:,end,:]
+    # ivpa_max = 1
+    # @loop_r ir begin
+    #     unity = 2.0
+    #     while unity > 1.0
+    #         unity = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
+    #         if unity > 1.0
+    #             pdf[ivpa_max,1,end,ir] = 0.0
+    #             ivpa_max += 1
+    #         end
+    #         phi[end,ir] = vthe[end,ir]^2 * vpa.grid[ivpa_max]^2
+    #     end
+    #     #println("unity: ", unity)
+    # end
+end
+
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
     # first enforce the boundary condition at z_min.
-    # this involves forcing the pdf to be zero for electron travelling faster than the max speed
+    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
     # they could attain by accelerating in the electric field between the wall and the simulation boundary;
     # for electrons with positive velocities less than this critical value, they must have the same
     # pdf as electrons with negative velocities of the same magnitude.
@@ -1239,7 +1779,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
     # add in the contribution to the residual from the wpa advection term
     electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
-                            vpa_advect, vpa, vpa_spectral, scratch_dummy)
+                            vpa_advect, vpa, vpa_spectral, scratch_dummy)#, z)
     #dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
     #@. single_term = residual - single_term
     #max_term .= max.(max_term, abs.(single_term))
@@ -1681,7 +2221,9 @@ function check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
         end
         @loop_vperp ivperp begin
             sum_residual += sum(abs.(@view residual[iv0_start:iv0_end,ivperp,iz,ir]))
-            sum_pdf += sum(abs.(@view pdf[iv0_start:iv0_end,ivperp,iz,ir]))
+            # account for the fact that we want dg/dt << vthe/L * g, but 
+            # residual is normalized by c_ref/L_ref * g
+            sum_pdf += sum(abs.(@view pdf[iv0_start:iv0_end,ivperp,iz,ir]) * vthe[iz,ir])
         end
     end
     sum_residual, sum_pdf = MPI.Allreduce([sum_residual, sum_pdf], +, comm_world)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index acb2fb7af..f0e07d75d 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -6,6 +6,7 @@ export allocate_pdf_and_moments
 export init_pdf_and_moments!
 export enforce_boundary_conditions!
 export enforce_neutral_boundary_conditions!
+export initialize_electrons!
 
 # functional testing 
 export create_boundary_distributions
@@ -275,24 +276,125 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
+    initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
+                          vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
+                          collisions, external_source_settings,
+                          scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
+                          io_input, input_dict)
+
+    # moments.electron.dens_updated[] = false
+    # # initialise the electron density profile
+    # init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
+    # # initialise the electron parallel flow profile
+    # init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
+    #     moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
+    # # initialise the electron thermal speed profile
+    # init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
+    # # calculate the electron temperature from the thermal speed
+    # @loop_r_z ir iz begin
+    #     moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
+    # end
+    # # the electron temperature has now been updated
+    # moments.electron.temp_updated[] = true
+    # # calculate the electron parallel pressure from the density and temperature
+    # @loop_r_z ir iz begin
+    #     moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
+    # end
+    # # the electron parallel pressure now been updated
+    # moments.electron.ppar_updated[] = true
+
+    # # calculate the zed derivative of the initial electron density
+    # @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
+    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # # calculate the zed derivative of the initial electron temperature
+    # @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
+    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # # calculate the zed derivative of the initial electron thermal speed
+    # @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, 
+    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # # calculate the zed derivative of the initial electron parallel pressure
+    # @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
+    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # if composition.electron_physics == kinetic_electrons
+    #     # Initialise the array for the electron pdf
+    #     init_electron_pdf_over_density_and_boundary_phi!(
+    #         pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
+    #         moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
+    # end
+    # # calculate the electron parallel heat flux;
+    # # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
+    # calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
+    #     moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
+    #     collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+    # # calculate the zed derivative of the initial electron parallel heat flux
+    # @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
+    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
+    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+    # # calculate the electron-ion parallel friction force
+    # calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+    #     moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+    #     composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
+    
+    # # initialize the scratch arrays containing pdfs and moments for the first RK stage
+    # # the electron pdf is yet to be initialised but with the current code logic, the scratch
+    # # arrays need to exist and be otherwise initialised in order to compute the initial
+    # # electron pdf. The electron arrays will be updated as necessary by
+    # # initialize_electron_pdf!().
+    # initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
+    # # get the initial electrostatic potential and parallel electric field
+    # update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+
+    # # initialize the electron pdf that satisfies the electron kinetic equation
+    # initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
+    #                          vzeta, vr, vz, z_spectral, vpa_spectral,
+    #                          advection_structs.electron_z_advect,
+    #                          advection_structs.electron_vpa_advect, scratch_dummy,
+    #                          collisions, composition, geometry, external_source_settings,
+    #                          num_diss_params, t_input.dt, io_input, input_dict)
+
+    return nothing
+end
+
+function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
+                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
+                               collisions, external_source_settings,
+                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
+                               io_input, input_dict; restart=false)
+    
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
     init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
     # initialise the electron parallel flow profile
     init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
         moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
-    # initialise the electron thermal speed profile
-    init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
-    # calculate the electron temperature from the thermal speed
-    @loop_r_z ir iz begin
-        moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
+    # different choices for initialization of electron temperature/pressure/vth depending on whether
+    # we are restarting from a previous simulation with Boltzmann electrons or not
+    if restart
+        # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
+        # that the electron parallel temperature is constant along the field line and equal to T_e
+        moments.electron.temp .= composition.T_e
+        # the thermal speed is related to the temperature by vth_e / v_ref = sqrt((T_e/T_ref) / (m_e/m_ref))
+        moments.electron.vth .= sqrt(composition.T_e / composition.me_over_mi)
+        # ppar = 0.5 * n * T, so we can calculate the parallel pressure from the density and T_e
+        moments.electron.ppar .= 0.5 * moments.electron.dens * composition.T_e
+    else
+        # initialise the electron thermal speed profile
+        init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
+        # calculate the electron temperature from the thermal speed
+        @loop_r_z ir iz begin
+            moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
+        end
+        # calculate the electron parallel pressure from the density and temperature
+        @loop_r_z ir iz begin
+            moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
+        end
     end
     # the electron temperature has now been updated
     moments.electron.temp_updated[] = true
-    # calculate the electron parallel pressure from the density and temperature
-    @loop_r_z ir iz begin
-        moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
-    end
     # the electron parallel pressure now been updated
     moments.electron.ppar_updated[] = true
 
@@ -318,7 +420,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
     end
-    # calculate the electron parallel heat flux;
+    # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
     calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
         moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
@@ -337,9 +439,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     # arrays need to exist and be otherwise initialised in order to compute the initial
     # electron pdf. The electron arrays will be updated as necessary by
     # initialize_electron_pdf!().
-    initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
-    # get the initial electrostatic potential and parallel electric field
-    update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+    if restart == false
+        initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
+        # get the initial electrostatic potential and parallel electric field
+        update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+    end
 
     # initialize the electron pdf that satisfies the electron kinetic equation
     initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
@@ -347,7 +451,8 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
-                             num_diss_params, t_input.dt, io_input, input_dict)
+                             num_diss_params, t_input.dt, io_input, input_dict,
+                             restart_from_boltzmann=restart)
 
     return nothing
 end
@@ -474,14 +579,15 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                   vz, z_spectral, vpa_spectral, z_advect, vpa_advect,
                                   scratch_dummy, collisions, composition, geometry,
                                   external_source_settings, num_diss_params, dt, io_input,
-                                  input_dict)
+                                  input_dict; restart_from_boltzmann=false)
+
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
         begin_serial_region()
         restart_filename = get_default_restart_filename(io_input, "initial_electron";
                                                         error_if_no_file_found=false)
-        if restart_filename === nothing
+        if restart_filename === nothing || restart_from_boltzmann
             # No file to restart from
             previous_runs_info = nothing
             code_time = 0.0
@@ -1226,53 +1332,78 @@ NB: as the electron pdf is obtained via a time-independent equation,
 this 'initital' value for the electron will just be the first guess in an iterative solution
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
-        vpa, vperp, vpa_spectral, me_over_mi)
+        vpa, vperp, vpa_spectral, me_over_mi, restart_from_boltzmann=false)
 
     if z.bc == "wall"
-        @loop_r ir begin
-            # Initialise an unshifted Maxwellian as a first step
-            @loop_z iz begin
-                vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
+        if restart_from_boltzmann
+            # initialize a Maxwellian of the form g = (n_0/n_e) / (sqrt(pi)*vth_e) * exp(-v^2/vth_e^2)
+            # and solve for the values of n_0 satisfies int dvpa g  = 1, given the boundary values of phi
+            # that were obtained from a simulation with Boltzmann electrons.
+            # note that the associated upar_e may not be consistent as a consequence.
+            @loop_r ir begin
+                # first initialize a Maxwellian for the normalized pdf
+                @loop_z iz begin
+                    @loop_vperp ivperp begin
+                        @. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2)
+                    end
+                end
+                # obtain the cutoff velocities at z_min and z_max
+                vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
+                # next apply a density scale factor to ensure that the density at the boundaries is consistent with
+                # a cutoff Maxwellian with the appropriate cutoff velocities
+                density_scale_factor_zmin = 2 / (1 + erf(vpa_crit_zmin))
+                density_scale_factor_zmax = 2 / (1 + erf(vpa_crit_zmax))
                 @loop_vperp ivperp begin
-                    @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
+                    @. pdf[:,ivperp,1,ir] *= density_scale_factor_zmin
+                    @. pdf[:,ivperp,end,ir] *= density_scale_factor_zmax
                 end
             end
-            # Apply the sheath boundary condition to get cut-off boundary distribution
-            # functions and boundary values of phi
-            enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vpa, vpa_spectral, me_over_mi)
-            # get critical velocities beyond which electrons are lost to the wall
-            #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
-            #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
-            # Blend boundary distribution function into bulk of domain to avoid
-            # discontinuities (as much as possible)
-            blend_fac = 100
-            if z.nrank > 1
-                error("Distributed MPI not supported in this init yet")
-            end
-            @loop_z_vperp iz ivperp begin
-                #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
-                #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
-                #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
-                blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
-                blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
-                @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
-                                         blend_fac_lower * pdf[:,ivperp,1,ir] +
-                                         blend_fac_upper * pdf[:,ivperp,end,ir]
-                #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
-                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
-                #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *
-                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
-                #                          tanh(-sharp_fac*(vpa.grid-vpa_crit_zmax)))) #/
-                                         #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
-                                         #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
-                                         #exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
-                                         #exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
-
-                # ensure that the normalised electron pdf integrates to unity
-                norm_factor = integrate_over_vspace(pdf[:,ivperp,iz,ir], vpa.wgts)
-                @. pdf[:,ivperp,iz,ir] /= norm_factor
-                #println("TMP FOR TESTING -- init electron pdf")
-                #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
+        else
+            @loop_r ir begin
+                # Initialise an unshifted Maxwellian as a first step
+                @loop_z iz begin
+                    vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
+                    @loop_vperp ivperp begin
+                        @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
+                    end
+                end
+                # Apply the sheath boundary condition to get cut-off boundary distribution
+                # functions and boundary values of phi
+                enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vpa, vpa_spectral, me_over_mi)
+                # get critical velocities beyond which electrons are lost to the wall
+                #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
+                #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
+                # Blend boundary distribution function into bulk of domain to avoid
+                # discontinuities (as much as possible)
+                blend_fac = 100
+                if z.nrank > 1
+                    error("Distributed MPI not supported in this init yet")
+                end
+                @loop_z_vperp iz ivperp begin
+                    #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
+                    #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
+                    #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
+                    blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
+                    blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
+                    @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
+                                            blend_fac_lower * pdf[:,ivperp,1,ir] +
+                                            blend_fac_upper * pdf[:,ivperp,end,ir]
+                    #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
+                    #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
+                    #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *
+                    #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
+                    #                          tanh(-sharp_fac*(vpa.grid-vpa_crit_zmax)))) #/
+                                            #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
+                                            #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
+                                            #exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
+                                            #exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
+
+                    # ensure that the normalised electron pdf integrates to unity
+                    norm_factor = integrate_over_vspace(pdf[:,ivperp,iz,ir], vpa.wgts)
+                    @. pdf[:,ivperp,iz,ir] /= norm_factor
+                    #println("TMP FOR TESTING -- init electron pdf")
+                    #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
+                end
             end
         end
     else
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 18c75b7ed..c137695a7 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -463,7 +463,7 @@ function load_ion_moments_data(fid; printout=false, extended_moments = false)
         println("done.")
     end
     if extended_moments
-        density, parallel_flow, parallel_pressure, perpendicular_pressure, parallel_heat_flux, thermal_speed, entropy_production
+        return density, parallel_flow, parallel_pressure, perpendicular_pressure, parallel_heat_flux, thermal_speed, entropy_production
     else
         return density, parallel_flow, parallel_pressure, parallel_heat_flux, thermal_speed
     end
@@ -540,7 +540,7 @@ Reload pdf and moments from an existing output file.
 """
 function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_prefix_iblock,
                                  time_index, composition, geometry, r, z, vpa, vperp,
-                                 vzeta, vr, vz)
+                                 vzeta, vr, vz, initialize_electrons_from_boltzmann=false)
     code_time = 0.0
     previous_runs_info = nothing
     begin_serial_region()
@@ -718,7 +718,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                 restart_electron_evolve_ppar = true, true, true
             electron_evolve_density, electron_evolve_upar, electron_evolve_ppar =
                 true, true, true
-            if pdf.electron !== nothing
+            if (pdf.electron !== nothing) && (initialize_electrons_from_boltzmann == false)
                 pdf.electron.norm .=
                     reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
                                         r_range, z_range, vperp_range, vpa_range,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 248b0038f..cdaa6a2e9 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -84,7 +84,8 @@ using .debugging
 using .external_sources
 using .input_structs
 using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!,
-                           enforce_boundary_conditions!, initialize_scratch_arrays!
+                           enforce_boundary_conditions!, initialize_scratch_arrays!,
+                           initialize_electrons!
 using .load_data: reload_evolving_fields!
 using .looping
 using .moment_constraints: hard_force_moment_constraints!
@@ -294,12 +295,14 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         backup_prefix_iblock = get_prefix_iblock_and_move_existing_file(restart_filename,
                                                                         io_input.output_dir)
 
+        # TMP FOR TESTING -- MAB
+        initialize_electrons_from_boltzmann = true                                                                
         # Reload pdf and moments from an existing output file
         code_time, previous_runs_info, restart_time_index =
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
-                                    vz)
+                                    vz, initialize_electrons_from_boltzmann)
 
         # Re-initialize the source amplitude here instead of loading it from the restart
         # file so that we can change the settings between restarts.
@@ -309,6 +312,17 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         # Copy the reloaded values into the `scratch` struct
         initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
 
+        if initialize_electrons_from_boltzmann
+            # If we are initializing kinetic electrons using info from a simulation
+            # where electrons have a Boltzmann distribution, there is missing information
+            # that still needs to be specified for the electrons
+            initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
+                                 vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
+                                 collisions, external_source_settings,
+                                 scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
+                                 io_input, input_dict, restart=true)
+        end
+
         _block_synchronize()
     end
 

From 6860604ca88683270da4ac5835772f07e85f575d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 26 Feb 2024 20:05:51 +0000
Subject: [PATCH 138/428] Print a message when restarting kinetic electrons

---
 moment_kinetics/src/initial_conditions.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 4944d124e..d8aeec226 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -491,6 +491,9 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
             code_time = 0.0
             restart_time_index = -1
         else
+            if global_rank[] == 0
+                println("Restarting electrons from $restart_filename")
+            end
             # Previously-created electron distribution function exists, so use it as
             # the initial guess.
             backup_prefix_iblock =

From 46523b3d0811f86eb1f1ba898ef17a0cc6f48fd9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 28 Feb 2024 13:50:11 +0000
Subject: [PATCH 139/428] Latest input files for kinetic electron
 initialisation

---
 .../wall+sheath-bc_kinetic.toml               |  6 +-
 ...wall+sheath-bc_kinetic_krook_loworder.toml | 89 +++++++++++++++++++
 .../wall+sheath-bc_kinetic_loworder.toml      | 88 ++++++++++++++++++
 3 files changed, 181 insertions(+), 2 deletions(-)
 create mode 100644 examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
 create mode 100644 examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index bdbdf5ae6..cd65ce083 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -81,6 +81,8 @@ ascii_output = true
 
 [numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
-moment_dissipation_coefficient = 1.0
+#moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
-vpa_dissipation_coefficient = 0.2
+#vpa_dissipation_coefficient = 0.2
+#vpa_dissipation_coefficient = 2.0
+vpa_dissipation_coefficient = 20.0
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
new file mode 100644
index 000000000..f0ae0e9c2
--- /dev/null
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
@@ -0,0 +1,89 @@
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+electron_physics = "kinetic_electrons"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+krook_collisions_option = "reference_parameters"
+ionization_frequency = 2.0
+electron_ionization_frequency = 2.0
+ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+#z_nelement = 32
+z_nelement = 64
+#z_nelement = 128
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_spacing_option = "sqrt"
+vpa_ngrid = 5
+#vpa_nelement = 40
+vpa_nelement = 80
+vpa_L = 12.0 #8.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+#vpa_discretization = "gausslegendre_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+#moment_dissipation_coefficient = 0.0001
+#moment_dissipation_coefficient = 1.0
+#vpa_dissipation_coefficient = 0.002
+#vpa_dissipation_coefficient = 0.2
+#vpa_dissipation_coefficient = 2.0
+vpa_dissipation_coefficient = 20.0
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
new file mode 100644
index 000000000..1cad81bb1
--- /dev/null
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -0,0 +1,88 @@
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+electron_physics = "kinetic_electrons"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+#electron_ionization_frequency = 2.0
+#ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+#z_nelement = 32
+z_nelement = 64
+#z_nelement = 128
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_spacing_option = "sqrt"
+vpa_ngrid = 5
+#vpa_nelement = 40
+vpa_nelement = 80
+vpa_L = 12.0 #8.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+#vpa_discretization = "gausslegendre_pseudospectral"
+vz_ngrid = 17
+vz_nelement = 10
+vz_L = 8.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+#moment_dissipation_coefficient = 0.0001
+#moment_dissipation_coefficient = 1.0
+#vpa_dissipation_coefficient = 0.002
+#vpa_dissipation_coefficient = 0.2
+#vpa_dissipation_coefficient = 2.0
+vpa_dissipation_coefficient = 20.0

From a466f1d94ff87f3c8c816026e18aa99dfde7a55f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 09:02:47 +0000
Subject: [PATCH 140/428] Remove begin_r_z_region() from
 calculate_electron_qpar_from_pdf!()

...in case calculate_electron_qpar_from_pdf!() is called from inside an
`@serial_region`.
---
 moment_kinetics/src/electron_fluid_equations.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 64d999c9a..2c0190aee 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -316,7 +316,6 @@ calculate the parallel component of the electron heat flux,
 defined as qpar = 2 * ppar * vth * int dwpa (pdf * wpa^3)
 """
 function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
-    begin_r_z_region()
     # specialise to 1D for now
     ivperp = 1
     @loop_r_z ir iz begin

From 06b356d1b96826154fd6157308e7c32830d29380 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 09:04:21 +0000
Subject: [PATCH 141/428] Remove try-catch from
 update_electron_pdf_with_time_advance!()

The try-catch was added to allow the function to return some arrays as
results in case an error happened. Now that output is written to a
binary file, there is no need to do this.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 703bcf7b6..a67674662 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -245,7 +245,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
-    try
     while !electron_pdf_converged && (iteration <= max_electron_pdf_iterations)
         #dt_energy = dt_electron * 10.0
 
@@ -427,9 +426,6 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         end
         iteration += 1
     end
-    catch e
-        println("Error: $e")
-    end
     # Update the 'scratch' arrays with the final result
     begin_r_z_vperp_vpa_region()
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin

From 7840853c1bdeabe6b3f6a40b687843da357b2dee Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 11:44:58 +0000
Subject: [PATCH 142/428] Write out most recent version of constraints
 coefficients as diagnostic

It can be useful to know how much of a change the 'moment constraint
corrections' are having to apply. To help diagnose this, store the
values of the coefficients used to apply the constraints (for each
species of ion, neutral, electron, at each spatial point) into arrays in
the `moments` structs, and write out the most recent values at each
output.
---
 .../src/electron_kinetic_equation.jl          |  39 +++--
 moment_kinetics/src/file_io.jl                | 161 +++++++++++++++++-
 moment_kinetics/src/initial_conditions.jl     |   7 +-
 moment_kinetics/src/moment_constraints.jl     |  20 ++-
 .../src/moment_kinetics_structs.jl            |  27 +++
 moment_kinetics/src/time_advance.jl           |  40 +++--
 moment_kinetics/src/velocity_moments.jl       |  33 +++-
 7 files changed, 288 insertions(+), 39 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index a67674662..672fbabed 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -325,19 +325,23 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # enforce the boundary condition(s) on the electron pdf
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
                                                     vperp_spectral, vpa_spectral,
-                                                    vpa_advect,
+                                                    vpa_advect, moments,
                                                     num_diss_params.vpa_dissipation_coefficient > 0.0,
                                                     composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
 
         begin_r_z_region()
+        A = moments.electron.constraints_A_coefficient
+        B = moments.electron.constraints_B_coefficient
+        C = moments.electron.constraints_C_coefficient
         @loop_r_z ir iz begin
             if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
                 continue
             end
-            #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
-            @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
+            (A[iz,ir], B[iz,ir], C[iz,ir]) =
+                #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
+                @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
         end
         #println("B pdf 1 ", pdf[:,1,1,1])
         #println("B pdf end ", pdf[:,1,end,1])
@@ -492,7 +496,8 @@ end
 
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
-                                                     vpa_adv, vpa_diffusion, me_over_mi)
+                                                     vpa_adv, moments, vpa_diffusion,
+                                                     me_over_mi)
     # Enforce velocity-space boundary conditions
     if vpa.n > 1
         begin_r_z_vperp_region()
@@ -652,6 +657,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
         vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
         
+        normalisation_constant_A = 0.0
+        normalisation_constant_B = 0.0
+        normalisation_constant_C = 0.0
         if pdf_adjustment_option == "absvpa"
             # calculate int dwpa |vpa| * pdf = absvpa_moment
             @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
@@ -731,15 +739,19 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
             @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+            normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
                                        (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
+            normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
+            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
         else
             println("pdf_adjustment_option not recognised")
             stop()
         end
 
+        moments.electron.constraints_A_coefficient[1,ir] = normalisation_constant_A
+        moments.electron.constraints_B_coefficient[1,ir] = normalisation_constant_B
+        moments.electron.constraints_C_coefficient[1,ir] = normalisation_constant_C
+
         # smooth the pdf at the boundary
         #for ivpa ∈ 2:ivpa_max-1
         #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
@@ -893,6 +905,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
         vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
         
+        normalisation_constant_A = 0.0
+        normalisation_constant_B = 0.0
+        normalisation_constant_C = 0.0
         if pdf_adjustment_option == "absvpa"
             # calculate int dwpa |vpa| * pdf = absvpa_moment
             @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
@@ -972,15 +987,19 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
             @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
                                 + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
         elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+            normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
                                        (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
+            normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
+            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
         else
             println("pdf_adjustment_option not recognised")
             stop()
         end
 
+        moments.electron.constraints_A_coefficient[end,ir] = normalisation_constant_A
+        moments.electron.constraints_B_coefficient[end,ir] = normalisation_constant_B
+        moments.electron.constraints_C_coefficient[end,ir] = normalisation_constant_C
+
         # smooth the pdf at the boundary
         #for ivpa ∈ ivpa_min+1:vpa.n-1
         #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index e932f8a37..c8d947447 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -55,7 +55,8 @@ moments & fields only
 """
 struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
                        Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
-                       Texti5, Textn1, Textn2, Textn3, Textn4, Textn5}
+                       Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Tconstri, Tconstrn,
+                       Tconstre}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the time variable
@@ -114,6 +115,17 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     external_source_neutral_pressure_amplitude::Textn4
     external_source_neutral_controller_integral::Textn5
 
+    # handles for constraint coefficients
+    ion_constraints_A_coefficient::Tconstri
+    ion_constraints_B_coefficient::Tconstri
+    ion_constraints_C_coefficient::Tconstri
+    neutral_constraints_A_coefficient::Tconstrn
+    neutral_constraints_B_coefficient::Tconstrn
+    neutral_constraints_C_coefficient::Tconstrn
+    electron_constraints_A_coefficient::Tconstre
+    electron_constraints_B_coefficient::Tconstre
+    electron_constraints_C_coefficient::Tconstre
+
     # cumulative wall clock time taken by the run
     time_for_run
 
@@ -146,7 +158,7 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom}
+struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the pseudotime variable
@@ -163,6 +175,10 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom}
     electron_parallel_heat_flux::Tmom
     # handle for the electron thermal speed variable
     electron_thermal_speed::Tmom
+    # handles for constraint coefficients
+    electron_constraints_A_coefficient::Tconstr
+    electron_constraints_B_coefficient::Tconstr
+    electron_constraints_C_coefficient::Tconstr
 
     # Use parallel I/O?
     parallel_io::Bool
@@ -346,7 +362,11 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_parallel_flow"),
                                         getvar("electron_parallel_pressure"),
                                         getvar("electron_parallel_heat_flux"),
-                                        getvar("electron_thermal_speed"), parallel_io)
+                                        getvar("electron_thermal_speed"),
+                                        getvar("electron_constraints_A_coefficient"),
+                                        getvar("electron_constraints_B_coefficient"),
+                                        getvar("electron_constraints_C_coefficient"),
+                                        parallel_io)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -774,13 +794,16 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
         external_source_amplitude, external_source_density_amplitude,
         external_source_momentum_amplitude, external_source_pressure_amplitude,
-        external_source_controller_integral, io_chodura_lower, io_chodura_upper =
+        external_source_controller_integral, io_chodura_lower, io_chodura_upper,
+        ion_constraints_A_coefficient, ion_constraints_B_coefficient,
+        ion_constraints_C_coefficient =
             define_dynamic_ion_moment_variables!(fid, n_ion_species, r, z, parallel_io,
                                                  external_source_settings, evolve_density,
                                                  evolve_upar, evolve_ppar)
 
         io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-        io_electron_vth =
+        io_electron_vth, electron_constraints_A_coefficient,
+        electron_constraints_B_coefficient, electron_constraints_C_coefficient =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -791,7 +814,8 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         external_source_neutral_density_amplitude,
         external_source_neutral_momentum_amplitude,
         external_source_neutral_pressure_amplitude,
-        external_source_neutral_controller_integral =
+        external_source_neutral_controller_integral, neutral_constraints_A_coefficient,
+        neutral_constraints_B_coefficient, neutral_constraints_C_coefficient =
             define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r, z,
                                                      parallel_io,
                                                      external_source_settings,
@@ -818,6 +842,15 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                external_source_neutral_momentum_amplitude,
                                external_source_neutral_pressure_amplitude,
                                external_source_neutral_controller_integral,
+                               ion_constraints_A_coefficient,
+                               ion_constraints_B_coefficient,
+                               ion_constraints_C_coefficient,
+                               neutral_constraints_A_coefficient,
+                               neutral_constraints_B_coefficient,
+                               neutral_constraints_C_coefficient,
+                               electron_constraints_A_coefficient,
+                               electron_constraints_B_coefficient,
+                               electron_constraints_C_coefficient,
                                io_time_for_run, parallel_io)
     end
 
@@ -986,10 +1019,34 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
         io_chodura_upper = nothing
     end
 
+    if evolve_density || evolve_upar || evolve_ppar
+        ion_constraints_A_coefficient =
+            create_dynamic_variable!(dynamic, "ion_constraints_A_coefficient", mk_float, z, r;
+                                   n_ion_species=n_ion_species,
+                                   parallel_io=parallel_io,
+                                   description="'A' coefficient enforcing density constraint for ions")
+        ion_constraints_B_coefficient =
+            create_dynamic_variable!(dynamic, "ion_constraints_B_coefficient", mk_float, z, r;
+                                   n_ion_species=n_ion_species,
+                                   parallel_io=parallel_io,
+                                   description="'B' coefficient enforcing flow constraint for ions")
+        ion_constraints_C_coefficient =
+            create_dynamic_variable!(dynamic, "ion_constraints_C_coefficient", mk_float, z, r;
+                                   n_ion_species=n_ion_species,
+                                   parallel_io=parallel_io,
+                                   description="'C' coefficient enforcing pressure constraint for ions")
+    else
+           ion_constraints_A_coefficient = nothing
+           ion_constraints_B_coefficient = nothing
+           ion_constraints_C_coefficient = nothing
+    end
+
     return io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
            external_source_amplitude, external_source_density_amplitude,
            external_source_momentum_amplitude, external_source_pressure_amplitude,
-           external_source_controller_integral, io_chodura_lower, io_chodura_upper
+           external_source_controller_integral, io_chodura_lower, io_chodura_upper,
+           ion_constraints_A_coefficient, ion_constraints_B_coefficient,
+           ion_constraints_C_coefficient
 end
 
 """
@@ -1030,8 +1087,22 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                       description="electron species thermal speed",
                                       units="c_ref")
 
+    electron_constraints_A_coefficient =
+        create_dynamic_variable!(dynamic, "electron_constraints_A_coefficient", mk_float, z, r;
+                               parallel_io=parallel_io,
+                               description="'A' coefficient enforcing density constraint for electrons")
+    electron_constraints_B_coefficient =
+        create_dynamic_variable!(dynamic, "electron_constraints_B_coefficient", mk_float, z, r;
+                               parallel_io=parallel_io,
+                               description="'B' coefficient enforcing flow constraint for electrons")
+    electron_constraints_C_coefficient =
+        create_dynamic_variable!(dynamic, "electron_constraints_C_coefficient", mk_float, z, r;
+                               parallel_io=parallel_io,
+                               description="'C' coefficient enforcing pressure constraint for electrons")
+
     return io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-           io_electron_vth
+           io_electron_vth, electron_constraints_A_coefficient, electron_constraints_B_coefficient,
+           electron_constraints_C_coefficient
 end
 
 """
@@ -1132,12 +1203,35 @@ function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coo
         external_source_neutral_controller_integral = nothing
     end
 
+    if evolve_density || evolve_upar || evolve_ppar
+        neutral_constraints_A_coefficient =
+            create_dynamic_variable!(dynamic, "neutral_constraints_A_coefficient", mk_float, z, r;
+                                   n_neutral_species=n_neutral_species,
+                                   parallel_io=parallel_io,
+                                   description="'A' coefficient enforcing density constraint for neutrals")
+        neutral_constraints_B_coefficient =
+            create_dynamic_variable!(dynamic, "neutral_constraints_B_coefficient", mk_float, z, r;
+                                   n_neutral_species=n_neutral_species,
+                                   parallel_io=parallel_io,
+                                   description="'B' coefficient enforcing flow constraint for neutrals")
+        neutral_constraints_C_coefficient =
+            create_dynamic_variable!(dynamic, "neutral_constraints_C_coefficient", mk_float, z, r;
+                                   n_neutral_species=n_neutral_species,
+                                   parallel_io=parallel_io,
+                                   description="'C' coefficient enforcing pressure constraint for neutrals")
+    else
+           neutral_constraints_A_coefficient = nothing
+           neutral_constraints_B_coefficient = nothing
+           neutral_constraints_C_coefficient = nothing
+    end
+
     return io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
            io_thermal_speed_neutral, external_source_neutral_amplitude,
            external_source_neutral_density_amplitude,
            external_source_neutral_momentum_amplitude,
            external_source_neutral_pressure_amplitude,
-           external_source_neutral_controller_integral
+           external_source_neutral_controller_integral, neutral_constraints_A_coefficient,
+           neutral_constraints_B_coefficient, neutral_constraints_C_coefficient
 end
 
 """
@@ -1317,6 +1411,15 @@ function reopen_moments_io(file_info)
                                getvar("external_source_neutral_momentum_amplitude"),
                                getvar("external_source_neutral_pressure_amplitude"),
                                getvar("external_source_neutral_controller_integral"),
+                               getvar("ion_constraints_A_coefficient"),
+                               getvar("ion_constraints_B_coefficient"),
+                               getvar("ion_constraints_C_coefficient"),
+                               getvar("neutral_constraints_A_coefficient"),
+                               getvar("neutral_constraints_B_coefficient"),
+                               getvar("neutral_constraints_C_coefficient"),
+                               getvar("electron_constraints_A_coefficient"),
+                               getvar("electron_constraints_B_coefficient"),
+                               getvar("electron_constraints_C_coefficient"),
                                getvar("time_for_run"), parallel_io)
     end
 
@@ -1421,6 +1524,15 @@ function reopen_dfns_io(file_info)
                                      getvar("external_source_neutral_momentum_amplitude"),
                                      getvar("external_source_neutral_pressure_amplitude"),
                                      getvar("external_source_neutral_controller_integral"),
+                                     getvar("ion_constraints_A_coefficient"),
+                                     getvar("ion_constraints_B_coefficient"),
+                                     getvar("ion_constraints_C_coefficient"),
+                                     getvar("neutral_constraints_A_coefficient"),
+                                     getvar("neutral_constraints_B_coefficient"),
+                                     getvar("neutral_constraints_C_coefficient"),
+                                     getvar("electron_constraints_A_coefficient"),
+                                     getvar("electron_constraints_B_coefficient"),
+                                     getvar("electron_constraints_C_coefficient"),
                                      getvar("time_for_run"),
                                      parallel_io)
 
@@ -1585,6 +1697,17 @@ function write_ion_moments_data_to_binary(moments, n_ion_species,
                                       t_idx, parallel_io, z, r)
             end
         end
+        if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+            append_to_dynamic_var(io_moments.ion_constraints_A_coefficient,
+                                  moments.ion.constraints_A_coefficient, t_idx,
+                                  parallel_io, z, r, n_ion_species)
+            append_to_dynamic_var(io_moments.ion_constraints_B_coefficient,
+                                  moments.ion.constraints_B_coefficient, t_idx,
+                                  parallel_io, z, r, n_ion_species)
+            append_to_dynamic_var(io_moments.ion_constraints_C_coefficient,
+                                  moments.ion.constraints_C_coefficient, t_idx,
+                                  parallel_io, z, r, n_ion_species)
+        end
     end
 
     return nothing
@@ -1613,6 +1736,15 @@ function write_electron_moments_data_to_binary(moments,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
                               t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_constraints_A_coefficient,
+                              moments.electron.constraints_A_coefficient, t_idx,
+                              parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_constraints_B_coefficient,
+                              moments.electron.constraints_B_coefficient, t_idx,
+                              parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_constraints_C_coefficient,
+                              moments.electron.constraints_C_coefficient, t_idx,
+                              parallel_io, z, r)
     end
 
     return nothing
@@ -1676,6 +1808,17 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
                                       t_idx, parallel_io, z, r)
             end
         end
+        if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+            append_to_dynamic_var(io_moments.neutral_constraints_A_coefficient,
+                                  moments.neutral.constraints_A_coefficient, t_idx,
+                                  parallel_io, z, r, n_neutral_species)
+            append_to_dynamic_var(io_moments.neutral_constraints_B_coefficient,
+                                  moments.neutral.constraints_B_coefficient, t_idx,
+                                  parallel_io, z, r, n_neutral_species)
+            append_to_dynamic_var(io_moments.neutral_constraints_C_coefficient,
+                                  moments.neutral.constraints_C_coefficient, t_idx,
+                                  parallel_io, z, r, n_neutral_species)
+        end
     end
 
     return nothing
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index d8aeec226..a34dc9235 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -319,7 +319,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
         init_electron_pdf_over_density_and_boundary_phi!(
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
-            [(speed=speed,)], num_diss_params,
+            [(speed=speed,)], moments, num_diss_params,
             composition.me_over_mi)
     end
     # calculate the electron parallel heat flux;
@@ -1234,7 +1234,8 @@ NB: as the electron pdf is obtained via a time-independent equation,
 this 'initital' value for the electron will just be the first guess in an iterative solution
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
-        vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, num_diss_params, me_over_mi)
+        vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, moments, num_diss_params,
+        me_over_mi)
 
     if z.bc == "wall"
         begin_r_region()
@@ -1251,7 +1252,7 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
         # functions and boundary values of phi
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
                                                     vperp_spectral, vpa_spectral,
-                                                    vpa_advect,
+                                                    vpa_advect, moments,
                                                     num_diss_params.vpa_dissipation_coefficient > 0.0,
                                                     me_over_mi)
         begin_r_region()
diff --git a/moment_kinetics/src/moment_constraints.jl b/moment_kinetics/src/moment_constraints.jl
index eba1bc12a..ae49821dd 100644
--- a/moment_kinetics/src/moment_constraints.jl
+++ b/moment_kinetics/src/moment_constraints.jl
@@ -68,12 +68,18 @@ function hard_force_moment_constraints!(f, moments, vpa)
         B = -A*I1/I2
 
         @. f1d = A*f1d + B*vpa.grid*f1d
+
+        C = NaN
     elseif moments.evolve_density
         I0 = integrate_over_vspace(f1d, vpa.wgts)
-        @. f1d = f1d / I0
+        A = 1.0 / I0
+        @. f1d = A * f1d
+
+        B = NaN
+        C = NaN
     end
 
-    return nothing
+    return A, B, C
 end
 
 """
@@ -112,12 +118,18 @@ function hard_force_moment_constraints_neutral!(f, moments, vz)
         B = -A*I1/I2
 
         @. f1d = A*f1d + B*vz.grid*f1d
+
+        C = NaN
     elseif moments.evolve_density
         I0 = integrate_over_vspace(f1d, vz.wgts)
-        @. f1d = f1d / I0
+        A = 1.0 / I0
+        @. f1d = A * f1d
+
+        B = NaN
+        C = NaN
     end
 
-    return nothing
+    return A, B, C
 end
 
 end
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index 97354ee5e..7fe440e40 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -132,6 +132,15 @@ struct moments_ion_substruct
     external_source_pressure_amplitude::MPISharedArray{mk_float,2}
     # Integral term for the PID controller of the external source term
     external_source_controller_integral::MPISharedArray{mk_float,2}
+    # Store coefficient 'A' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_A_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
+    # Store coefficient 'B' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_B_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
+    # Store coefficient 'C' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_C_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
 end
 
 """
@@ -187,6 +196,15 @@ struct moments_electron_substruct
     dT_dz_upwind::Union{MPISharedArray{mk_float,2},Nothing}
     # this is the z-derivative of the electron thermal speed vth = sqrt(2*Tpar/m)
     dvth_dz::Union{MPISharedArray{mk_float,2},Nothing}
+    # Store coefficient 'A' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_A_coefficient::Union{MPISharedArray{mk_float,2},Nothing}
+    # Store coefficient 'B' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_B_coefficient::Union{MPISharedArray{mk_float,2},Nothing}
+    # Store coefficient 'C' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_C_coefficient::Union{MPISharedArray{mk_float,2},Nothing}
 end
 
 """
@@ -270,6 +288,15 @@ struct moments_neutral_substruct
     external_source_pressure_amplitude::MPISharedArray{mk_float,2}
     # Integral term for the PID controller of the external source term
     external_source_controller_integral::MPISharedArray{mk_float,2}
+    # Store coefficient 'A' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_A_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
+    # Store coefficient 'B' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_B_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
+    # Store coefficient 'C' from applying moment constraints so we can write it out as a
+    # diagnostic
+    constraints_C_coefficient::Union{MPISharedArray{mk_float,3},Nothing}
 end
 
 """
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index ccd0b1d31..5a65b2bee 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -435,9 +435,15 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
             composition, scratch_dummy, advance.r_diffusion, advance.vpa_diffusion)
         # Ensure normalised pdf exactly obeys integral constraints if evolving moments
         begin_s_r_z_region()
-        @loop_s_r_z is ir iz begin
-            @views hard_force_moment_constraints!(pdf.ion.norm[:,:,iz,ir,is], moments,
-                                                  vpa)
+        if moments.evolve_density && moments.enforce_conservation
+            A = moments.ion.constraints_A_coefficient
+            B = moments.ion.constraints_B_coefficient
+            C = moments.ion.constraints_C_coefficient
+            @loop_s_r_z is ir iz begin
+                (A[iz,ir,is], B[iz,ir,is], C[iz,ir,is]) =
+                    @views hard_force_moment_constraints!(pdf.ion.norm[:,:,iz,ir,is],
+                                                          moments, vpa)
+            end
         end
         # update moments in case they were affected by applying boundary conditions or
         # constraints to the pdf
@@ -458,9 +464,15 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
                 nothing, neutral_vz_advect, r, z, vzeta, vr, vz, composition, geometry,
                 scratch_dummy, advance.r_diffusion, advance.vz_diffusion)
             begin_sn_r_z_region()
-            @loop_sn_r_z isn ir iz begin
-                @views hard_force_moment_constraints_neutral!(
-                    pdf.neutral.norm[:,:,:,iz,ir,isn], moments, vz)
+            if moments.evolve_density && moments.enforce_conservation
+                A = moments.neutral.constraints_A_coefficient
+                B = moments.neutral.constraints_B_coefficient
+                C = moments.neutral.constraints_C_coefficient
+                @loop_sn_r_z isn ir iz begin
+                    (A[iz,ir,isn], B[iz,ir,isn], C[iz,ir,isn]) =
+                        @views hard_force_moment_constraints_neutral!(
+                            pdf.neutral.norm[:,:,:,iz,ir,isn], moments, vz)
+                end
             end
             update_moments_neutral!(moments, pdf.neutral.norm, vz, vr, vzeta, z, r,
                                     composition)
@@ -1346,9 +1358,13 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 
     if moments.evolve_density && moments.enforce_conservation
         begin_s_r_z_region()
+        A = moments.ion.constraints_A_coefficient
+        B = moments.ion.constraints_B_coefficient
+        C = moments.ion.constraints_C_coefficient
         @loop_s_r_z is ir iz begin
-            @views hard_force_moment_constraints!(new_scratch.pdf[:,:,iz,ir,is], moments,
-                                                  vpa)
+            (A[iz,ir,is], B[iz,ir,is], C[iz,ir,is]) =
+                @views hard_force_moment_constraints!(new_scratch.pdf[:,:,iz,ir,is],
+                                                     moments, vpa)
         end
     end
     # update remaining velocity moments that are calculable from the evolved pdf
@@ -1448,9 +1464,13 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
 
         if moments.evolve_density && moments.enforce_conservation
             begin_sn_r_z_region()
+            A = moments.neutral.constraints_A_coefficient
+            B = moments.neutral.constraints_B_coefficient
+            C = moments.neutral.constraints_C_coefficient
             @loop_sn_r_z isn ir iz begin
-                @views hard_force_moment_constraints_neutral!(
-                    new_scratch.pdf_neutral[:,:,:,iz,ir,isn], moments, vz)
+                (A[iz,ir,isn], B[iz,ir,isn], C[iz,ir,isn]) =
+                    @views hard_force_moment_constraints_neutral!(
+                        new_scratch.pdf_neutral[:,:,:,iz,ir,isn], moments, vz)
             end
         end
 
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index d223d3319..3770c92fb 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -174,6 +174,16 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
         external_source_controller_integral = allocate_shared_float(1, 1)
     end
 
+    if evolve_density || evolve_upar || evolve_ppar
+        constraints_A_coefficient = allocate_shared_float(nz, nr, n_species)
+        constraints_B_coefficient = allocate_shared_float(nz, nr, n_species)
+        constraints_C_coefficient = allocate_shared_float(nz, nr, n_species)
+    else
+        constraints_A_coefficient = nothing
+        constraints_B_coefficient = nothing
+        constraints_C_coefficient = nothing
+    end
+
     # return struct containing arrays needed to update moments
     return moments_ion_substruct(density, density_updated, parallel_flow,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,perpendicular_pressure,
@@ -183,7 +193,8 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
         dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz, entropy_production,
         external_source_amplitude, external_source_density_amplitude,
         external_source_momentum_amplitude, external_source_pressure_amplitude,
-        external_source_controller_integral)
+        external_source_controller_integral, constraints_A_coefficient,
+        constraints_B_coefficient, constraints_C_coefficient)
 end
 
 """
@@ -239,6 +250,10 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     dT_dz = allocate_shared_float(nz, nr)
     dvth_dz = allocate_shared_float(nz, nr)
     
+    constraints_A_coefficient = allocate_shared_float(nz, nr)
+    constraints_B_coefficient = allocate_shared_float(nz, nr)
+    constraints_C_coefficient = allocate_shared_float(nz, nr)
+
     # return struct containing arrays needed to update moments
     return moments_electron_substruct(density, density_updated, parallel_flow,
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,
@@ -246,7 +261,8 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
         parallel_friction_force, heat_source, v_norm_fac,
         ddens_dz, dupar_dz, dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, 
-        dT_dz, dT_dz_upwind, dvth_dz)
+        dT_dz, dT_dz_upwind, dvth_dz, constraints_A_coefficient,
+        constraints_B_coefficient, constraints_C_coefficient)
 end
 
 # neutral particles have natural mean velocities 
@@ -375,6 +391,16 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
         external_source_controller_integral = allocate_shared_float(1, 1)
     end
 
+    if evolve_density || evolve_upar || evolve_ppar
+        constraints_A_coefficient = allocate_shared_float(nz, nr, n_species)
+        constraints_B_coefficient = allocate_shared_float(nz, nr, n_species)
+        constraints_C_coefficient = allocate_shared_float(nz, nr, n_species)
+    else
+        constraints_A_coefficient = nothing
+        constraints_B_coefficient = nothing
+        constraints_C_coefficient = nothing
+    end
+
     # return struct containing arrays needed to update moments
     return moments_neutral_substruct(density, density_updated, uz, uz_updated, ur,
         ur_updated, uzeta, uzeta_updated, pz, pz_updated, pr, pr_updated, pzeta,
@@ -382,7 +408,8 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
         d2dens_dz2, duz_dz, duz_dz_upwind, d2uz_dz2, dpz_dz, dpz_dz_upwind, d2pz_dz2,
         dqz_dz, dvth_dz, external_source_amplitude, external_source_density_amplitude,
         external_source_momentum_amplitude, external_source_pressure_amplitude,
-        external_source_controller_integral)
+        external_source_controller_integral, constraints_A_coefficient,
+        constraints_B_coefficient, constraints_C_coefficient)
 end
 
 """

From 7c1d073b5b21e7fe78b3f62b024c13d840d2f841 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 13:16:17 +0000
Subject: [PATCH 143/428] Fix formatting of docstrings for plot_vs_* and
 animate_vs_* functions

---
 .../src/makie_post_processing.jl              | 110 +++++++++---------
 1 file changed, 55 insertions(+), 55 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index d36f25506..8a34fc124 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1243,7 +1243,7 @@ end
 for dim ∈ one_dimension_combinations
     function_name_str = "plot_vs_$dim"
     function_name = Symbol(function_name_str)
-    spaces = " " ^ length(function_name_str)
+    spaces = " " ^ (length(function_name_str) + 1)
     dim_str = String(dim)
     if dim == :t
         dim_grid = :( run_info.time )
@@ -1255,17 +1255,17 @@ for dim ∈ one_dimension_combinations
              export $function_name
 
              """
-             function $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, outfile=nothing, yscale=nothing,
-                      transform=identity, axis_args=Dict{Symbol,Any}(), it=nothing,
-                      $($spaces)ir=nothing, iz=nothing, ivperp=nothing, ivpa=nothing,
-                      $($spaces)ivzeta=nothing, ivr=nothing, ivz=nothing, kwargs...)
-             function $($function_name_str)(run_info, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, ax=nothing, label=nothing,
-                      $($spaces)outfile=nothing, yscale=nothing, transform=identity,
-                      $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing,
-                      $($spaces)iz=nothing, ivperp=nothing, ivpa=nothing, ivzeta=nothing,
-                      $($spaces)ivr=nothing, ivz=nothing, kwargs...)
+                 $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, outfile=nothing, yscale=nothing,
+                 transform=identity, axis_args=Dict{Symbol,Any}(), it=nothing,
+                 $($spaces)ir=nothing, iz=nothing, ivperp=nothing, ivpa=nothing,
+                 $($spaces)ivzeta=nothing, ivr=nothing, ivz=nothing, kwargs...)
+                 $($function_name_str)(run_info, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, ax=nothing, label=nothing,
+                 $($spaces)outfile=nothing, yscale=nothing, transform=identity,
+                 $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing,
+                 $($spaces)iz=nothing, ivperp=nothing, ivpa=nothing, ivzeta=nothing,
+                 $($spaces)ivr=nothing, ivz=nothing, kwargs...)
 
              Plot `var_name` from the run(s) represented by `run_info` (as returned by
              [`get_run_info`](@ref)) vs $($dim_str).
@@ -1438,7 +1438,7 @@ end
 for (dim1, dim2) ∈ two_dimension_combinations
     function_name_str = "plot_vs_$(dim2)_$(dim1)"
     function_name = Symbol(function_name_str)
-    spaces = " " ^ length(function_name_str)
+    spaces = " " ^ (length(function_name_str) + 1)
     dim1_str = String(dim1)
     dim2_str = String(dim2)
     if dim1 == :t
@@ -1453,19 +1453,19 @@ for (dim1, dim2) ∈ two_dimension_combinations
              export $function_name
 
              """
-             function $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, outfile=nothing, colorscale=identity,
-                      $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
-                      $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
-                      $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
-                      $($spaces)kwargs...)
-             function $($function_name_str)(run_info, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, ax=nothing,
-                      $($spaces)colorbar_place=nothing, title=nothing,
-                      $($spaces)outfile=nothing, colorscale=identity, transform=identity,
-                      $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing,
-                      $($spaces)iz=nothing, ivperp=nothing, ivpa=nothing, ivzeta=nothing,
-                      $($spaces)ivr=nothing, ivz=nothing, kwargs...)
+                 $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, outfile=nothing, colorscale=identity,
+                 $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
+                 $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
+                 $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
+                 $($spaces)kwargs...)
+                 $($function_name_str)(run_info, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, ax=nothing,
+                 $($spaces)colorbar_place=nothing, title=nothing,
+                 $($spaces)outfile=nothing, colorscale=identity, transform=identity,
+                 $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing,
+                 $($spaces)iz=nothing, ivperp=nothing, ivpa=nothing, ivzeta=nothing,
+                 $($spaces)ivr=nothing, ivz=nothing, kwargs...)
 
              Plot `var_name` from the run(s) represented by `run_info` (as returned by
              [`get_run_info`](@ref))vs $($dim1_str) and $($dim2_str).
@@ -1637,7 +1637,7 @@ end
 for dim ∈ one_dimension_combinations_no_t
     function_name_str = "animate_vs_$dim"
     function_name = Symbol(function_name_str)
-    spaces = " " ^ length(function_name_str)
+    spaces = " " ^ (length(function_name_str) + 1)
     dim_str = String(dim)
     dim_grid = :( run_info.$dim.grid )
     idim = Symbol(:i, dim)
@@ -1645,19 +1645,19 @@ for dim ∈ one_dimension_combinations_no_t
              export $function_name
 
              """
-             function $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, outfile=nothing, yscale=nothing,
-                      $($spaces)transform=identity, ylims=nothing,
-                      $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing, iz=nothing,
-                      $($spaces)ivperp=nothing, ivpa=nothing, ivzeta=nothing, ivr=nothing,
-                      $($spaces)ivz=nothing, kwargs...)
-             function $($function_name_str)(run_info, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, frame_index=nothing, ax=nothing,
-                      $($spaces)fig=nothing, outfile=nothing, yscale=nothing,
-                      $($spaces)transform=identity, ylims=nothing,
-                      $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing, iz=nothing,
-                      $($spaces)ivperp=nothing, ivpa=nothing, ivzeta=nothing, ivr=nothing,
-                      $($spaces)ivz=nothing, kwargs...)
+                 $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, outfile=nothing, yscale=nothing,
+                 $($spaces)transform=identity, ylims=nothing,
+                 $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing, iz=nothing,
+                 $($spaces)ivperp=nothing, ivpa=nothing, ivzeta=nothing, ivr=nothing,
+                 $($spaces)ivz=nothing, kwargs...)
+                 $($function_name_str)(run_info, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, frame_index=nothing, ax=nothing,
+                 $($spaces)fig=nothing, outfile=nothing, yscale=nothing,
+                 $($spaces)transform=identity, ylims=nothing,
+                 $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing, iz=nothing,
+                 $($spaces)ivperp=nothing, ivpa=nothing, ivzeta=nothing, ivr=nothing,
+                 $($spaces)ivz=nothing, kwargs...)
 
              Animate `var_name` from the run(s) represented by `run_info` (as returned by
              [`get_run_info`](@ref))vs $($dim_str).
@@ -1870,7 +1870,7 @@ end
 for (dim1, dim2) ∈ two_dimension_combinations_no_t
     function_name_str = "animate_vs_$(dim2)_$(dim1)"
     function_name = Symbol(function_name_str)
-    spaces = " " ^ length(function_name_str)
+    spaces = " " ^ (length(function_name_str) + 1)
     dim1_str = String(dim1)
     dim2_str = String(dim2)
     dim1_grid = :( run_info.$dim1.grid )
@@ -1881,20 +1881,20 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
              export $function_name
 
              """
-             function $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, outfile=nothing, colorscale=identity,
-                      $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
-                      $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
-                      $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
-                      $($spaces)kwargs...)
-             function $($function_name_str)(run_info, var_name; is=1, data=nothing,
-                      $($spaces)input=nothing, frame_index=nothing, ax=nothing,
-                      $($spaces)fig=nothing, colorbar_place=colorbar_place,
-                      $($spaces)title=nothing, outfile=nothing, colorscale=identity,
-                      $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
-                      $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
-                      $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
-                      $($spaces)kwargs...)
+                 $($function_name_str)(run_info::Tuple, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, outfile=nothing, colorscale=identity,
+                 $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
+                 $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
+                 $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
+                 $($spaces)kwargs...)
+                 $($function_name_str)(run_info, var_name; is=1, data=nothing,
+                 $($spaces)input=nothing, frame_index=nothing, ax=nothing,
+                 $($spaces)fig=nothing, colorbar_place=colorbar_place,
+                 $($spaces)title=nothing, outfile=nothing, colorscale=identity,
+                 $($spaces)transform=identity, axis_args=Dict{Symbol,Any}(),
+                 $($spaces)it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
+                 $($spaces)ivpa=nothing, ivzeta=nothing, ivr=nothing, ivz=nothing,
+                 $($spaces)kwargs...)
 
              Animate `var_name` from the run(s) represented by `run_info` (as returned by
              [`get_run_info`](@ref))vs $($dim1_str) and $($dim2_str).

From 02b8e5e049d3a5a66335634aa725e08f71d02599 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 13:20:35 +0000
Subject: [PATCH 144/428] Allow passing a `label` kwarg to 1d animation
 functions

---
 .../src/makie_post_processing.jl               | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 8a34fc124..ca2de4122 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1654,7 +1654,7 @@ for dim ∈ one_dimension_combinations_no_t
                  $($function_name_str)(run_info, var_name; is=1, data=nothing,
                  $($spaces)input=nothing, frame_index=nothing, ax=nothing,
                  $($spaces)fig=nothing, outfile=nothing, yscale=nothing,
-                 $($spaces)transform=identity, ylims=nothing,
+                 $($spaces)transform=identity, ylims=nothing, label=nothing,
                  $($spaces)axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing, iz=nothing,
                  $($spaces)ivperp=nothing, ivpa=nothing, ivzeta=nothing, ivr=nothing,
                  $($spaces)ivz=nothing, kwargs...)
@@ -1689,6 +1689,9 @@ for dim ∈ one_dimension_combinations_no_t
              When a single `run_info` is passed, an `Axis` can be passed to `ax`. If it
              is, the plot will be added to `ax`.
 
+             When a single `run_info` is passed, `label` can be passed to set a custom
+             label for the line. By default the `run_info.run_name` is used.
+
              `outfile` is required for animations unless `ax` is passed. The animation
              will be saved to a file named `outfile`.  The suffix determines the file
              type. If both `outfile` and `ax` are passed, then the `Figure` containing
@@ -1785,10 +1788,10 @@ for dim ∈ one_dimension_combinations_no_t
              function $function_name(run_info, var_name; is=1, data=nothing,
                                      input=nothing, frame_index=nothing, ax=nothing,
                                      fig=nothing, outfile=nothing, yscale=nothing,
-                                     ylims=nothing, axis_args=Dict{Symbol,Any}(),
-                                     it=nothing, ir=nothing, iz=nothing, ivperp=nothing,
-                                     ivpa=nothing, ivzeta=nothing, ivr=nothing,
-                                     ivz=nothing, kwargs...)
+                                     ylims=nothing, label=nothing,
+                                     axis_args=Dict{Symbol,Any}(), it=nothing, ir=nothing,
+                                     iz=nothing, ivperp=nothing, ivpa=nothing,
+                                     ivzeta=nothing, ivr=nothing, ivz=nothing, kwargs...)
                  if input === nothing
                      if run_info.dfns
                          input = input_dict_dfns[var_name]
@@ -1827,13 +1830,16 @@ for dim ∈ one_dimension_combinations_no_t
                  else
                      fig = nothing
                  end
+                 if label === nothing
+                     label = run_info.run_name
+                 end
 
                  x = $dim_grid
                  if $idim !== nothing
                      x = x[$idim]
                  end
                  animate_1d(x, data; ax=ax, ylims=ylims, frame_index=ind,
-                            label=run_info.run_name, kwargs...)
+                            label=label, kwargs...)
 
                  if input.show_element_boundaries && fig !== nothing
                      element_boundary_inds =

From 11e23656d9bf2d0d660bb2c6e53168eb93208913 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 14:28:35 +0000
Subject: [PATCH 145/428] Make plots of coefficients for moment constraints

---
 .../src/makie_post_processing.jl              | 354 +++++++++++++++++-
 1 file changed, 351 insertions(+), 3 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index ca2de4122..654c0cba2 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -30,9 +30,7 @@ using moment_kinetics.analysis: analyze_fields_data, check_Chodura_condition,
                                 get_unnormalised_f_2d
 using moment_kinetics.array_allocation: allocate_float
 using moment_kinetics.coordinates: define_coordinate
-using moment_kinetics.input_structs: grid_input, advection_input,
-                                     set_defaults_and_check_top_level!,
-                                     set_defaults_and_check_section!, Dict_to_NamedTuple
+using moment_kinetics.input_structs
 using moment_kinetics.looping: all_dimensions, ion_dimensions, neutral_dimensions
 using moment_kinetics.manufactured_solns: manufactured_solutions,
                                           manufactured_electric_fields
@@ -297,6 +295,8 @@ function makie_post_process(run_dir::Union{String,Tuple},
         plot_neutral_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
     end
 
+    constraints_plots(run_info; plot_prefix=plot_prefix)
+
     if has_rdim
         # Plots for 2D instability do not make sense for 1D simulations
         instability_input = input_dict["instability2D"]
@@ -671,6 +671,22 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         animation_ext=this_input_dict["animation_ext"],
        )
 
+    set_defaults_and_check_section!(
+        this_input_dict, "constraints";
+        plot=false,
+        animate=true,
+        it0=this_input_dict["it0"],
+        ir0=this_input_dict["ir0"],
+        iz0=this_input_dict["iz0"],
+        ivperp0=this_input_dict["ivperp0"],
+        ivpa0=this_input_dict["ivpa0"],
+        ivzeta0=this_input_dict["ivzeta0"],
+        ivr0=this_input_dict["ivr0"],
+        ivz0=this_input_dict["ivz0"],
+        animation_ext=this_input_dict["animation_ext"],
+        show_element_boundaries=this_input_dict["show_element_boundaries"],
+       )
+
     set_defaults_and_check_section!(
         this_input_dict, "Chodura_condition";
         plot_vs_t=false,
@@ -4511,6 +4527,338 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
     return nothing
 end
 
+"""
+    constraints_plots(run_info; plot_prefix=plot_prefix)
+
+Plot and/or animate the coefficients used to correct the normalised distribution
+function(s) (aka shape functions) to obey the moment constraints.
+
+If there were no discretisation errors, we would have \$A=1\$, \$B=0\$, \$C=0\$. The
+plots/animations show \$(A-1)\$ so that all three coefficients can be shown nicely on the
+same axes.
+"""
+function constraints_plots(run_info; plot_prefix=plot_prefix)
+    input = Dict_to_NamedTuple(input_dict["constraints"])
+
+    if !(input.plot || input.animate)
+        return nothing
+    end
+
+    try
+        println("Making plots of moment constraints coefficients")
+
+        if !isa(run_info, Tuple)
+            run_info = (run_info,)
+        end
+
+        it0 = input.it0
+        ir0 = input.ir0
+
+        if input.plot
+            if any(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar
+                   for ri ∈ run_info)
+
+                # Ions
+                frame_index = Observable(1)
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+                for ri ∈ run_info
+                    if !(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                        continue
+                    end
+                    nspecies = ri.n_ion_species
+                    for is ∈ 1:nspecies
+                        if length(run_info) > 1
+                            prefix = ri.run_name * ", "
+                        else
+                            prefix = ""
+                        end
+                        if nspecies > 1
+                            suffix = ", species $is"
+                        else
+                            suffix = ""
+                        end
+
+                        varname = "ion_constraints_A_coefficient"
+                        label = prefix * "(A-1)" * suffix
+                        data = get_variable(ri, varname; it=it0, is=is, ir=ir0)
+                        data .-= 1.0
+                        plot_vs_z(ri, varname; label=label, data=data, ax=ax, input=input)
+
+                        varname = "ion_constraints_B_coefficient"
+                        label = prefix * "B" * suffix
+                        plot_vs_z(ri, varname; label=label, ax=ax, it=it0, is=is, ir=ir0,
+                                  input=input)
+
+                        varname = "ion_constraints_C_coefficient"
+                        label = prefix * "C" * suffix
+                        plot_vs_z(ri, varname; label=label, ax=ax, it=it0, is=is, ir=ir0,
+                                  input=input)
+                    end
+                end
+                put_legend_right(fig, ax)
+                save(plot_prefix * "ion_constraints.pdf", fig)
+            end
+
+            # Neutrals
+            if any(ri.n_neutral_species > 1
+                   && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                   for ri ∈ run_info)
+
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+                for ri ∈ run_info
+                    if !(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                        continue
+                    end
+                    nspecies = ri.n_neutral_species
+                    for is ∈ 1:nspecies
+                        if length(run_info) > 1
+                            prefix = ri.run_name * ", "
+                        else
+                            prefix = ""
+                        end
+                        if nspecies > 1
+                            suffix = ", species $is"
+                        else
+                            suffix = ""
+                        end
+
+                        varname = "neutral_constraints_A_coefficient"
+                        label = prefix * "(A-1)" * suffix
+                        data = get_variable(ri, varname; it=it0, is=is, ir=ir0)
+                        data .-= 1.0
+                        plot_vs_z(ri, varname; label=label, data=data, ax=ax, input=input)
+
+                        varname = "neutral_constraints_B_coefficient"
+                        label = prefix * "B" * suffix
+                        plot_vs_z(ri, varname; label=label, ax=ax, it=it0, is=is, ir=ir0,
+                                  input=input)
+
+                        varname = "neutral_constraints_C_coefficient"
+                        label = prefix * "C" * suffix
+                        plot_vs_z(ri, varname; label=label, ax=ax, it=it0, is=is, ir=ir0,
+                                  input=input)
+                    end
+                end
+                put_legend_right(fig, ax)
+                save(plot_prefix * "neutral_constraints.pdf", fig)
+            end
+
+            # Electrons
+            if any(ri.composition.electron_physics == kinetic_electrons for ri ∈ run_info)
+
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+                for ri ∈ run_info
+                    if length(run_info) > 1
+                        prefix = ri.run_name * ", "
+                    else
+                        prefix = ""
+                    end
+
+                    varname = "electron_constraints_A_coefficient"
+                    label = prefix * "(A-1)"
+                    data = get_variable(ri, varname; it=it0, ir=ir0)
+                    data .-= 1.0
+                    plot_vs_z(ri, varname; label=label, data=data, ax=ax, input=input)
+
+                    varname = "electron_constraints_B_coefficient"
+                    label = prefix * "B"
+                    plot_vs_z(ri, varname; label=label, ax=ax, it=it0, ir=ir0,
+                              input=input)
+
+                    varname = "electron_constraints_C_coefficient"
+                    label = prefix * "C"
+                    plot_vs_z(ri, varname; label=label, ax=ax, it=it0, ir=ir0,
+                              input=input)
+                end
+                put_legend_right(fig, ax)
+                save(plot_prefix * "electron_constraints.pdf", fig)
+            end
+        end
+
+        if input.animate
+            nt = minimum(ri.nt for ri ∈ run_info)
+
+            if any(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar
+                   for ri ∈ run_info)
+
+                # Ions
+                frame_index = Observable(1)
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+
+                # Calculate plot limits manually so we can exclude the first time point, which
+                # often has a large value for (A-1) due to the way initialisation is done,
+                # which can make the subsequent values hard to see.
+                ymin = Inf
+                ymax = -Inf
+                for ri ∈ run_info
+                    if !(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                        continue
+                    end
+                    nspecies = ri.n_ion_species
+                    for is ∈ 1:nspecies
+                        if length(run_info) > 1
+                            prefix = ri.run_name * ", "
+                        else
+                            prefix = ""
+                        end
+                        if nspecies > 1
+                            suffix = ", species $is"
+                        else
+                            suffix = ""
+                        end
+
+                        varname = "ion_constraints_A_coefficient"
+                        label = prefix * "(A-1)" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        data .-= 1.0
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, input=input)
+
+                        varname = "ion_constraints_B_coefficient"
+                        label = prefix * "B" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, is=is, ir=ir0,
+                                     input=input)
+
+                        varname = "ion_constraints_C_coefficient"
+                        label = prefix * "C" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, is=is, ir=ir0,
+                                     input=input)
+                    end
+                end
+                put_legend_right(fig, ax)
+                ylims!(ax, ymin, ymax)
+                save_animation(fig, frame_index, nt,
+                               plot_prefix * "ion_constraints." * input.animation_ext)
+            end
+
+            # Neutrals
+            if any(ri.n_neutral_species > 1
+                   && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                   for ri ∈ run_info)
+
+                frame_index = Observable(1)
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+
+                # Calculate plot limits manually so we can exclude the first time point, which
+                # often has a large value for (A-1) due to the way initialisation is done,
+                # which can make the subsequent values hard to see.
+                ymin = Inf
+                ymax = -Inf
+                for ri ∈ run_info
+                    if !(ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                        continue
+                    end
+                    nspecies = ri.n_neutral_species
+                    for is ∈ 1:nspecies
+                        if length(run_info) > 1
+                            prefix = ri.run_name * ", "
+                        else
+                            prefix = ""
+                        end
+                        if nspecies > 1
+                            suffix = ", species $is"
+                        else
+                            suffix = ""
+                        end
+
+                        varname = "neutral_constraints_A_coefficient"
+                        label = prefix * "(A-1)" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        data .-= 1.0
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, input=input)
+
+                        varname = "neutral_constraints_B_coefficient"
+                        label = prefix * "B" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, is=is, ir=ir0,
+                                     input=input)
+
+                        varname = "neutral_constraints_C_coefficient"
+                        label = prefix * "C" * suffix
+                        data = get_variable(ri, varname; is=is, ir=ir0)
+                        ymin = min(ymin, minimum(data[:,2:end]))
+                        ymax = max(ymax, maximum(data[:,2:end]))
+                        animate_vs_z(ri, varname; label=label, data=data,
+                                     frame_index=frame_index, ax=ax, is=is, ir=ir0,
+                                     input=input)
+                    end
+                end
+                put_legend_right(fig, ax)
+                ylims!(ax, ymin, ymax)
+                save_animation(fig, frame_index, nt,
+                               plot_prefix * "neutral_constraints." * input.animation_ext)
+            end
+
+            # Electrons
+            if any(ri.composition.electron_physics == kinetic_electrons for ri ∈ run_info)
+
+                frame_index = Observable(1)
+                fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
+
+                # Calculate plot limits manually so we can exclude the first time point, which
+                # often has a large value for (A-1) due to the way initialisation is done,
+                # which can make the subsequent values hard to see.
+                ymin = Inf
+                ymax = -Inf
+                for ri ∈ run_info
+                    if length(run_info) > 1
+                        prefix = ri.run_name * ", "
+                    else
+                        prefix = ""
+                    end
+
+                    varname = "electron_constraints_A_coefficient"
+                    label = prefix * "(A-1)"
+                    data = get_variable(ri, varname; ir=ir0)
+                    data .-= 1.0
+                    ymin = min(ymin, minimum(data[:,2:end]))
+                    ymax = max(ymax, maximum(data[:,2:end]))
+                    animate_vs_z(ri, varname; label=label, data=data,
+                                 frame_index=frame_index, ax=ax, input=input)
+
+                    varname = "electron_constraints_B_coefficient"
+                    label = prefix * "B"
+                    data = get_variable(ri, varname; ir=ir0)
+                    ymin = min(ymin, minimum(data[:,2:end]))
+                    ymax = max(ymax, maximum(data[:,2:end]))
+                    animate_vs_z(ri, varname; label=label, data=data,
+                                 frame_index=frame_index, ax=ax, ir=ir0, input=input)
+
+                    varname = "electron_constraints_C_coefficient"
+                    label = prefix * "C"
+                    data = get_variable(ri, varname; ir=ir0)
+                    ymin = min(ymin, minimum(data[:,2:end]))
+                    ymax = max(ymax, maximum(data[:,2:end]))
+                    animate_vs_z(ri, varname; label=label, data=data,
+                                 frame_index=frame_index, ax=ax, ir=ir0, input=input)
+                end
+                put_legend_right(fig, ax)
+                ylims!(ax, ymin, ymax)
+                save_animation(fig, frame_index, nt,
+                               plot_prefix * "electron_constraints." * input.animation_ext)
+            end
+        end
+    catch e
+        println("Error in constraints_plots(). Error was ", e)
+    end
+end
+
 """
     Chodura_condition_plots(run_info::Tuple; plot_prefix)
     Chodura_condition_plots(run_info; plot_prefix=nothing, axes=nothing)

From 771082cd58f8a6266c00e7f7a1acbcd3980bed28 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 11 Mar 2024 15:14:17 +0000
Subject: [PATCH 146/428] Remove
 enforce_boundary_condition_on_electron_pdf_noshift!() copy of bc

Rather than duplicating all the code, would be nicer to have an argument
or option flag to turn the shift by upar_e on/off. For now, just have
lines with/without shift present and comment/uncomment them.
---
 .../src/electron_kinetic_equation.jl          | 536 +-----------------
 1 file changed, 4 insertions(+), 532 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index ccf1bb31b..b2a576468 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -504,536 +504,6 @@ function speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
     return nothing
 end
 
-function enforce_boundary_condition_on_electron_pdf_noshift!(pdf, phi, vthe, upar, vpa, vpa_spectral, me_over_mi)
-    # first enforce the boundary condition at z_min.
-    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
-    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
-    # for electrons with positive velocities less than this critical value, they must have the same
-    # pdf as electrons with negative velocities of the same magnitude.
-    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
-    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
-
-    begin_r_region()
-
-    # pdf_adjustment_option determines the velocity-dependent pre-factor for the
-    # corrections to the pdf needed to ensure moment constraints are satisfied
-    #pdf_adjustment_option = "vpa4"
-    #pdf_adjustment_option = "vpa4_gaussian"
-    pdf_adjustment_option = "no1st_vpa2"
-
-    cutoff_step_width = 0.1
-
-    # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
-    #wpa_values = vpa.scratch
-    # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
-    #interpolated_pdf = vpa.scratch2
-    reversed_pdf = vpa.scratch
-
-    # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
-    #ivpa_zero = (vpa.n+1)÷2
-
-    @loop_r ir begin
-        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
-        #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
-
-        # Want to construct the w-grid corresponding to -vpa.
-        #   wpa(vpa) = (vpa - upar)/vth
-        #   ⇒ vpa = vth*wpa(vpa) + upar
-        #   wpa(-vpa) = (-vpa - upar)/vth
-        #             = (-(vth*wpa(vpa) + upar) - upar)/vth
-        #             = (-vth*wpa - 2*upar)/vth
-        #             = -wpa - 2*upar/vth
-        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        # reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
-        # #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
-        # # interpolate the pdf onto this grid
-        # #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
-        # @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-        # reverse!(reversed_pdf)
-        reversed_pdf .= pdf[:,1,1,ir]
-        # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
-        #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
-        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
-        # interpolate back onto the original wpa grid
-        #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-        # construct wpa * pdf
-        #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
-        # calculate the first moment of the normalised pdf
-        #first_vspace_moment = 0.0
-        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
-        ivpa = 0
-        ivpa_max = vpa.n + 1
-        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-        # if the first moment is positive, then the cutoff speed needs to be reduced
-        upar_over_vth = upar[1,ir] / vthe[1,ir]
-        #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
-        #println("$first_vspace_moment, u/vth=$upar_over_vth")
-        vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid
-        upar_integral = 0.0
-        #while first_vspace_moment > upar_over_vth # > 0.0
-        #    # zero out the pdf at the current cutoff velocity
-        #    pdf[ivpa_max,1,1,ir] = 0.0
-        #    # update wpa * pdf
-        #    vpa.scratch3[ivpa_max] = 0.0
-        #    # calculate the updated first moment of the normalised pdf
-        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
-        #    if first_vspace_moment > upar_over_vth #0.0
-        #        ivpa_max -= 1
-        #    end
-        #end
-        upar0 = upar[1,ir]
-        #println("before pdf left ", pdf[:,1,1,ir])
-        while upar_integral > upar0 && ivpa_max > 1
-            ivpa += 1
-            ivpa_max -= 1
-            # zero out the reversed pdf at the current cutoff velocity
-            #reversed_pdf[ivpa_max] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
-            #println("left ", ivpa, " ", upar_integral, " ", upar0)
-        end
-        integral_excess = upar_integral - upar0
-        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
-        #println("fraction_of_pdf=", fraction_of_pdf)
-
-        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-        # integral from the point at ivpa is required to make upar_integral<upar0) the
-        # cut-off velocity is half way between ivpa and ivpa+1, while when
-        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-        # ivpa is required to make upar_integral<upar0) the cut-off is half way between
-        # ivpa-1 and ivpa.
-        vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
-               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
-
-        #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
-        #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
-        #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
-        #@loop_vpa ivpa begin
-        #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
-        #end
-        reversed_pdf[ivpa_max+1:end] .= 0
-        reversed_pdf[ivpa_max] *= fraction_of_pdf
-        #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
-        #println("done first cutoff loop")
-        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
-        phi[1,ir] = me_over_mi * vmax^2
-        iv0 = findfirst(x -> x>0.0, vpa_unnorm)
-        pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
-        #println("check reversed change ", reversed_pdf[iv0:end])
-        #println("reversed_pdf ", reversed_pdf)
-        #println("after pdf left ", pdf[:,1,1,ir])
-        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-        # of the modified pdf are 1, 0 and 1/2 respectively
-        # will need vpa / vthe = wpa + upar/vthe
-        @. vpa.scratch2 = vpa.grid
-        # first need to calculate int dwpa pdf = zeroth_moment
-        zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
-        # calculate int dwpa wpa^2 * pdf = wpa2_moment
-        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
-        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * pdf = vpa2_moment
-        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
-        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
-        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
-        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        
-        if pdf_adjustment_option == "absvpa"
-            # calculate int dwpa |vpa| * pdf = absvpa_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
-            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            # calculate the 'B' normalisation constant
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-            # calculate the 'A' normalisation constant
-            normalisation_constant_A = (1 + normalisation_constant_B * 
-                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-            # calculate the 'C' normalisation constant
-            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                + vpa.scratch2^2 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4"
-            # calculate int dwpa vpa^4 * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                + vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4_gaussian"
-            afac = 0.1
-            bfac = 0.2
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
-            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #                           (vpa2_wpa2_moment
-            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-            #normalisation_constant_A = (1 - normalisation_constant_B 
-            #                                * vpa2_moment) / zeroth_moment
-            #normalisation_constant_C = 0.0
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
-            # TMP FOR TESTING -- MAB
-            #@. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
-        else
-            println("pdf_adjustment_option not recognised")
-            stop()
-        end
-
-        # smooth the pdf at the boundary
-        #for ivpa ∈ 2:ivpa_max-1
-        #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
-        #end
-    end
-
-    # next enforce the boundary condition at z_max.
-    # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
-    # they could attain by accelerating in the electric field between the wall and the simulation boundary;
-    # for electrons with negative velocities less than this critical value, they must have the same
-    # pdf as electrons with positive velocities of the same magnitude.
-    # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
-    # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
-    
-    # io_pdf_stages = open("pdf_stages.txt", "w")
-    # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-    # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid^2
-    # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-    # @loop_vpa ivpa begin
-    #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-    #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 0)
-    # end
-    # println(io_pdf_stages,"")
-
-    @loop_r ir begin
-        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-
-        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        # reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
-        # #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
-        # # interpolate the pdf onto this grid
-        # #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
-        # @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-        # # reverse!(reversed_pdf)
-        reversed_pdf .= pdf[:,1,end,ir]
-        # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
-        #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
-        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
-        # interpolate back onto the original wpa grid
-        #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
-        # end
-        # println(io_pdf_stages,"")
-
-        # construct wpa * pdf
-        #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
-        # calculate the first moment of the normalised pdf
-        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
-        ivpa = vpa.n+1
-        ivpa_min = 0
-        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-        # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
-        upar_over_vth = upar[end,ir] / vthe[end,ir]
-        #println("$first_vspace_moment, u/vth=$upar_over_vth")
-        vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid
-        upar_integral = 0.0
-        #while first_vspace_moment < upar_over_vth # < 0.0
-        #    # zero out the pdf at the current cutoff velocity
-        #    pdf[ivpa_min,1,end,ir] = 0.0
-        #    # update wpa * pdf
-        #    vpa.scratch3[ivpa_min] = 0.0
-        #    # calculate the updated first moment of the normalised pdf
-        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        #    if first_vspace_moment < upar_over_vth
-        #        ivpa_min += 1
-        #    end
-        #end
-        upar_end = upar[end,ir]
-        #println("before pdf ", pdf[:,1,end,ir])
-        while upar_integral < upar_end && ivpa > 1
-            ivpa -= 1
-            ivpa_min += 1
-            # zero out the reversed pdf at the current cutoff velocity
-            #reversed_pdf[ivpa_min] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
-            #println("right ", ivpa, " ", upar_integral, " ", upar_end)
-        end
-        integral_excess = upar_integral - upar_end
-        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
-        #println("B fraction_of_pdf=", fraction_of_pdf)
-
-        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-        # integral from the point at ivpa is required to make upar_integral>upar_end) the
-        # cut-off velocity is half way between ivpa-1 and ivpa, while when
-        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-        # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
-        # ivpa and ivpa+1.
-        vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
-               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
-
-        #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
-        #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
-        #@loop_vpa ivpa begin
-        #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
-        #end
-        reversed_pdf[1:ivpa_min-1] .= 0
-        reversed_pdf[ivpa_min] *= fraction_of_pdf
-        #println("done second cutoff loop")
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
-        # end
-        # println(io_pdf_stages,"")
-
-        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
-        phi[end,ir] = me_over_mi * vmin^2
-        iv0 = findlast(x -> x<0.0, vpa_unnorm)
-        pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
-        #println("after pdf ", pdf[:,1,end,ir])
-        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-        # of the modified pdf are 1, 0 and 1/2 respectively
-        # will need vpa / vthe = wpa + upar/vthe
-        @. vpa.scratch2 = vpa.grid
-        # first need to calculate int dwpa pdf = zeroth_moment
-        zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
-        # calculate int dwpa wpa^2 * pdf = wpa2_moment
-        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
-        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * pdf = vpa2_moment
-        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
-        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
-        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
-        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        
-        if pdf_adjustment_option == "absvpa"
-            # calculate int dwpa |vpa| * pdf = absvpa_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
-            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            # calculate the 'B' normalisation constant
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-            # calculate the 'A' normalisation constant
-            normalisation_constant_A = (1 + normalisation_constant_B * 
-                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-            # calculate the 'C' normalisation constant
-            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                + vpa.scratch2^2 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4"
-            # calculate int dwpa vpa^4 * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                + vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4_gaussian"
-            afac = 0.1
-            bfac = 0.2
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
-            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #                           (vpa2_wpa2_moment
-            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-            #normalisation_constant_A = (1 - normalisation_constant_B 
-            #                                * vpa2_moment) / zeroth_moment
-            #normalisation_constant_C = 0.0
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_B = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_B*vpa2_wpa2_moment) / wpa2_moment
-            # TMP FOR TESTING -- MAB
-            #@. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B)
-        else
-            println("pdf_adjustment_option not recognised")
-            stop()
-        end
-
-        # smooth the pdf at the boundary
-        #for ivpa ∈ ivpa_min+1:vpa.n-1
-        #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
-        #end
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
-        # end
-        # println(io_pdf_stages,"")
-    end
-
-    # # initialise the electron pdf for positive vpa to be the mirror reflection about vpa = 0
-    # ivpa_zero = (vpa.n+1)÷2
-    # ivpa_max = vpa.n
-    # @. pdf[ivpa_zero+1:end,:,1,:] = pdf[ivpa_zero-1:-1:1,:,1,:]
-    # # calculate the zeroth v-space moment of the normalised electron pdf:
-    # # if unity to within specified tolerance, then the boundary condition is satisfied;
-    # # otherwise, modify the cutoff velocity and repeat
-    # @loop_r ir begin
-    #     unity = 2.0
-    #     while unity > 1.0
-    #         unity = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
-    #         # if unity > 1.0, then the cutoff velocity is too high so reduce it
-    #         if unity > 1.0
-    #             pdf[ivpa_max,1,1,ir] = 0.0
-    #             ivpa_max -= 1
-    #         end
-    #     end
-    #     phi[1,ir] = vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
-    # end
-    # # repeat the above procedure for the boundary at z_max
-    # @. pdf[ivpa_zero-1:-1:1,:,end,:] = pdf[ivpa_zero+1:end,:,end,:]
-    # ivpa_max = 1
-    # @loop_r ir begin
-    #     unity = 2.0
-    #     while unity > 1.0
-    #         unity = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
-    #         if unity > 1.0
-    #             pdf[ivpa_max,1,end,ir] = 0.0
-    #             ivpa_max += 1
-    #         end
-    #         phi[end,ir] = vthe[end,ir]^2 * vpa.grid[ivpa_max]^2
-    #     end
-    #     #println("unity: ", unity)
-    # end
-end
-
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, moments, vpa_diffusion,
@@ -1097,8 +567,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
         # Need to reverse vpa.grid because the grid passed as the second argument of
         # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
+        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
         #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
         @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
@@ -1328,8 +799,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
         # Need to reverse vpa.grid because the grid passed as the second argument of
         # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
+        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
         #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
+        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
         # interpolate the pdf onto this grid
         #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
         @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).

From 7a29a8ccd936978bb1253a0520ed043f1e07263c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 12 Mar 2024 11:24:13 +0000
Subject: [PATCH 147/428] More example input files for electron physics

One with Boltzmann electrons to create a better initial condition, one
with kinetic electrons that can restart from the Boltzmann electrons.
---
 .../wall+sheath-bc_boltzmann_loworder.toml    | 92 +++++++++++++++++++
 .../wall+sheath-bc_kinetic_loworder.toml      |  9 +-
 2 files changed, 98 insertions(+), 3 deletions(-)
 create mode 100644 examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml

diff --git a/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
new file mode 100644
index 000000000..aced83ff7
--- /dev/null
+++ b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
@@ -0,0 +1,92 @@
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+n_ion_species = 1
+n_neutral_species = 1
+#boltzmann_electron_response = false
+#electron_physics = "kinetic_electrons"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 2.0
+electron_charge_exchange_frequency = 0.0
+nu_ei = 0.0
+ionization_frequency = 2.0
+#electron_ionization_frequency = 2.0
+#ionization_energy = 1.0
+constant_ionization_rate = false
+nstep = 1000000
+#nstep = 1
+dt = 1.0e-5
+nwrite = 10000
+nwrite_dfns = 10000
+use_semi_lagrange = false
+n_rk_stages = 4
+split_operators = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+#z_nelement = 32
+z_nelement = 64
+#z_nelement = 128
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_spacing_option = "sqrt"
+vpa_ngrid = 5
+#vpa_nelement = 40
+vpa_nelement = 80
+vpa_L = 12.0 #8.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+#vpa_discretization = "gausslegendre_pseudospectral"
+vz_ngrid = 5
+vz_nelement = 80
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[output]
+ascii_output = true
+
+[numerical_dissipation]
+#moment_dissipation_coefficient = 0.0001
+#moment_dissipation_coefficient = 1.0
+vpa_dissipation_coefficient = 0.002
+#vpa_dissipation_coefficient = 0.1
+#vpa_dissipation_coefficient = 0.2
+#vpa_dissipation_coefficient = 2.0
+#vpa_dissipation_coefficient = 20.0
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index 1cad81bb1..0f0d11fe3 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -1,3 +1,5 @@
+steady_state_residual = true
+converged_residual_value = 1.0e-3
 n_ion_species = 1
 n_neutral_species = 1
 #boltzmann_electron_response = false
@@ -49,8 +51,9 @@ ionization_frequency = 2.0
 constant_ionization_rate = false
 nstep = 40000
 #nstep = 1
-dt = 0.0005
+dt = 5.0e-5
 nwrite = 200
+nwrite_dfns = 10000
 use_semi_lagrange = false
 n_rk_stages = 4
 split_operators = false
@@ -83,6 +86,6 @@ ascii_output = true
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
-#vpa_dissipation_coefficient = 0.2
+vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
-vpa_dissipation_coefficient = 20.0
+#vpa_dissipation_coefficient = 20.0

From 80a9c134fb6149f21b472fb847b1e051caa88104 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 12 Mar 2024 15:18:32 +0000
Subject: [PATCH 148/428] Fix order of arguments to
 electron_kinetic_equation_residual!()

---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index b2a576468..5d113cf3d 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -426,7 +426,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         # TMP FOR TESTING
         #dqpar_dz .= 0.0
         # calculate the residual of the electron kinetic equation for the updated electron pdf
-        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, ddens_dz, upar_ion,
+        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, upar_ion, ddens_dz,
                                             dppar_dz, dqpar_dz, dvth_dz, 
                                             z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                             collisions, num_diss_params, dt_max)

From 64cc4f0109ff404a6077eb121f75948e25670c28 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 12 Mar 2024 15:18:57 +0000
Subject: [PATCH 149/428] Re-activate evolution of ppar with electron kinetic
 equation

---
 .../src/electron_kinetic_equation.jl          | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 5d113cf3d..f7eebb389 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -283,18 +283,18 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
             end
 
             # TMP FOR TESTING -- MAB
-            # #dt_energy = dt_electron
-            # electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
-
-            # # Apply same 'speed up' hack to ppar that we do to the distribution function,
-            # # but without the wpa dependence.
-            # @loop_r_z ir iz begin
-            #     zval = z.grid[iz]
-            #     znorm = 2.0*zval/z.L
-            #     ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
-            #                   (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
-            #                   (1.0 + z_speedup_fac*(1.0 - znorm^2))
-            # end
+            #dt_energy = dt_electron
+            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
+
+            # Apply same 'speed up' hack to ppar that we do to the distribution function,
+            # but without the wpa dependence.
+            @loop_r_z ir iz begin
+                zval = z.grid[iz]
+                znorm = 2.0*zval/z.L
+                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
+                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
+                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
+            end
             begin_r_z_region()
             @loop_r_z ir iz begin
                 fvec.electron_ppar[iz,ir] = ppar[iz,ir]

From c68d98174b5714012d69028ba6eba757617a4e27 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 12 Mar 2024 15:24:15 +0000
Subject: [PATCH 150/428] Include vperp_spectral in arguments to
 initialise_electrons!()

Fixes bad merge.
---
 moment_kinetics/src/initial_conditions.jl | 17 +++++++++--------
 moment_kinetics/src/moment_kinetics.jl    |  9 +++++----
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 90ec68ce1..6d2cf907c 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -274,10 +274,10 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                                  z, r, composition)
 
     initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                          vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
-                          collisions, external_source_settings,
-                          scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
-                          io_input, input_dict)
+                          vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                          vperp_spectral, vpa_spectral, collisions,
+                          external_source_settings, scratch_dummy, scratch, t_input,
+                          num_diss_params, advection_structs, io_input, input_dict)
 
     # moments.electron.dens_updated[] = false
     # # initialise the electron density profile
@@ -357,10 +357,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
 end
 
 function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
-                               collisions, external_source_settings,
-                               scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
-                               io_input, input_dict; restart=false)
+                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                               vperp_spectral, vpa_spectral, collisions,
+                               external_source_settings, scratch_dummy, scratch, t_input,
+                               num_diss_params, advection_structs, io_input, input_dict;
+                               restart=false)
     
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 06ec662dc..8ed626cd3 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -318,10 +318,11 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             # where electrons have a Boltzmann distribution, there is missing information
             # that still needs to be specified for the electrons
             initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                                 vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral, vpa_spectral, 
-                                 collisions, external_source_settings,
-                                 scratch_dummy, scratch, t_input, num_diss_params, advection_structs,
-                                 io_input, input_dict, restart=true)
+                                  vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                                  vperp_spectral, vpa_spectral, collisions,
+                                  external_source_settings, scratch_dummy, scratch,
+                                  t_input, num_diss_params, advection_structs, io_input,
+                                  input_dict, restart=true)
         end
 
         _block_synchronize()

From 9ce1fc1520524aa55498aaad49f563be9dfff2c5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 12 Mar 2024 15:25:11 +0000
Subject: [PATCH 151/428] `moments.electron.upar_updated[] = false` correct
 initial electron upar

---
 moment_kinetics/src/initial_conditions.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6d2cf907c..108b33059 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -367,6 +367,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # initialise the electron density profile
     init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
     # initialise the electron parallel flow profile
+    moments.electron.upar_updated[] = false
     init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
         moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
     # different choices for initialization of electron temperature/pressure/vth depending on whether

From 95c313641ebffa83ce3acb0c7added955576383d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 10:15:49 +0000
Subject: [PATCH 152/428] Tidy up restarting of kinetic electrons from
 Boltzmann

---
 moment_kinetics/src/initial_conditions.jl | 13 ++++++-------
 moment_kinetics/src/load_data.jl          |  8 +++++---
 moment_kinetics/src/moment_kinetics.jl    | 10 ++++------
 3 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 108b33059..3818a3cd6 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -361,7 +361,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                                vperp_spectral, vpa_spectral, collisions,
                                external_source_settings, scratch_dummy, scratch, t_input,
                                num_diss_params, advection_structs, io_input, input_dict;
-                               restart=false)
+                               restart_from_Boltzmann_electrons=false)
     
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
@@ -372,7 +372,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
     # different choices for initialization of electron temperature/pressure/vth depending on whether
     # we are restarting from a previous simulation with Boltzmann electrons or not
-    if restart
+    if restart_from_Boltzmann_electrons
         # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
         # that the electron parallel temperature is constant along the field line and equal to T_e
         moments.electron.temp .= composition.T_e
@@ -445,7 +445,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # arrays need to exist and be otherwise initialised in order to compute the initial
     # electron pdf. The electron arrays will be updated as necessary by
     # initialize_electron_pdf!().
-    if restart == false
+    if !restart_from_Boltzmann_electrons
         initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
         # get the initial electrostatic potential and parallel electric field
         update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
@@ -457,8 +457,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
-                             num_diss_params, t_input.dt, io_input, input_dict,
-                             restart_from_boltzmann=restart)
+                             num_diss_params, t_input.dt, io_input, input_dict)
 
     return nothing
 end
@@ -585,7 +584,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                   vz, z_spectral, vperp_spectral, vpa_spectral, z_advect,
                                   vpa_advect, scratch_dummy, collisions, composition,
                                   geometry, external_source_settings, num_diss_params, dt,
-                                  io_input, input_dict; restart_from_boltzmann=false)
+                                  io_input, input_dict)
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
@@ -593,7 +592,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         begin_serial_region()
         restart_filename = get_default_restart_filename(io_input, "initial_electron";
                                                         error_if_no_file_found=false)
-        if restart_filename === nothing || restart_from_boltzmann
+        if restart_filename === nothing
             # No file to restart from
             previous_runs_info = nothing
             code_time = 0.0
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index b3dead265..628484462 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -542,9 +542,10 @@ Reload pdf and moments from an existing output file.
 """
 function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_prefix_iblock,
                                  time_index, composition, geometry, r, z, vpa, vperp,
-                                 vzeta, vr, vz, initialize_electrons_from_boltzmann=false)
+                                 vzeta, vr, vz)
     code_time = 0.0
     previous_runs_info = nothing
+    restart_had_kinetic_electrons = false
     begin_serial_region()
     @serial_region begin
         fid = open_readonly_output_file(restart_prefix_iblock[1], "dfns";
@@ -720,7 +721,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
                 restart_electron_evolve_ppar = true, true, true
             electron_evolve_density, electron_evolve_upar, electron_evolve_ppar =
                 true, true, true
-            if (pdf.electron !== nothing) && (initialize_electrons_from_boltzmann == false)
+            restart_had_kinetic_electrons = ("f_electron" ∈ keys(dynamic))
+            if pdf.electron !== nothing && restart_had_kinetic_electrons
                 pdf.electron.norm .=
                     reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
                                         r_range, z_range, vperp_range, vpa_range,
@@ -817,7 +819,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions, restart_p
         end
     end
 
-    return code_time, previous_runs_info, time_index
+    return code_time, previous_runs_info, time_index, restart_had_kinetic_electrons
 end
 
 """
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 8ed626cd3..4cf09b4a9 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -296,14 +296,12 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         backup_prefix_iblock = get_prefix_iblock_and_move_existing_file(restart_filename,
                                                                         io_input.output_dir)
 
-        # TMP FOR TESTING -- MAB
-        initialize_electrons_from_boltzmann = true                                                                
         # Reload pdf and moments from an existing output file
-        code_time, previous_runs_info, restart_time_index =
+        code_time, previous_runs_info, restart_time_index, restart_had_kinetic_electrons =
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
-                                    vz, initialize_electrons_from_boltzmann)
+                                    vz)
 
         # Re-initialize the source amplitude here instead of loading it from the restart
         # file so that we can change the settings between restarts.
@@ -313,7 +311,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         # Copy the reloaded values into the `scratch` struct
         initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
 
-        if initialize_electrons_from_boltzmann
+        if !restart_had_kinetic_electrons
             # If we are initializing kinetic electrons using info from a simulation
             # where electrons have a Boltzmann distribution, there is missing information
             # that still needs to be specified for the electrons
@@ -322,7 +320,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
                                   vperp_spectral, vpa_spectral, collisions,
                                   external_source_settings, scratch_dummy, scratch,
                                   t_input, num_diss_params, advection_structs, io_input,
-                                  input_dict, restart=true)
+                                  input_dict, restart_from_Boltzmann_electrons=true)
         end
 
         _block_synchronize()

From ae45762d9643d61d41d81e39ec346cd41432054e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 10:30:51 +0000
Subject: [PATCH 153/428] Delete unused electron-initialization block

The code block was nominally for restarting from Boltzmann electron
simulations, but was not being used and looks (?) only partially
implemented. Seems not to be needed, so remove.
---
 moment_kinetics/src/initial_conditions.jl | 119 +++++++++-------------
 1 file changed, 47 insertions(+), 72 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 3818a3cd6..3d1f67157 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -1346,82 +1346,57 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
 
     if z.bc == "wall"
         begin_r_region()
-        if restart_from_boltzmann
-            # initialize a Maxwellian of the form g = (n_0/n_e) / (sqrt(pi)*vth_e) * exp(-v^2/vth_e^2)
-            # and solve for the values of n_0 satisfies int dvpa g  = 1, given the boundary values of phi
-            # that were obtained from a simulation with Boltzmann electrons.
-            # note that the associated upar_e may not be consistent as a consequence.
-            @loop_r ir begin
-                # first initialize a Maxwellian for the normalized pdf
-                @loop_z iz begin
-                    @loop_vperp ivperp begin
-                        @. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2)
-                    end
-                end
-                # obtain the cutoff velocities at z_min and z_max
-                vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
-                # next apply a density scale factor to ensure that the density at the boundaries is consistent with
-                # a cutoff Maxwellian with the appropriate cutoff velocities
-                density_scale_factor_zmin = 2 / (1 + erf(vpa_crit_zmin))
-                density_scale_factor_zmax = 2 / (1 + erf(vpa_crit_zmax))
+        @loop_r ir begin
+            # Initialise an unshifted Maxwellian as a first step
+            @loop_z iz begin
+                vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
                 @loop_vperp ivperp begin
-                    @. pdf[:,ivperp,1,ir] *= density_scale_factor_zmin
-                    @. pdf[:,ivperp,end,ir] *= density_scale_factor_zmax
+                    @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
                 end
             end
-        else
-            @loop_r ir begin
-                # Initialise an unshifted Maxwellian as a first step
-                @loop_z iz begin
-                    vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
-                    @loop_vperp ivperp begin
-                        @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
-                    end
-                end
+        end
+        # Apply the sheath boundary condition to get cut-off boundary distribution
+        # functions and boundary values of phi
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
+                                                    vperp_spectral, vpa_spectral,
+                                                    vpa_advect, moments,
+                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                    me_over_mi)
+        begin_r_region()
+        @loop_r ir begin
+            # get critical velocities beyond which electrons are lost to the wall
+            #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
+            #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
+            # Blend boundary distribution function into bulk of domain to avoid
+            # discontinuities (as much as possible)
+            blend_fac = 100
+            if z.nrank > 1
+                error("Distributed MPI not supported in this init yet")
             end
-            # Apply the sheath boundary condition to get cut-off boundary distribution
-            # functions and boundary values of phi
-            enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
-                                                        vperp_spectral, vpa_spectral,
-                                                        vpa_advect, moments,
-                                                        num_diss_params.vpa_dissipation_coefficient > 0.0,
-                                                        me_over_mi)
-            begin_r_region()
-            @loop_r ir begin
-                # get critical velocities beyond which electrons are lost to the wall
-                #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
-                #println("vpa_crit_zmin = ", vpa_crit_zmin, " vpa_crit_zmax = ", vpa_crit_zmax)
-                # Blend boundary distribution function into bulk of domain to avoid
-                # discontinuities (as much as possible)
-                blend_fac = 100
-                if z.nrank > 1
-                    error("Distributed MPI not supported in this init yet")
-                end
-                @loop_z_vperp iz ivperp begin
-                    #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
-                    #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
-                    #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
-                    blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
-                    blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
-                    @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
-                                            blend_fac_lower * pdf[:,ivperp,1,ir] +
-                                            blend_fac_upper * pdf[:,ivperp,end,ir]
-                    #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
-                    #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
-                    #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *
-                    #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
-                    #                          tanh(-sharp_fac*(vpa.grid-vpa_crit_zmax)))) #/
-                                            #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
-                                            #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
-                                            #exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
-                                            #exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
-
-                    # ensure that the normalised electron pdf integrates to unity
-                    norm_factor = integrate_over_vspace(pdf[:,ivperp,iz,ir], vpa.wgts)
-                    @. pdf[:,ivperp,iz,ir] /= norm_factor
-                    #println("TMP FOR TESTING -- init electron pdf")
-                    #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
-                end
+            @loop_z_vperp iz ivperp begin
+                #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
+                #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
+                #@. pdf[:,ivperp,iz] = exp(-vpa.grid[:]^2)
+                blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
+                blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
+                @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
+                                        blend_fac_lower * pdf[:,ivperp,1,ir] +
+                                        blend_fac_upper * pdf[:,ivperp,end,ir]
+                #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
+                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
+                #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *
+                #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) *
+                #                          tanh(-sharp_fac*(vpa.grid-vpa_crit_zmax)))) #/
+                                        #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) * tanh(-sharp_fac*vpa_crit_zmin)) /
+                                        #(1 - exp(-blend_fac*(z.grid[iz] - z.grid[end])^2) * tanh(sharp_fac*vpa_crit_zmax)))
+                                        #exp(-((vpa.grid[:] - upar[iz])^2) / vth[iz]^2)
+                                        #exp(-((vpa.grid - upar[iz])^2 + vperp.grid[ivperp]^2) / vth[iz]^2)
+
+                # ensure that the normalised electron pdf integrates to unity
+                norm_factor = integrate_over_vspace(pdf[:,ivperp,iz,ir], vpa.wgts)
+                @. pdf[:,ivperp,iz,ir] /= norm_factor
+                #println("TMP FOR TESTING -- init electron pdf")
+                #@. pdf[:,ivperp,iz] = exp(-2*vpa.grid[:]^2)*exp(-z.grid[iz]^2)
             end
         end
     else

From 63f714e1af0ae7c52e943f5181d383e86e0e6df3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 11:35:00 +0000
Subject: [PATCH 154/428] Update electron vth and temperature with a dedicated
 function

---
 moment_kinetics/src/electron_fluid_equations.jl | 15 +++++++++++++++
 moment_kinetics/src/time_advance.jl             |  7 +++----
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 2c0190aee..0175bf3f7 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -6,6 +6,7 @@ export electron_energy_equation!
 export calculate_electron_qpar!
 export calculate_electron_parallel_friction_force!
 export calculate_electron_qpar_from_pdf!
+export update_electron_vth_temperature!
 
 using ..communication
 using ..looping
@@ -362,4 +363,18 @@ end
 #    end
 #end
 
+function update_electron_vth_temperature!(moments, ppar, dens)
+    begin_r_z_region()
+
+    temp = moments.electron.temp
+    vth = moments.electron.vth
+    @loop_r_z ir iz begin
+        temp[iz,ir] = 2 * ppar[iz,ir] / dens[iz,ir]
+        vth[iz,ir] = sqrt(temp[iz,ir])
+    end
+    moments.electron.temp_updated[] = true
+
+    return nothing
+end
+
 end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 5a65b2bee..380a141a2 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -64,7 +64,7 @@ using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
-using ..electron_fluid_equations: electron_energy_equation!
+using ..electron_fluid_equations: electron_energy_equation!, update_electron_vth_temperature!
 using ..input_structs: braginskii_fluid
 using ..derivatives: derivative_z!
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
@@ -1410,11 +1410,10 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     else
         @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
     end
-    @. moments.electron.temp = 2 * new_scratch.electron_ppar / new_scratch.electron_density
-    moments.electron.temp_updated[] = true
-    @. moments.electron.vth = sqrt(moments.electron.temp)
     # regardless of electron model, electron ppar is now updated
     moments.electron.ppar_updated[] = true
+    update_electron_vth_temperature!(moments, new_scratch.electron_ppar,
+                                     new_scratch.electron_density)
     # calculate the corresponding zed derivatives of the moments
     calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                            num_diss_params, composition.electron_physics)

From 87e5fb54069b3b5855991b2ebbeda8bc1c5f742e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 15:46:08 +0000
Subject: [PATCH 155/428] Make vz resolution consistent with vpa in
 wall+sheath-bc_kinetic_loworder.toml

---
 .../kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml  | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index 0f0d11fe3..bc30e359d 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -73,9 +73,9 @@ vpa_L = 12.0 #8.0
 vpa_bc = "zero"
 vpa_discretization = "chebyshev_pseudospectral"
 #vpa_discretization = "gausslegendre_pseudospectral"
-vz_ngrid = 17
-vz_nelement = 10
-vz_L = 8.0
+vz_ngrid = 80
+vz_nelement = 5
+vz_L = 12.0
 vz_bc = "zero"
 vz_discretization = "chebyshev_pseudospectral"
 

From bd3c962a18a667f0358a2cbb7b3b3b4ab5a273e4 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 15:46:57 +0000
Subject: [PATCH 156/428] Do not animate constraints by default in
 makie_post_process()

---
 .../makie_post_processing/src/makie_post_processing.jl          | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 654c0cba2..c55e0a8d0 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -674,7 +674,7 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
     set_defaults_and_check_section!(
         this_input_dict, "constraints";
         plot=false,
-        animate=true,
+        animate=false,
         it0=this_input_dict["it0"],
         ir0=this_input_dict["ir0"],
         iz0=this_input_dict["iz0"],

From 3c969ee1e879f7716dc4366b4894ef289b5b6f5e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 15:47:59 +0000
Subject: [PATCH 157/428] Some shared-memory parallelisation and initialisation
 fixes

---
 moment_kinetics/src/coordinates.jl            |  7 +++++++
 .../src/electron_fluid_equations.jl           |  8 +++++++
 moment_kinetics/src/initial_conditions.jl     | 21 ++++++++++++-------
 moment_kinetics/src/time_advance.jl           | 14 ++++++++++---
 moment_kinetics/src/velocity_moments.jl       |  5 +++++
 5 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index d2ef23b94..ee208e33f 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -150,6 +150,13 @@ function define_coordinate(input, parallel_io::Bool=false; ignore_MPI=false, ini
         scratch_shared = allocate_shared_float(n_local)
         scratch_shared2 = allocate_shared_float(n_local)
     end
+    # Initialise scratch_shared and scratch_shared2 so that the debug checks do not
+    # complain when they get printed by `println(io, all_inputs)` in mk_input().
+    if block_rank[] == 0
+        scratch_shared .= NaN
+        scratch_shared2 .= NaN
+    end
+    _block_synchronize()
     # scratch_2d is an array used for intermediate calculations requiring ngrid x nelement entries
     scratch_2d = allocate_float(input.ngrid, input.nelement_local)
     # struct containing the advection speed options/inputs for this coordinate
diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 0175bf3f7..9febf93f9 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -305,6 +305,13 @@ function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth
                 electron_pdf = pdf
             end
             calculate_electron_qpar_from_pdf!(qpar_e, ppar_e, vth_e, electron_pdf, vpa)
+        else
+            begin_r_z_region()
+            # qpar_e is not used. Initialize to 0.0 to avoid failure of
+            # @debug_track_initialized check
+            @loop_r_z ir iz begin
+                qpar_e[iz,ir] = 0.0
+            end
         end
     end
     # qpar has been updated
@@ -318,6 +325,7 @@ defined as qpar = 2 * ppar * vth * int dwpa (pdf * wpa^3)
 """
 function calculate_electron_qpar_from_pdf!(qpar, ppar, vth, pdf, vpa)
     # specialise to 1D for now
+    begin_r_z_region()
     ivperp = 1
     @loop_r_z ir iz begin
         @views qpar[iz, ir] = 2*ppar[iz,ir]*vth[iz,ir]*integrate_over_vspace(pdf[:, ivperp, iz, ir], vpa.grid.^3, vpa.wgts)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 3d1f67157..9d0b396d2 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -373,16 +373,20 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # different choices for initialization of electron temperature/pressure/vth depending on whether
     # we are restarting from a previous simulation with Boltzmann electrons or not
     if restart_from_Boltzmann_electrons
-        # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
-        # that the electron parallel temperature is constant along the field line and equal to T_e
-        moments.electron.temp .= composition.T_e
-        # the thermal speed is related to the temperature by vth_e / v_ref = sqrt((T_e/T_ref) / (m_e/m_ref))
-        moments.electron.vth .= sqrt(composition.T_e / composition.me_over_mi)
-        # ppar = 0.5 * n * T, so we can calculate the parallel pressure from the density and T_e
-        moments.electron.ppar .= 0.5 * moments.electron.dens * composition.T_e
+        begin_serial_region()
+        @serial_region begin
+            # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
+            # that the electron parallel temperature is constant along the field line and equal to T_e
+            moments.electron.temp .= composition.T_e
+            # the thermal speed is related to the temperature by vth_e / v_ref = sqrt((T_e/T_ref) / (m_e/m_ref))
+            moments.electron.vth .= sqrt(composition.T_e / composition.me_over_mi)
+            # ppar = 0.5 * n * T, so we can calculate the parallel pressure from the density and T_e
+            moments.electron.ppar .= 0.5 * moments.electron.dens * composition.T_e
+        end
     else
         # initialise the electron thermal speed profile
         init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
+        begin_r_z_region()
         # calculate the electron temperature from the thermal speed
         @loop_r_z ir iz begin
             moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
@@ -732,8 +736,8 @@ function init_vth!(vth, z, r, spec, n_species)
                 end
             end
         end
-        @. vth = sqrt(vth)
     end
+    @. vth = sqrt(vth)
     return nothing
 end
 
@@ -936,6 +940,7 @@ for now the only initialisation option for the temperature is constant in z.
 returns vth0 = sqrt(2*Ts/Te)
 """
 function init_electron_vth!(vth_e, vth_i, T_e, me_over_mi, z)
+    begin_r_z_region()
     # @loop_r_z ir iz begin
     #     vth_e[iz,ir] = sqrt(T_e)
     # end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 380a141a2..9b4678576 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -495,13 +495,18 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
         calculate_electron_upar_from_charge_conservation!(moments.electron.upar, moments.electron.upar_updated,
                                                           moments.electron.dens, moments.ion.upar, moments.ion.dens,
                                                           composition.electron_physics, r, z)
+        begin_serial_region()
         # compute the updated electron temperature
         # NB: not currently necessary, as initial vth is not directly dependent on ion quantities
-        @. moments.electron.temp = moments.electron.vth^2
+        @serial_region begin
+            @. moments.electron.temp = moments.electron.vth^2
+        end
         # as the electron temperature has now been updated, set the appropriate flag
         moments.electron.temp_updated[] = true
         # compute the updated electron parallel pressure
-        @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
+        @serial_region begin
+            @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
+        end
         # as the electron ppar has now been updated, set the appropriate flag
         moments.electron.ppar_updated[] = true
         # calculate the zed derivative of the initial electron temperature, potentially
@@ -1408,7 +1413,10 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                 + rk_coefs[2]*old_scratch.electron_ppar[iz,ir] + rk_coefs[3]*new_scratch.electron_ppar[iz,ir])
         end
     else
-        @. new_scratch.electron_ppar = 0.5 * new_scratch.electron_density * moments.electron.vth^2
+        @loop_r_z ir iz begin
+            new_scratch.electron_ppar[iz,ir] = 0.5 * new_scratch.electron_density[iz,ir] *
+                                               moments.electron.vth[iz,ir]^2
+        end
     end
     # regardless of electron model, electron ppar is now updated
     moments.electron.ppar_updated[] = true
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 3770c92fb..b0230351c 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -253,6 +253,11 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     constraints_A_coefficient = allocate_shared_float(nz, nr)
     constraints_B_coefficient = allocate_shared_float(nz, nr)
     constraints_C_coefficient = allocate_shared_float(nz, nr)
+    @serial_region begin
+        constraints_A_coefficient .= 1.0
+        constraints_B_coefficient .= 0.0
+        constraints_C_coefficient .= 0.0
+    end
 
     # return struct containing arrays needed to update moments
     return moments_electron_substruct(density, density_updated, parallel_flow,

From a233eeb156c6657a94f9f545bfdf3121a6c2342b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 16:42:01 +0000
Subject: [PATCH 158/428] Remove unnecessary uses of @views in
 boundary_conditions.jl

---
 moment_kinetics/src/boundary_conditions.jl | 39 +++++++++++-----------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index adba3654e..61dff1a08 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -36,25 +36,28 @@ function enforce_boundary_conditions!(f, f_r_bc, density, upar, ppar, moments, v
     end
     if vperp.n > 1
         begin_s_r_z_vpa_region()
-        @views enforce_vperp_boundary_condition!(f, vperp.bc, vperp, vperp_spectral,
+        enforce_vperp_boundary_condition!(f, vperp.bc, vperp, vperp_spectral,
                              vperp_adv, vperp_diffusion)
     end
     if z.n > 1
         begin_s_r_vperp_vpa_region()
         # enforce the z BC on the evolved velocity space moments of the pdf
-        @views enforce_z_boundary_condition_moments!(density, moments, z_bc)
-        @views enforce_z_boundary_condition!(f, density, upar, ppar, moments, z_bc, z_adv, z,
-                                             vperp, vpa, composition,
-                                             scratch_dummy.buffer_vpavperprs_1, scratch_dummy.buffer_vpavperprs_2,
-                                             scratch_dummy.buffer_vpavperprs_3, scratch_dummy.buffer_vpavperprs_4)
+        enforce_z_boundary_condition_moments!(density, moments, z_bc)
+        enforce_z_boundary_condition!(f, density, upar, ppar, moments, z_bc, z_adv, z,
+                                      vperp, vpa, composition,
+                                      scratch_dummy.buffer_vpavperprs_1,
+                                      scratch_dummy.buffer_vpavperprs_2,
+                                      scratch_dummy.buffer_vpavperprs_3,
+                                      scratch_dummy.buffer_vpavperprs_4)
 
     end
     if r.n > 1
         begin_s_z_vperp_vpa_region()
-        @views enforce_r_boundary_condition!(f, f_r_bc, r_bc, r_adv, vpa, vperp, z, r, composition,
-            scratch_dummy.buffer_vpavperpzs_1, scratch_dummy.buffer_vpavperpzs_2,
-            scratch_dummy.buffer_vpavperpzs_3, scratch_dummy.buffer_vpavperpzs_4,
-            r_diffusion)
+        enforce_r_boundary_condition!(f, f_r_bc, r_bc, r_adv, vpa, vperp, z, r,
+                                      composition, scratch_dummy.buffer_vpavperpzs_1,
+                                      scratch_dummy.buffer_vpavperpzs_2,
+                                      scratch_dummy.buffer_vpavperpzs_3,
+                                      scratch_dummy.buffer_vpavperpzs_4, r_diffusion)
     end
 end
 function enforce_boundary_conditions!(fvec_out::scratch_pdf, moments, f_r_bc, vpa_bc,
@@ -83,8 +86,7 @@ function enforce_r_boundary_condition!(f::AbstractArray{mk_float,5}, f_r_bc, bc:
             end1[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,1,is]
             end2[ivpa,ivperp,iz,is] = f[ivpa,ivperp,iz,nr,is]
         end
-        @views reconcile_element_boundaries_MPI!(f,
-            end1, end2,	buffer1, buffer2, r)
+        reconcile_element_boundaries_MPI!(f, end1, end2, buffer1, buffer2, r)
     end
 
     # 'periodic' BC enforces periodicity by taking the average of the boundary points
@@ -133,8 +135,7 @@ function enforce_z_boundary_condition!(pdf, density, upar, ppar, moments, bc::St
             end2[ivpa,ivperp,ir,is] = pdf[ivpa,ivperp,nz,ir,is]
         end
         # check on periodic bc happens inside this call below
-        @views reconcile_element_boundaries_MPI!(pdf,
-            end1, end2,	buffer1, buffer2, z)
+        reconcile_element_boundaries_MPI!(pdf, end1, end2, buffer1, buffer2, z)
     end
     # define a zero that accounts for finite precision
     zero = 1.0e-14
@@ -233,7 +234,7 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
     # f_initial contains the initial condition for enforcing a fixed-boundary-value condition
     if z.n > 1
         begin_sn_r_vzeta_vr_vz_region()
-        @views enforce_neutral_z_boundary_condition!(f_neutral, density_neutral, uz_neutral,
+        enforce_neutral_z_boundary_condition!(f_neutral, density_neutral, uz_neutral,
             pz_neutral, moments, density_ion, upar_ion, Er, boundary_distributions,
             z_adv, z, vzeta, vr, vz, composition, geometry,
             scratch_dummy.buffer_vzvrvzetarsn_1, scratch_dummy.buffer_vzvrvzetarsn_2,
@@ -241,7 +242,7 @@ function enforce_neutral_boundary_conditions!(f_neutral, f_ion,
     end
     if r.n > 1
         begin_sn_z_vzeta_vr_vz_region()
-        @views enforce_neutral_r_boundary_condition!(f_neutral, boundary_distributions.pdf_rboundary_neutral,
+        enforce_neutral_r_boundary_condition!(f_neutral, boundary_distributions.pdf_rboundary_neutral,
                                     r_adv, vz, vr, vzeta, z, r, composition,
                                     scratch_dummy.buffer_vzvrvzetazsn_1, scratch_dummy.buffer_vzvrvzetazsn_2,
                                     scratch_dummy.buffer_vzvrvzetazsn_3, scratch_dummy.buffer_vzvrvzetazsn_4,
@@ -265,8 +266,7 @@ function enforce_neutral_r_boundary_condition!(f::AbstractArray{mk_float,6},
             end1[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,1,isn]
             end2[ivz,ivr,ivzeta,iz,isn] = f[ivz,ivr,ivzeta,iz,nr,isn]
         end
-        @views reconcile_element_boundaries_MPI!(f,
-            end1, end2,	buffer1, buffer2, r)
+        reconcile_element_boundaries_MPI!(f, end1, end2, buffer1, buffer2, r)
     end
     # 'periodic' BC enforces periodicity by taking the average of the boundary points
     # local case only when no communication required
@@ -316,8 +316,7 @@ function enforce_neutral_z_boundary_condition!(pdf, density, uz, pz, moments, de
             end2[ivz,ivr,ivzeta,ir,isn] = pdf[ivz,ivr,ivzeta,nz,ir,isn]
         end
         # check on periodic bc occurs within this call below
-        @views reconcile_element_boundaries_MPI!(pdf,
-            end1, end2,	buffer1, buffer2, z)
+        reconcile_element_boundaries_MPI!(pdf, end1, end2, buffer1, buffer2, z)
     end
 
     zero = 1.0e-14

From 95adbfe1f0fe9c5f654d8d36b021a291cca9c312 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 17:12:13 +0000
Subject: [PATCH 159/428] Respect ignore_MPI arg when initializing
 scratch_shared/scratch_shared2

---
 moment_kinetics/src/coordinates.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index ee208e33f..ee01f93dd 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -156,7 +156,9 @@ function define_coordinate(input, parallel_io::Bool=false; ignore_MPI=false, ini
         scratch_shared .= NaN
         scratch_shared2 .= NaN
     end
-    _block_synchronize()
+    if !ignore_MPI
+        _block_synchronize()
+    end
     # scratch_2d is an array used for intermediate calculations requiring ngrid x nelement entries
     scratch_2d = allocate_float(input.ngrid, input.nelement_local)
     # struct containing the advection speed options/inputs for this coordinate

From c3d5c5259d82f3367f97d27553b1def20e56b13b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 17:46:43 +0000
Subject: [PATCH 160/428] Filter out size-1 dimensions from combinations to
 check in debug tests

If dimension has size 1 for some input it cannot be parallelised, so
there is no need to check it (for that case). By skipping size 1
dimensions, we can reduce the number of combinations that need debug
tests, which speeds up the tests.
---
 .../debug_test/runtest_template.jl            | 34 ++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/debug_test/runtest_template.jl b/moment_kinetics/debug_test/runtest_template.jl
index 779be8929..6f1cb0b4e 100644
--- a/moment_kinetics/debug_test/runtest_template.jl
+++ b/moment_kinetics/debug_test/runtest_template.jl
@@ -1,7 +1,8 @@
 using moment_kinetics: setup_moment_kinetics, cleanup_moment_kinetics!
 using moment_kinetics.time_advance: time_advance!
 using moment_kinetics.communication
-using moment_kinetics.looping: dimension_combinations, anyv_dimension_combinations
+using moment_kinetics.looping: all_dimensions, dimension_combinations,
+                               anyv_dimension_combinations
 using moment_kinetics.Glob
 using moment_kinetics.Primes
 
@@ -71,8 +72,39 @@ function runtests(; restart=false)
             else
                 dims_to_test = debug_loop_type
             end
+            for d ∈ all_dimensions
+                nelement_name = "$(d)_nelement"
+                if nelement_name ∈ keys(input)
+                    nelement = input[nelement_name]
+                elseif d ∈ (:vperp, :vzeta, :vr)
+                    nelement = 1
+                else
+                    # Dummy value, here it only matters if this is 1 or greater than 1
+                    nelement = 2
+                end
+
+                ngrid_name = "$(d)_ngrid"
+                if ngrid_name ∈ keys(input)
+                    ngrid = input[ngrid_name]
+                elseif d ∈ (:vperp, :vzeta, :vr)
+                    ngrid = 1
+                else
+                    # Dummy value, here it only matters if this is 1 or greater than 1
+                    ngrid = 2
+                end
+
+                if nelement == 1 && ngrid == 1
+                    # Dimension has only one point, so cannot be parallelised - no need to
+                    # test
+                    dims_to_test = Tuple(x for x ∈ dims_to_test if x != d)
+                end
+            end
 
             ndims = length(dims_to_test)
+            if ndims == 0
+                # No dimensions to test here
+                continue
+            end
             for i ∈ 1:(ndims+n_factors-1)÷n_factors
                 debug_loop_parallel_dims =
                     dims_to_test[(i-1)*n_factors+1:min(i*n_factors, ndims)]

From 6a86f05b6aeee6e4b06c48dfed3d24801c51ca53 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 20:49:16 +0000
Subject: [PATCH 161/428] Tests must use ignore_MPI=true when loading or
 defining coordinates

...because loading or defining coordinates is done in an
`@serial_region`, so the tests would hang without `ignore_MPI=true`.
---
 moment_kinetics/test/Krook_collisions_tests.jl             | 4 ++--
 moment_kinetics/test/fokker_planck_time_evolution_tests.jl | 4 ++--
 moment_kinetics/test/harrisonthompson.jl                   | 4 ++--
 moment_kinetics/test/nonlinear_sound_wave_tests.jl         | 4 ++--
 moment_kinetics/test/sound_wave_tests.jl                   | 4 ++--
 moment_kinetics/test/wall_bc_tests.jl                      | 2 +-
 6 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/test/Krook_collisions_tests.jl b/moment_kinetics/test/Krook_collisions_tests.jl
index ef2528814..eef5f3d3b 100644
--- a/moment_kinetics/test/Krook_collisions_tests.jl
+++ b/moment_kinetics/test/Krook_collisions_tests.jl
@@ -226,7 +226,7 @@ function run_test(test_input, rtol, atol; args...)
             # load velocity moments data
             n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
-            z, z_spectral = load_coordinate_data(fid, "z")
+            z, z_spectral = load_coordinate_data(fid, "z"; ignore_MPI=true)
 
             close(fid)
             
@@ -236,7 +236,7 @@ function run_test(test_input, rtol, atol; args...)
             # load particle distribution function (pdf) data
             f_ion_vpavperpzrst = load_pdf_data(fid)
             f_neutral_vzvrvzetazrst = load_neutral_pdf_data(fid)
-            vpa, vpa_spectral = load_coordinate_data(fid, "vpa")
+            vpa, vpa_spectral = load_coordinate_data(fid, "vpa"; ignore_MPI=true)
 
             close(fid)
             
diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index 12302d51d..2f4281a41 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -253,8 +253,8 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # open the netcdf file containing pdf data
             fid = open_readonly_output_file(path, "dfns")
             # load coordinates
-            vpa, vpa_spectral = load_coordinate_data(fid, "vpa")
-            vperp, vperp_spectral = load_coordinate_data(fid, "vperp")
+            vpa, vpa_spectral = load_coordinate_data(fid, "vpa"; ignore_MPI=true)
+            vperp, vperp_spectral = load_coordinate_data(fid, "vperp"; ignore_MPI=true)
 
             # load particle distribution function (pdf) data
             f_ion_vpavperpzrst = load_pdf_data(fid)
diff --git a/moment_kinetics/test/harrisonthompson.jl b/moment_kinetics/test/harrisonthompson.jl
index 59ea8aeed..1f196e32b 100644
--- a/moment_kinetics/test/harrisonthompson.jl
+++ b/moment_kinetics/test/harrisonthompson.jl
@@ -197,8 +197,8 @@ function run_test(test_input, analytic_rtol, analytic_atol, expected_phi,
             fid = open_readonly_output_file(path,"moments")
 
             # load space-time coordinate data
-            z, z_spectral, z_chunk_size = load_coordinate_data(fid, "z")
-            r, r_spectral, r_chunk_size = load_coordinate_data(fid, "r")
+            z, z_spectral, z_chunk_size = load_coordinate_data(fid, "z"; ignore_MPI=true)
+            r, r_spectral, r_chunk_size = load_coordinate_data(fid, "r"; ignore_MPI=true)
             ntime, time = load_time_data(fid)
             n_ion_species, n_neutral_species = load_species_data(fid)
             
diff --git a/moment_kinetics/test/nonlinear_sound_wave_tests.jl b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
index f6e5d7334..52ef17409 100644
--- a/moment_kinetics/test/nonlinear_sound_wave_tests.jl
+++ b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
@@ -91,7 +91,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # load velocity moments data
             n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, qpar_ion_zrst, v_t_ion_zrst = load_ion_moments_data(fid)
             n_neutral_zrst, upar_neutral_zrst, ppar_neutral_zrst, qpar_neutral_zrst, v_t_neutral_zrst = load_neutral_particle_moments_data(fid)
-            z, z_spectral = load_coordinate_data(fid, "z")
+            z, z_spectral = load_coordinate_data(fid, "z"; ignore_MPI=true)
 
             close(fid)
             
@@ -101,7 +101,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
             # load particle distribution function (pdf) data
             f_ion_vpavperpzrst = load_pdf_data(fid)
             f_neutral_vzvrvzetazrst = load_neutral_pdf_data(fid)
-            vpa, vpa_spectral = load_coordinate_data(fid, "vpa")
+            vpa, vpa_spectral = load_coordinate_data(fid, "vpa"; ignore_MPI=true)
 
             close(fid)
             
diff --git a/moment_kinetics/test/sound_wave_tests.jl b/moment_kinetics/test/sound_wave_tests.jl
index 9ffdf25ed..cb88d2e64 100644
--- a/moment_kinetics/test/sound_wave_tests.jl
+++ b/moment_kinetics/test/sound_wave_tests.jl
@@ -178,8 +178,8 @@ function run_test(test_input, analytic_frequency, analytic_growth_rate,
             fid = open_readonly_output_file(path,"moments")
 
             # load space-time coordinate data
-            z, z_spectral = load_coordinate_data(fid, "z")
-            r, r_spectral = load_coordinate_data(fid, "r")
+            z, z_spectral = load_coordinate_data(fid, "z"; ignore_MPI=true)
+            r, r_spectral = load_coordinate_data(fid, "r"; ignore_MPI=true)
             n_ion_species, n_neutral_species = load_species_data(fid)
             ntime, time = load_time_data(fid)
             
diff --git a/moment_kinetics/test/wall_bc_tests.jl b/moment_kinetics/test/wall_bc_tests.jl
index b87e26fd9..f4228f0d9 100644
--- a/moment_kinetics/test/wall_bc_tests.jl
+++ b/moment_kinetics/test/wall_bc_tests.jl
@@ -216,7 +216,7 @@ function run_test(test_input, expected_phi, tolerance; args...)
                            test_input["z_nelement"], nrank_per_block, irank, 1.0,
                            test_input["z_discretization"], "", cheb_option, test_input["z_bc"],
                            adv_input, comm, test_input["z_element_spacing_option"])
-        z, z_spectral = define_coordinate(input)
+        z, z_spectral = define_coordinate(input; ignore_MPI=true)
 
         # Cross comparison of all discretizations to same benchmark
         if test_input["z_element_spacing_option"] == "uniform" 

From 2738d41fb69af62fbb34e193d8399ba985f7ac21 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 22:09:34 +0000
Subject: [PATCH 162/428] Fix restarting when not using kinetic electrons

---
 moment_kinetics/src/moment_kinetics.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 4e224ede8..4ff037e01 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -312,7 +312,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         # Copy the reloaded values into the `scratch` struct
         initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
 
-        if !restart_had_kinetic_electrons
+        if composition.electron_physics == kinetic_electrons && !restart_had_kinetic_electrons
             # If we are initializing kinetic electrons using info from a simulation
             # where electrons have a Boltzmann distribution, there is missing information
             # that still needs to be specified for the electrons

From 79e7cf8b9ba3d928eb4a52f7c81e0d304a9ceaeb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 15 Mar 2024 22:25:48 +0000
Subject: [PATCH 163/428] Template the type of shared_scratch in coordinate
 struct

When running in debug mode, the arrays might be DebugMPISharedArray (for
the runs) or Array (when loading coordinates with ignore_MPI=true for
testing or post-processing).
---
 moment_kinetics/src/coordinates.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index ee01f93dd..da7c7b85b 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -23,7 +23,7 @@ using MPI
 """
 structure containing basic information related to coordinates
 """
-struct coordinate
+struct coordinate{T <: AbstractVector{mk_float}}
     # name is the name of the variable associated with this coordiante
     name::String
     # n_global is the total number of grid points associated with this coordinate
@@ -80,10 +80,10 @@ struct coordinate
     scratch3::Array{mk_float,1}
     # scratch_shared is a shared-memory array used for intermediate calculations requiring
     # n entries
-    scratch_shared::MPISharedArray{mk_float,1}
+    scratch_shared::T
     # scratch_shared2 is a shared-memory array used for intermediate calculations requiring
     # n entries
-    scratch_shared2::MPISharedArray{mk_float,1}
+    scratch_shared2::T
     # scratch_2d and scratch2_2d are arrays used for intermediate calculations requiring
     # ngrid x nelement entries
     scratch_2d::Array{mk_float,2}

From 431d0c7da8a590ebee34fa6d717760ca9b077a2e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 15:27:04 +0000
Subject: [PATCH 164/428] Fix / tidy up initialization of electrons after merge

---
 moment_kinetics/src/initial_conditions.jl | 47 ++++-------------------
 moment_kinetics/src/input_structs.jl      |  1 +
 moment_kinetics/src/moment_kinetics.jl    | 19 ++-------
 moment_kinetics/src/time_advance.jl       | 18 ++++++++-
 4 files changed, 29 insertions(+), 56 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index d8cf86a44..7466064ad 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -219,12 +219,6 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
-    initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                          vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
-                          vperp_spectral, vpa_spectral, collisions,
-                          external_source_settings, scratch_dummy, scratch, t_input,
-                          num_diss_params, advection_structs, io_input, input_dict)
-
     # moments.electron.dens_updated[] = false
     # # initialise the electron density profile
     # init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
@@ -378,6 +372,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
+    moments.electron.qpar_updated[] = false
     calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
         moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
         collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
@@ -396,7 +391,13 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # electron pdf. The electron arrays will be updated as necessary by
     # initialize_electron_pdf!().
     if !restart_from_Boltzmann_electrons
-        initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
+        begin_serial_region()
+        @serial_region begin
+            scratch[1].electron_density .= moments.electron.dens
+            scratch[1].electron_upar .= moments.electron.upar
+            scratch[1].electron_ppar .= moments.electron.ppar
+            scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
+        end
         # get the initial electrostatic potential and parallel electric field
         update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
     end
@@ -412,38 +413,6 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     return nothing
 end
 
-"""
-initialize the array of structs containing scratch arrays for the normalised pdf and low-order moments
-that may be evolved separately via fluid equations
-"""
-function initialize_scratch_arrays!(scratch, moments, pdf, n_rk_stages)
-    # populate each of the structs
-    begin_serial_region()
-    @serial_region begin
-        for istage ∈ 1:n_rk_stages+1
-            # initialise the scratch arrays for the ion pdf and velocity moments
-            scratch[istage].pdf .= pdf.ion.norm
-            scratch[istage].density .= moments.ion.dens
-            scratch[istage].upar .= moments.ion.upar
-            scratch[istage].ppar .= moments.ion.ppar
-            # initialise the scratch arrays for the electron pdf and velocity moments
-            if pdf.electron !== nothing
-                scratch[istage].pdf_electron .= pdf.electron.norm
-            end
-            scratch[istage].electron_density .= moments.electron.dens
-            scratch[istage].electron_upar .= moments.electron.upar
-            scratch[istage].electron_ppar .= moments.electron.ppar
-            scratch[istage].electron_temp .= moments.electron.temp
-            # initialise the scratch arrays for the neutral velocity moments and pdf
-            scratch[istage].pdf_neutral .= pdf.neutral.norm
-            scratch[istage].density_neutral .= moments.neutral.dens
-            scratch[istage].uz_neutral .= moments.neutral.uz
-            scratch[istage].pz_neutral .= moments.neutral.pz
-        end
-    end
-    return nothing
-end
-
 """
 """
 function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z, vperp,
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index a9351a1f9..982575bfd 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -2,6 +2,7 @@
 """
 module input_structs
 
+export advance_info
 export evolve_moments_options
 export time_info
 export advection_input, advection_input_mutable
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 4ecd4395a..8bba7e0ff 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -87,8 +87,7 @@ using .communication: _block_synchronize
 using .debugging
 using .external_sources
 using .input_structs
-using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!,
-                           initialize_electrons!
+using .initial_conditions: allocate_pdf_and_moments, init_pdf_and_moments!
 using .load_data: reload_evolving_fields!
 using .looping
 using .moment_constraints: hard_force_moment_constraints!
@@ -288,6 +287,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         dt = nothing
         dt_before_last_fail = nothing
         previous_runs_info = nothing
+        restart_had_kinetic_electrons = false
     else
         restarting = true
 
@@ -313,18 +313,6 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         initialize_external_source_amplitude!(moments, external_source_settings, vperp,
                                               vzeta, vr, composition.n_neutral_species)
 
-        if composition.electron_physics == kinetic_electrons && !restart_had_kinetic_electrons
-            # If we are initializing kinetic electrons using info from a simulation
-            # where electrons have a Boltzmann distribution, there is missing information
-            # that still needs to be specified for the electrons
-            initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                                  vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
-                                  vperp_spectral, vpa_spectral, collisions,
-                                  external_source_settings, scratch_dummy, scratch,
-                                  t_input, num_diss_params, advection_structs, io_input,
-                                  input_dict, restart_from_Boltzmann_electrons=true)
-        end
-
         _block_synchronize()
     end
 
@@ -342,7 +330,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             r_spectral, composition, drive_input, moments, t_input, code_time, dt,
             dt_before_last_fail, collisions, species, geometry, boundary_distributions,
             external_source_settings, num_diss_params, manufactured_solns_input,
-            advection_structs, scratch_dummy, restarting)
+            advection_structs, scratch_dummy, io_input, restarting,
+            restart_had_kinetic_electrons, input_dict)
 
     # This is the closest we can get to the end time of the setup before writing it to the
     # output file
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index c3467b9a7..cef4f29a6 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -15,6 +15,7 @@ using ..communication
 using ..communication: _block_synchronize
 using ..debugging
 using ..file_io: write_data_to_ascii, write_all_moments_data_to_binary, write_all_dfns_data_to_binary, debug_dump
+using ..initial_conditions: initialize_electrons!
 using ..looping
 using ..moment_kinetics_structs: scratch_pdf
 using ..velocity_moments: update_moments!, update_moments_neutral!, reset_moments_status!
@@ -28,7 +29,7 @@ using ..velocity_moments: update_chodura!
 using ..velocity_grid_transforms: vzvrvzeta_to_vpavperp!, vpavperp_to_vzvrvzeta!
 using ..boundary_conditions: enforce_boundary_conditions!
 using ..boundary_conditions: enforce_neutral_boundary_conditions!
-using ..input_structs: advance_info, time_info
+using ..input_structs
 using ..moment_constraints: hard_force_moment_constraints!,
                             hard_force_moment_constraints_neutral!
 using ..advection: setup_advection
@@ -276,7 +277,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                              collisions, species, geometry, boundary_distributions,
                              external_source_settings, num_diss_params,
                              manufactured_solns_input, advection_structs, scratch_dummy,
-                             restarting)
+                             io_input, restarting, restart_had_kinetic_electrons,
+                             input_dict)
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
@@ -395,6 +397,18 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
     else
         fp_arrays = nothing
     end
+
+    # Now that `t_params` and `scratch` have been created, initialize electrons if
+    # necessary
+    if !restarting || (composition.electron_physics == kinetic_electrons &&
+                       !restart_had_kinetic_electrons)
+        initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
+                              vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                              vperp_spectral, vpa_spectral, collisions,
+                              external_source_settings, scratch_dummy, scratch, t_params,
+                              num_diss_params, advection_structs, io_input, input_dict)
+    end
+
     # update the derivatives of the electron moments as these may be needed when
     # computing the electrostatic potential (and components of the electric field)
     calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 

From 0ab74bcb3ade433351913fa62593446a1ec04b01 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 15:27:40 +0000
Subject: [PATCH 165/428] Remove commented-out electron initialization code

This code was moved into initialize_electrons!() function.
---
 moment_kinetics/src/initial_conditions.jl | 74 -----------------------
 1 file changed, 74 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 7466064ad..b5df32128 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -219,80 +219,6 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
     init_boundary_distributions!(boundary_distributions, pdf, vz, vr, vzeta, vpa, vperp,
                                  z, r, composition)
 
-    # moments.electron.dens_updated[] = false
-    # # initialise the electron density profile
-    # init_electron_density!(moments.electron.dens, moments.electron.dens_updated, moments.ion.dens)
-    # # initialise the electron parallel flow profile
-    # init_electron_upar!(moments.electron.upar, moments.electron.upar_updated, moments.electron.dens, 
-    #     moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
-    # # initialise the electron thermal speed profile
-    # init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
-    # # calculate the electron temperature from the thermal speed
-    # @loop_r_z ir iz begin
-    #     moments.electron.temp[iz,ir] = composition.me_over_mi * moments.electron.vth[iz,ir]^2
-    # end
-    # # the electron temperature has now been updated
-    # moments.electron.temp_updated[] = true
-    # # calculate the electron parallel pressure from the density and temperature
-    # @loop_r_z ir iz begin
-    #     moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
-    # end
-    # # the electron parallel pressure now been updated
-    # moments.electron.ppar_updated[] = true
-
-    # # calculate the zed derivative of the initial electron density
-    # @views derivative_z!(moments.electron.ddens_dz, moments.electron.dens, 
-    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # # calculate the zed derivative of the initial electron temperature
-    # @views derivative_z!(moments.electron.dT_dz, moments.electron.temp, 
-    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # # calculate the zed derivative of the initial electron thermal speed
-    # @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, 
-    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # # calculate the zed derivative of the initial electron parallel pressure
-    # @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
-    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # if composition.electron_physics == kinetic_electrons
-    #     # Initialise the array for the electron pdf
-    #     init_electron_pdf_over_density_and_boundary_phi!(
-    #         pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
-    #         moments.electron.vth, z, vpa, vperp, vpa_spectral, composition.me_over_mi)
-    # end
-    # # calculate the electron parallel heat flux;
-    # # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
-    # calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
-    #     moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
-    #     collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
-    # # calculate the zed derivative of the initial electron parallel heat flux
-    # @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
-    #     scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-    #     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    # # calculate the electron-ion parallel friction force
-    # calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
-    #     moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
-    #     composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
-    
-    # # initialize the scratch arrays containing pdfs and moments for the first RK stage
-    # # the electron pdf is yet to be initialised but with the current code logic, the scratch
-    # # arrays need to exist and be otherwise initialised in order to compute the initial
-    # # electron pdf. The electron arrays will be updated as necessary by
-    # # initialize_electron_pdf!().
-    # initialize_scratch_arrays!(scratch, moments, pdf, t_input.n_rk_stages)
-    # # get the initial electrostatic potential and parallel electric field
-    # update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
-
-    # # initialize the electron pdf that satisfies the electron kinetic equation
-    # initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
-    #                          vzeta, vr, vz, z_spectral, vpa_spectral,
-    #                          advection_structs.electron_z_advect,
-    #                          advection_structs.electron_vpa_advect, scratch_dummy,
-    #                          collisions, composition, geometry, external_source_settings,
-    #                          num_diss_params, t_input.dt, io_input, input_dict)
-
     return nothing
 end
 

From 60b712968fc92db25bba5f9d0c09fbc4c50ed5ba Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 17:36:59 +0000
Subject: [PATCH 166/428] Fix lookup of dfn vars from makie_post_processing
 input

Fixes changes from 399e521acca84d50d5635affa004223bcc6128b4.
---
 .../src/makie_post_processing.jl                     | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 78cdd230e..d4598fbf0 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1368,7 +1368,7 @@ for dim ∈ one_dimension_combinations
                      if input === nothing
                          if run_info[1].dfns
                              if var_name ∈ keys(input_dict_dfns)
-                                 input = input_dict[var_name]
+                                 input = input_dict_dfns[var_name]
                              else
                                  input = input_dict_dfns
                              end
@@ -1430,7 +1430,7 @@ for dim ∈ one_dimension_combinations
                  if input === nothing
                      if run_info.dfns
                          if var_name ∈ keys(input_dict_dfns)
-                             input = input_dict[var_name]
+                             input = input_dict_dfns[var_name]
                          else
                              input = input_dict_dfns
                          end
@@ -1615,7 +1615,7 @@ for (dim1, dim2) ∈ two_dimension_combinations
                  if input === nothing
                      if run_info.dfns
                          if var_name ∈ keys(input_dict_dfns)
-                             input = input_dict[var_name]
+                             input = input_dict_dfns[var_name]
                          else
                              input = input_dict_dfns
                          end
@@ -1798,7 +1798,7 @@ for dim ∈ one_dimension_combinations_no_t
                      if input === nothing
                          if run_info[1].dfns
                              if var_name ∈ keys(input_dict_dfns)
-                                 input = input_dict[var_name]
+                                 input = input_dict_dfns[var_name]
                              else
                                  input = input_dict_dfns
                              end
@@ -1879,7 +1879,7 @@ for dim ∈ one_dimension_combinations_no_t
                  if input === nothing
                      if run_info.dfns
                          if var_name ∈ keys(input_dict_dfns)
-                             input = input_dict[var_name]
+                             input = input_dict_dfns[var_name]
                          else
                              input = input_dict_dfns
                          end
@@ -2109,7 +2109,7 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
                  if input === nothing
                      if run_info.dfns
                          if var_name ∈ keys(input_dict_dfns)
-                             input = input_dict[var_name]
+                             input = input_dict_dfns[var_name]
                          else
                              input = input_dict_dfns
                          end

From 4089b824231dd990b0dae226bf49479fca0f1eba Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 17:44:30 +0000
Subject: [PATCH 167/428] Add select_slice() and select_slice_of_variable() for
 electron pdf

---
 .../src/makie_post_processing.jl              | 65 ++++++++++++++++++-
 moment_kinetics/src/load_data.jl              | 24 +++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index d4598fbf0..7e265a7ae 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -3193,10 +3193,71 @@ function select_slice(variable::AbstractArray{T,4}, dims::Symbol...; input=nothi
     return slice
 end
 
+function select_slice(variable::AbstractArray{T,5}, dims::Symbol...; input=nothing,
+                      it=nothing, is=1, ir=nothing, iz=nothing, ivperp=nothing,
+                      ivpa=nothing, kwargs...) where T
+    # Array is (vpa,vperp,z,r,t)
+
+    if it !== nothing
+        it0 = it
+    elseif input === nothing || :it0 ∉ input
+        it0 = size(variable, 5)
+    else
+        it0 = input.it0
+    end
+    if ir !== nothing
+        ir0 = ir
+    elseif input === nothing || :ir0 ∉ input
+        ir0 = max(size(variable, 4) ÷ 3, 1)
+    else
+        ir0 = input.ir0
+    end
+    if iz !== nothing
+        iz0 = iz
+    elseif input === nothing || :iz0 ∉ input
+        iz0 = max(size(variable, 3) ÷ 3, 1)
+    else
+        iz0 = input.iz0
+    end
+    if ivperp !== nothing
+        ivperp0 = ivperp
+    elseif input === nothing || :ivperp0 ∉ input
+        ivperp0 = max(size(variable, 2) ÷ 3, 1)
+    else
+        ivperp0 = input.ivperp0
+    end
+    if ivpa !== nothing
+        ivpa0 = ivpa
+    elseif input === nothing || :ivpa0 ∉ input
+        ivpa0 = max(size(variable, 1) ÷ 3, 1)
+    else
+        ivpa0 = input.ivpa0
+    end
+
+    slice = variable
+    if :t ∉ dims || it !== nothing
+        slice = selectdim(slice, 5, it0)
+    end
+    if :r ∉ dims || ir !== nothing
+        slice = selectdim(slice, 4, ir0)
+    end
+    if :z ∉ dims || iz !== nothing
+        slice = selectdim(slice, 3, iz0)
+    end
+    if :vperp ∉ dims || ivperp !== nothing
+        slice = selectdim(slice, 2, ivperp0)
+    end
+    if :vpa ∉ dims || ivpa !== nothing
+        slice = selectdim(slice, 1, ivpa0)
+    end
+
+    return slice
+end
+
 function select_slice(variable::AbstractArray{T,6}, dims::Symbol...; input=nothing,
                       it=nothing, is=1, ir=nothing, iz=nothing, ivperp=nothing,
                       ivpa=nothing, kwargs...) where T
-    # Array is (z,r,species,t)
+    # Array is (vpa,vperp,z,r,species,t)
 
     if it !== nothing
         it0 = it
@@ -3258,7 +3319,7 @@ end
 function select_slice(variable::AbstractArray{T,7}, dims::Symbol...; input=nothing,
                       it=nothing, is=1, ir=nothing, iz=nothing, ivzeta=nothing,
                       ivr=nothing, ivz=nothing, kwargs...) where T
-    # Array is (z,r,species,t)
+    # Array is (vz,vr,vzeta,z,r,species,t)
 
     if it !== nothing
         it0 = it
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 6aa3d00d5..635389461 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3812,6 +3812,30 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         return variable
     end
 
+    # Select a slice of an electron distribution function sized variable
+    function select_slice_of_variable(variable::AbstractArray{T,5} where T; it=nothing,
+                                      is=nothing, ir=nothing, iz=nothing, ivperp=nothing,
+                                      ivpa=nothing, ivzeta=nothing, ivr=nothing,
+                                      ivz=nothing)
+        if it !== nothing
+            variable = selectdim(variable, 5, kwargs[:it])
+        end
+        if ir !== nothing
+            variable = selectdim(variable, 4, kwargs[:ir])
+        end
+        if iz !== nothing
+            variable = selectdim(variable, 3, kwargs[:iz])
+        end
+        if ivperp !== nothing
+            variable = selectdim(variable, 2, kwargs[:ivperp])
+        end
+        if ivpa !== nothing
+            variable = selectdim(variable, 1, kwargs[:ivpa])
+        end
+
+        return variable
+    end
+
     # Select a slice of a neutral distribution function sized variable
     function select_slice_of_variable(variable::AbstractArray{T,7} where T; it=nothing,
                                       is=nothing, ir=nothing, iz=nothing, ivperp=nothing,

From cbe3e644f5e7ba3aa540f8351920e9580f576e52 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 13:09:13 +0000
Subject: [PATCH 168/428] Allow passing *.initial_electron.* output files to
 `get_run_info()`

---
 moment_kinetics/src/load_data.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 635389461..6f731d2fe 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3190,8 +3190,10 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
             run_name = split(filename, ".moments.")[1]
         elseif occursin(".dfns.", filename)
             run_name = split(filename, ".dfns.")[1]
+        elseif occursin(".initial_electron.", filename)
+            run_name = split(filename, ".initial_electron.")[1]
         else
-            error("Cannot recognise '$this_run_dir' as a moment_kinetics output file")
+            error("Cannot recognise '$this_run_dir/$filename' as a moment_kinetics output file")
         end
     elseif isdir(this_run_dir)
         # Normalise by removing any trailing slash - with a slash basename() would return an

From 2422cc4412085be01ffc1edf7f7e9708df9eba80 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 17:46:06 +0000
Subject: [PATCH 169/428] Timestep and CFL diagnostics for electrons

---
 .../src/makie_post_processing.jl              | 230 ++++++++++++------
 moment_kinetics/src/load_data.jl              | 181 +++++++++++++-
 moment_kinetics/src/utils.jl                  |  11 +
 3 files changed, 344 insertions(+), 78 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 7e265a7ae..b32058d76 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -6918,7 +6918,7 @@ will be saved with the format `plot_prefix_timestep_diagnostics.pdf`.
 
 `it` can be used to select a subset of the time points by passing a range.
 """
-function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
+function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electron=false)
     try
         if !isa(run_info, Tuple)
             run_info = (run_info,)
@@ -6928,9 +6928,15 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
 
         input = Dict_to_NamedTuple(input_dict["timestep_diagnostics"])
 
-         steps_fig = nothing
-         dt_fig = nothing
-         CFL_fig = nothing
+        steps_fig = nothing
+        dt_fig = nothing
+        CFL_fig = nothing
+
+        if electron
+            electron_prefix = "electron_"
+        else
+            electron_prefix = ""
+        end
 
         if input.plot
             # Plot numbers of steps and numbers of failures
@@ -6951,67 +6957,86 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
                     prefix = ri.run_name * " "
                 end
 
-                plot_1d(ri.time, get_variable(ri, "steps_per_output"; it=it);
-                        label=prefix * "steps", ax=ax)
+                if it !== nothing
+                    time = ri.time[it]
+                else
+                    time = ri.time
+                end
+                plot_1d(time, get_variable(ri, "$(electron_prefix)steps_per_output";
+                                              it=it); label=prefix * "steps", ax=ax)
                 # Fudge to create an invisible line on ax_failures that cycles the line colors
                 # and adds a label for "steps_per_output" to the plot because we create the
                 # legend from ax_failures.
                 plot_1d([ri.time[1]], [0]; label=prefix * "steps", ax=ax_failures)
-                plot_1d(ri.time, get_variable(ri, "failures_per_output"; it=it);
+                plot_1d(time,
+                        get_variable(ri, "$(electron_prefix)failures_per_output"; it=it);
                         label=prefix * "failures", ax=ax_failures)
 
-                failure_caused_by_per_output = get_variable(ri,
-                                                            "failure_caused_by_per_output";
-                                                            it=it)
+                failure_caused_by_per_output =
+                    get_variable(ri, "$(electron_prefix)failure_caused_by_per_output";
+                                 it=it)
                 counter = 0
-                # Ion pdf failure counter
+                # pdf failure counter
                 counter += 1
-                plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
-                        label=prefix * "failures caused by f_ion", ax=ax_failures)
-                if ri.evolve_density
+                if electron
+                    label = prefix * "failures caused by f_electron"
+                else
+                    label = prefix * "failures caused by f_ion"
+                end
+                plot_1d(time, @view failure_caused_by_per_output[counter,:];
+                        label=label, ax=ax_failures)
+                if !electron && ri.evolve_density
                     # Ion density failure counter
                     counter += 1
-                    plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                    plot_1d(time, @view failure_caused_by_per_output[counter,:];
                             linestyle=:dash, label=prefix * "failures caused by n_ion",
                             ax=ax_failures)
                 end
-                if ri.evolve_upar
+                if !electron && ri.evolve_upar
                     # Ion flow failure counter
                     counter += 1
-                    plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                    plot_1d(time, @view failure_caused_by_per_output[counter,:];
                             linestyle=:dash, label=prefix * "failures caused by u_ion",
                             ax=ax_failures)
                 end
-                if ri.evolve_ppar
-                    # Ion flow failure counter
+                if !electron && ri.evolve_ppar
+                    # Ion parallel pressure failure counter
                     counter += 1
-                    plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                    plot_1d(time, @view failure_caused_by_per_output[counter,:];
                             linestyle=:dash, label=prefix * "failures caused by p_ion",
                             ax=ax_failures)
                 end
-                if ri.n_neutral_species > 0
+                if electron || ri.composition.electron_physics ∈ (braginskii_fluid,
+                                                                  kinetic_electrons)
+                    # Electron parallel pressure failure counter
+                    counter += 1
+                    plot_1d(time, @view failure_caused_by_per_output[counter,:];
+                            linestyle=:dash, label=prefix * "failures caused by p_electron",
+                            ax=ax_failures)
+                end
+                if !electron && ri.n_neutral_species > 0
                     # Neutral pdf failure counter
                     counter += 1
-                    plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                    plot_1d(time, @view failure_caused_by_per_output[counter,:];
                             label=prefix * "failures caused by f_neutral", ax=ax_failures)
                     if ri.evolve_density
                         # Neutral density failure counter
                         counter += 1
-                        plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                        plot_1d(time, @view failure_caused_by_per_output[counter,:];
                                 linestyle=:dash,
                                 label=prefix * "failures caused by n_neutral", ax=ax_failures)
                     end
                     if ri.evolve_upar
                         # Neutral flow failure counter
                         counter += 1
-                        plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                        plot_1d(time, @view failure_caused_by_per_output[counter,:];
                                 linestyle=:dash,
                                 label=prefix * "failures caused by u_neutral", ax=ax_failures)
                     end
                     if ri.evolve_ppar
                         # Neutral flow failure counter
                         counter += 1
-                        plot_1d(ri.time, @view failure_caused_by_per_output[counter,:];
+                        plot_1d(time, @view failure_caused_by_per_output[counter,:];
                                 linestyle=:dash,
                                 label=prefix * "failures caused by p_neutral", ax=ax_failures)
                     end
@@ -7034,11 +7059,11 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
             ########################
 
             if plot_prefix !== nothing
-                outfile = plot_prefix * "successful_dt.pdf"
+                outfile = plot_prefix * "$(electron_prefix)successful_dt.pdf"
             else
                 outfile = nothing
             end
-            dt_fig = plot_vs_t(run_info, "average_successful_dt"; outfile=outfile)
+            dt_fig = plot_vs_t(run_info, "$(electron_prefix)average_successful_dt"; outfile=outfile)
 
             # PLot minimum CFL factors
             ##########################
@@ -7051,14 +7076,23 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
                 else
                     prefix = ri.run_name * " "
                 end
-                CFL_vars = ["minimum_CFL_ion_z", "minimum_CFL_ion_vpa"]
-                if ri.n_neutral_species > 0
-                    push!(CFL_vars, "minimum_CFL_neutral_z", "minimum_CFL_neutral_vz")
+                if electron
+                    CFL_vars = ["minimum_CFL_electron_z", "minimum_CFL_electron_vpa"]
+                else
+                    CFL_vars = ["minimum_CFL_ion_z", "minimum_CFL_ion_vpa"]
+                    if ri.n_neutral_species > 0
+                        push!(CFL_vars, "minimum_CFL_neutral_z", "minimum_CFL_neutral_vz")
+                    end
+                end
+                if it !== nothing
+                    time = ri.time[it]
+                else
+                    time = ri.time
                 end
                 for varname ∈ CFL_vars
-                    var = get_variable(ri, varname)
+                    var = get_variable(ri, varname; it=it)
                     maxval = min(maxval, maximum(var))
-                    plot_1d(ri.time, var; ax=ax, label=prefix*varname)
+                    plot_1d(time, var; ax=ax, label=prefix*electron_prefix*varname)
                 end
             end
             ylims!(ax, 0.0, 4.0 * maxval)
@@ -7072,56 +7106,71 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
                 else
                     prefix = ri.run_name * " "
                 end
+                if it !== nothing
+                    time = ri.time[it]
+                else
+                    time = ri.time
+                end
 
-                limit_caused_by_per_output = get_variable(ri,
-                                                          "limit_caused_by_per_output";
-                                                          it=it)
+                limit_caused_by_per_output =
+                    get_variable(ri, "$(electron_prefix)limit_caused_by_per_output";
+                                 it=it)
                 counter = 0
 
                 # Accuracy limit counter
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                plot_1d(time, @view limit_caused_by_per_output[counter,:];
                         label=prefix * "RK accuracy", ax=ax)
 
                 # Maximum timestep increase limit counter
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                plot_1d(time, @view limit_caused_by_per_output[counter,:];
                         label=prefix * "max timestep increase", ax=ax)
 
                 # Slower maximum timestep increase near last failure limit counter
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                plot_1d(time, @view limit_caused_by_per_output[counter,:];
                         label=prefix * "max timestep increase near last fail", ax=ax)
 
                 # Minimum timestep limit counter
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                plot_1d(time, @view limit_caused_by_per_output[counter,:];
                         label=prefix * "min timestep", ax=ax)
 
                 # Maximum timestep limit counter
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                plot_1d(time, @view limit_caused_by_per_output[counter,:];
                         label=prefix * "max timestep", ax=ax)
 
-                # Ion z advection
+                # z advection
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
-                        label=prefix * "ion z advect", ax=ax)
+                if electron
+                    label = prefix * "electron z advect"
+                else
+                    label = prefix * "ion z advect"
+                end
+                plot_1d(time, @view limit_caused_by_per_output[counter,:]; label=label,
+                        ax=ax)
 
-                # Ion vpa advection
+                # vpa advection
                 counter += 1
-                plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
-                        label=prefix * "ion vpa advect", ax=ax)
+                if electron
+                    label = prefix * "electron vpa advect"
+                else
+                    label = prefix * "ion vpa advect"
+                end
+                plot_1d(time, @view limit_caused_by_per_output[counter,:]; label=label,
+                        ax=ax)
 
-                if ri.n_neutral_species > 0
-                    # Ion z advection
+                if !electron && ri.n_neutral_species > 0
+                    # Neutral z advection
                     counter += 1
-                    plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                    plot_1d(time, @view limit_caused_by_per_output[counter,:];
                             label=prefix * "neutral z advect", ax=ax)
 
-                    # Ion vpa advection
+                    # Neutral vz advection
                     counter += 1
-                    plot_1d(ri.time, @view limit_caused_by_per_output[counter,:];
+                    plot_1d(time, @view limit_caused_by_per_output[counter,:];
                             label=prefix * "neutral vz advect", ax=ax)
                 end
 
@@ -7140,13 +7189,13 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
 
 
             if plot_prefix !== nothing
-                outfile = plot_prefix * "timestep_diagnostics.pdf"
+                outfile = plot_prefix * electron_prefix * "timestep_diagnostics.pdf"
                 save(outfile, steps_fig)
 
-                outfile = plot_prefix * "CFL_factors.pdf"
+                outfile = plot_prefix * electron_prefix * "CFL_factors.pdf"
                 save(outfile, CFL_fig)
 
-                outfile = plot_prefix * "timestep_limits.pdf"
+                outfile = plot_prefix * electron_prefix * "timestep_limits.pdf"
                 save(outfile, limits_fig)
             else
                 display(steps_fig)
@@ -7160,29 +7209,56 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
             if plot_prefix === nothing
                 error("plot_prefix is required when animate_CFL=true")
             end
-            data = get_variable(run_info, "CFL_ion_z")
-            datamin = minimum(minimum(d) for d ∈ data)
-            animate_vs_vpa_z(run_info, "CFL_ion_z"; data=data,
-                             outfile=plot_prefix * "CFL_ion_z_vs_vpa_z.gif",
-                             colorscale=log10,
-                             transform=x->positive_or_nan(x; epsilon=1.e-30),
-                             colorrange=(datamin, datamin * 1000.0),
-                             axis_args=Dict(:bottomspinevisible=>false,
-                                            :topspinevisible=>false,
-                                            :leftspinevisible=>false,
-                                            :rightspinevisible=>false))
-            data = get_variable(run_info, "CFL_ion_vpa")
-            datamin = minimum(minimum(d) for d ∈ data)
-            animate_vs_vpa_z(run_info, "CFL_ion_vpa"; data=data,
-                             outfile=plot_prefix * "CFL_ion_vpa_vs_vpa_z.gif",
-                             colorscale=log10,
-                             transform=x->positive_or_nan(x; epsilon=1.e-30),
-                             colorrange=(datamin, datamin * 1000.0),
-                             axis_args=Dict(:bottomspinevisible=>false,
-                                            :topspinevisible=>false,
-                                            :leftspinevisible=>false,
-                                            :rightspinevisible=>false))
-            if any(ri.n_neutral_species > 0 for ri ∈ run_info)
+            if !electron
+                data = get_variable(run_info, "CFL_ion_z")
+                datamin = minimum(minimum(d) for d ∈ data)
+                animate_vs_vpa_z(run_info, "CFL_ion_z"; data=data,
+                                 outfile=plot_prefix * "CFL_ion_z_vs_vpa_z.gif",
+                                 colorscale=log10,
+                                 transform=x->positive_or_nan(x; epsilon=1.e-30),
+                                 colorrange=(datamin, datamin * 1000.0),
+                                 axis_args=Dict(:bottomspinevisible=>false,
+                                                :topspinevisible=>false,
+                                                :leftspinevisible=>false,
+                                                :rightspinevisible=>false))
+                data = get_variable(run_info, "CFL_ion_vpa")
+                datamin = minimum(minimum(d) for d ∈ data)
+                animate_vs_vpa_z(run_info, "CFL_ion_vpa"; data=data,
+                                 outfile=plot_prefix * "CFL_ion_vpa_vs_vpa_z.gif",
+                                 colorscale=log10,
+                                 transform=x->positive_or_nan(x; epsilon=1.e-30),
+                                 colorrange=(datamin, datamin * 1000.0),
+                                 axis_args=Dict(:bottomspinevisible=>false,
+                                                :topspinevisible=>false,
+                                                :leftspinevisible=>false,
+                                                :rightspinevisible=>false))
+            end
+            if electron || any(ri.composition.electron_physics == kinetic_electrons for ri
+                               ∈ run_info)
+                data = get_variable(run_info, "CFL_electron_z")
+                datamin = minimum(minimum(d) for d ∈ data)
+                animate_vs_vpa_z(run_info, "CFL_electron_z"; data=data,
+                                 outfile=plot_prefix * "CFL_electron_z_vs_vpa_z.gif",
+                                 colorscale=log10,
+                                 transform=x->positive_or_nan(x; epsilon=1.e-30),
+                                 colorrange=(datamin, datamin * 1000.0),
+                                 axis_args=Dict(:bottomspinevisible=>false,
+                                                :topspinevisible=>false,
+                                                :leftspinevisible=>false,
+                                                :rightspinevisible=>false))
+                data = get_variable(run_info, "CFL_electron_vpa")
+                datamin = minimum(minimum(d) for d ∈ data)
+                animate_vs_vpa_z(run_info, "CFL_electron_vpa"; data=data,
+                                 outfile=plot_prefix * "CFL_electron_vpa_vs_vpa_z.gif",
+                                 colorscale=log10,
+                                 transform=x->positive_or_nan(x; epsilon=1.e-30),
+                                 colorrange=(datamin, datamin * 1000.0),
+                                 axis_args=Dict(:bottomspinevisible=>false,
+                                                :topspinevisible=>false,
+                                                :leftspinevisible=>false,
+                                                :rightspinevisible=>false))
+            end
+            if !electron && any(ri.n_neutral_species > 0 for ri ∈ run_info)
                 data = get_variable(run_info, "CFL_neutral_z")
                 datamin = minimum(minimum(d) for d ∈ data)
                 animate_vs_vz_z(run_info, "CFL_neutral_z"; data=data,
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 6f731d2fe..623a12e55 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -17,6 +17,8 @@ using ..array_allocation: allocate_float, allocate_int
 using ..calculus: derivative!
 using ..communication: setup_distributed_memory_MPI
 using ..coordinates: coordinate, define_coordinate
+using ..electron_vpa_advection: update_electron_speed_vpa!
+using ..electron_z_advection: update_electron_speed_z!
 using ..file_io: check_io_implementation, get_group, get_subgroup_keys, get_variable_keys
 using ..input_structs: advection_input, grid_input, hdf5, netcdf
 using ..interpolation: interpolate_to_grid_1d!
@@ -39,7 +41,13 @@ const timestep_diagnostic_variables = ("time_for_run", "step_counter", "dt",
                                        "failure_counter", "failure_caused_by",
                                        "steps_per_output", "failures_per_output",
                                        "failure_caused_by_per_output",
-                                       "average_successful_dt")
+                                       "average_successful_dt", "electron_step_counter",
+                                       "electron_dt", "electron_failure_counter",
+                                       "electron_failure_caused_by",
+                                       "electron_steps_per_output",
+                                       "electron_failures_per_output",
+                                       "electron_failure_caused_by_per_output",
+                                       "electron_average_successful_dt")
 const em_variables = ("phi", "Er", "Ez")
 const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "thermal_speed", "temperature", "parallel_heat_flux",
@@ -4016,6 +4024,85 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
                                      run_info.time[it], run_info.geometry)
         end
 
+        variable = speed
+        variable = select_slice_of_variable(variable; kwargs...)
+     elseif variable_name == "electron_z_advect_speed"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        upar = get_variable(run_info, "electron_parallel_flow")
+        vth = get_variable(run_info, "electron_thermal_speed")
+        nz, nr, nt = size(upar)
+        nvperp = run_info.vperp.n
+        nvpa = run_info.vpa.n
+
+        speed = allocate_float(nz, nvpa, nvperp, nr, nt)
+
+        setup_distributed_memory_MPI(1,1,1,1)
+        setup_loop_ranges!(0, 1; s=run_info.n_ion_species, sn=run_info.n_neutral_species,
+                           r=nr, z=nz, vperp=nvperp, vpa=nvpa, vzeta=run_info.vzeta.n,
+                           vr=run_info.vr.n, vz=run_info.vz.n)
+        for it ∈ 1:nt
+            begin_serial_region()
+            # Only need some struct with a 'speed' variable
+            advect = (speed=@view(speed[:,:,:,:,it]),)
+            @views update_electron_speed_z!(advect, upar[:,:,it], vth[:,:,it],
+                                            run_info.vpa.grid)
+        end
+
+        # Horrible hack so that we can get the speed back without rearranging the
+        # dimensions, if we want that to pass it to a utility function from the main code
+        # (e.g. to calculate a CFL limit).
+        if normalize_advection_speed_shape
+            variable = allocate_float(nvpa, nvperp, nz, nr, nspecies, nt)
+            for it ∈ 1:nt, ir ∈ 1:nr, iz ∈ 1:nz, ivperp ∈ 1:nvperp, ivpa ∈ 1:nvpa
+                variable[ivpa,ivperp,iz,ir,it] = speed[iz,ivpa,ivperp,ir,it]
+            end
+            variable = select_slice_of_variable(variable; kwargs...)
+        else
+            variable = speed
+            if :it ∈ keys(kwargs)
+                variable = selectdim(variable, 5, kwargs[:it])
+            end
+            if :ir ∈ keys(kwargs)
+                variable = selectdim(variable, 4, kwargs[:ir])
+            end
+            if :ivperp ∈ keys(kwargs)
+                variable = selectdim(variable, 3, kwargs[:ivperp])
+            end
+            if :ivpa ∈ keys(kwargs)
+                variable = selectdim(variable, 2, kwargs[:ivpa])
+            end
+            if :iz ∈ keys(kwargs)
+                variable = selectdim(variable, 1, kwargs[:iz])
+            end
+        end
+    elseif variable_name == "electron_vpa_advect_speed"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        ppar = get_variable(run_info, "electron_parallel_pressure")
+        vth = get_variable(run_info, "electron_thermal_speed")
+        dppar_dz = get_z_derivative(run_info, "electron_parallel_pressure")
+        dvth_dz = get_z_derivative(run_info, "electron_thermal_speed")
+        dqpar_dz = get_z_derivative(run_info, "electron_parallel_heat_flux")
+
+        nz, nr, nt = size(vth)
+        nvperp = run_info.vperp.n
+        nvpa = run_info.vpa.n
+
+        speed=allocate_float(nvpa, nvperp, nz, nr, nt)
+        setup_distributed_memory_MPI(1,1,1,1)
+        setup_loop_ranges!(0, 1; s=run_info.n_ion_species, sn=run_info.n_neutral_species,
+                           r=nr, z=nz, vperp=nvperp, vpa=nvpa, vzeta=run_info.vzeta.n,
+                           vr=run_info.vr.n, vz=run_info.vz.n)
+        for it ∈ 1:nt
+            begin_serial_region()
+            # Only need some struct with a 'speed' variable
+            advect = (speed=@view(speed[:,:,:,:,it]),)
+            @views update_electron_speed_vpa!(advect, ppar[:,:,it], vth[:,:,it],
+                                              dppar_dz[:,:,it], dqpar_dz[:,:,it],
+                                              dvth_dz[:,:,it], run_info.vpa.grid)
+        end
+
         variable = speed
         variable = select_slice_of_variable(variable; kwargs...)
     elseif variable_name == "neutral_z_advect_speed"
@@ -4175,6 +4262,45 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             # Don't want a meaningless Inf...
             variable[1] = 0.0
         end
+    elseif variable_name == "electron_steps_per_output"
+        electron_steps_per_output = get_variable(run_info, "electron_step_counter"; kwargs...)
+        for i ∈ length(electron_steps_per_output):-1:2
+            electron_steps_per_output[i] -= electron_steps_per_output[i-1]
+        end
+        variable = electron_steps_per_output
+    elseif variable_name == "electron_failures_per_output"
+        electron_failures_per_output = get_variable(run_info, "electron_failure_counter"; kwargs...)
+        for i ∈ length(electron_failures_per_output):-1:2
+            electron_failures_per_output[i] -= electron_failures_per_output[i-1]
+        end
+        variable = electron_failures_per_output
+    elseif variable_name == "electron_failure_caused_by_per_output"
+        electron_failure_caused_by_per_output = get_variable(run_info, "electron_failure_caused_by"; kwargs...)
+        for i ∈ size(electron_failure_caused_by_per_output,2):-1:2
+            electron_failure_caused_by_per_output[:,i] .-= electron_failure_caused_by_per_output[:,i-1]
+        end
+        variable = electron_failure_caused_by_per_output
+    elseif variable_name == "electron_limit_caused_by_per_output"
+        electron_limit_caused_by_per_output = get_variable(run_info, "electron_limit_caused_by"; kwargs...)
+        for i ∈ size(electron_limit_caused_by_per_output,2):-1:2
+            electron_limit_caused_by_per_output[:,i] .-= electron_limit_caused_by_per_output[:,i-1]
+        end
+        variable = electron_limit_caused_by_per_output
+    elseif variable_name == "electron_average_successful_dt"
+        electron_steps_per_output = get_variable(run_info, "electron_steps_per_output"; kwargs...)
+        electron_failures_per_output = get_variable(run_info, "electron_failures_per_output"; kwargs...)
+        electron_successful_steps_per_output = electron_steps_per_output - electron_failures_per_output
+
+        delta_t = copy(run_info.time)
+        for i ∈ length(delta_t):-1:2
+            delta_t[i] -= delta_t[i-1]
+        end
+
+        variable = delta_t ./ electron_successful_steps_per_output
+        if electron_successful_steps_per_output[1] == 0
+            # Don't want a meaningless Inf...
+            variable[1] = 0.0
+        end
     elseif variable_name == "CFL_ion_z"
         # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
         # to get_variable() in this case. Instead select a slice of the result.
@@ -4201,6 +4327,34 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             @views get_CFL!(CFL[:,:,:,:,:,it], speed[:,:,:,:,:,it], run_info.vpa)
         end
 
+        variable = CFL
+        variable = select_slice_of_variable(variable; kwargs...)
+    elseif variable_name == "CFL_electron_z"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        speed = get_variable(run_info, "electron_z_advect_speed";
+                             normalize_advection_speed_shape=false)
+        nz, nvpa, nvperp, nr, nt = size(speed)
+        CFL = similar(speed)
+        for it ∈ 1:nt
+            @views get_CFL!(CFL[:,:,:,:,it], speed[:,:,:,:,it], run_info.z)
+        end
+
+        variable = allocate_float(nvpa, nvperp, nz, nr, nt)
+        for it ∈ 1:nt, ir ∈ 1:nr, iz ∈ 1:nz, ivperp ∈ 1:nvperp, ivpa ∈ 1:nvpa
+            variable[ivpa,ivperp,iz,ir,it] = CFL[iz,ivpa,ivperp,ir,it]
+        end
+        variable = select_slice_of_variable(variable; kwargs...)
+    elseif variable_name == "CFL_electron_vpa"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        speed = get_variable(run_info, "electron_vpa_advect_speed")
+        nt = size(speed, 5)
+        CFL = similar(speed)
+        for it ∈ 1:nt
+            @views get_CFL!(CFL[:,:,:,:,it], speed[:,:,:,:,it], run_info.vpa)
+        end
+
         variable = CFL
         variable = select_slice_of_variable(variable; kwargs...)
     elseif variable_name == "CFL_neutral_z"
@@ -4264,6 +4418,31 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             variable[it] = min_CFL
         end
         variable = select_slice_of_variable(variable; kwargs...)
+    elseif variable_name == "minimum_CFL_electron_z"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        speed = get_variable(run_info, "electron_z_advect_speed";
+                             normalize_advection_speed_shape=false)
+        nt = size(speed, 5)
+        variable = allocate_float(nt)
+        begin_serial_region()
+        for it ∈ 1:nt
+            min_CFL = get_minimum_CFL_z(@view(speed[:,:,:,:,it]), run_info.z)
+            variable[it] = min_CFL
+        end
+        variable = select_slice_of_variable(variable; kwargs...)
+    elseif variable_name == "minimum_CFL_electron_vpa"
+        # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
+        # to get_variable() in this case. Instead select a slice of the result.
+        speed = get_variable(run_info, "electron_vpa_advect_speed")
+        nt = size(speed, 5)
+        variable = allocate_float(nt)
+        begin_serial_region()
+        for it ∈ 1:nt
+            min_CFL = get_minimum_CFL_vpa(@view(speed[:,:,:,:,it]), run_info.vpa)
+            variable[it] = min_CFL
+        end
+        variable = select_slice_of_variable(variable; kwargs...)
     elseif variable_name == "minimum_CFL_neutral_z"
         # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
         # to get_variable() in this case. Instead select a slice of the result.
diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl
index 2c98983ff..052f71687 100644
--- a/moment_kinetics/src/utils.jl
+++ b/moment_kinetics/src/utils.jl
@@ -404,6 +404,17 @@ post-processing.
 """
 function get_CFL end
 
+function get_CFL!(CFL::AbstractArray{T,4}, speed::AbstractArray{T,4}, coord) where T
+
+    nmain, n2, n3, n4 = size(speed)
+
+    for i4 ∈ 1:n4, i3 ∈ 1:n3, i2 ∈ 1:n2, imain ∈ 1:nmain
+        CFL[imain,i2,i3,i4] = abs(coord.cell_width[imain] / speed[imain,i2,i3,i4])
+    end
+
+    return CFL
+end
+
 function get_CFL!(CFL::AbstractArray{T,5}, speed::AbstractArray{T,5}, coord) where T
 
     nmain, n2, n3, n4, n5 = size(speed)

From aaf8e8342708d08c1a9ebae65ebaa6526c8a2b8b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 14:08:21 +0000
Subject: [PATCH 170/428] Support `it` argument for CFL animations in
 `timestep_diagnostics()`

---
 .../src/makie_post_processing.jl                     | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index b32058d76..01aaf558d 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -7212,7 +7212,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
             if !electron
                 data = get_variable(run_info, "CFL_ion_z")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vpa_z(run_info, "CFL_ion_z"; data=data,
+                animate_vs_vpa_z(run_info, "CFL_ion_z"; data=data, it=it,
                                  outfile=plot_prefix * "CFL_ion_z_vs_vpa_z.gif",
                                  colorscale=log10,
                                  transform=x->positive_or_nan(x; epsilon=1.e-30),
@@ -7223,7 +7223,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                                                 :rightspinevisible=>false))
                 data = get_variable(run_info, "CFL_ion_vpa")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vpa_z(run_info, "CFL_ion_vpa"; data=data,
+                animate_vs_vpa_z(run_info, "CFL_ion_vpa"; data=data, it=it,
                                  outfile=plot_prefix * "CFL_ion_vpa_vs_vpa_z.gif",
                                  colorscale=log10,
                                  transform=x->positive_or_nan(x; epsilon=1.e-30),
@@ -7237,7 +7237,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                                ∈ run_info)
                 data = get_variable(run_info, "CFL_electron_z")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vpa_z(run_info, "CFL_electron_z"; data=data,
+                animate_vs_vpa_z(run_info, "CFL_electron_z"; data=data, it=it,
                                  outfile=plot_prefix * "CFL_electron_z_vs_vpa_z.gif",
                                  colorscale=log10,
                                  transform=x->positive_or_nan(x; epsilon=1.e-30),
@@ -7248,7 +7248,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                                                 :rightspinevisible=>false))
                 data = get_variable(run_info, "CFL_electron_vpa")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vpa_z(run_info, "CFL_electron_vpa"; data=data,
+                animate_vs_vpa_z(run_info, "CFL_electron_vpa"; data=data, it=it,
                                  outfile=plot_prefix * "CFL_electron_vpa_vs_vpa_z.gif",
                                  colorscale=log10,
                                  transform=x->positive_or_nan(x; epsilon=1.e-30),
@@ -7261,7 +7261,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
             if !electron && any(ri.n_neutral_species > 0 for ri ∈ run_info)
                 data = get_variable(run_info, "CFL_neutral_z")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vz_z(run_info, "CFL_neutral_z"; data=data,
+                animate_vs_vz_z(run_info, "CFL_neutral_z"; data=data, it=it,
                                 outfile=plot_prefix * "CFL_neutral_z_vs_vz_z.gif",
                                 colorscale=log10,
                                 transform=x->positive_or_nan(x; epsilon=1.e-30),
@@ -7272,7 +7272,7 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                                                :rightspinevisible=>false))
                 data = get_variable(run_info, "CFL_neutral_vz")
                 datamin = minimum(minimum(d) for d ∈ data)
-                animate_vs_vz_z(run_info, "CFL_neutral_vz"; data=data,
+                animate_vs_vz_z(run_info, "CFL_neutral_vz"; data=data, it=it,
                                 outfile=plot_prefix * "CFL_neutral_vz_vs_vz_z.gif",
                                 colorscale=log10,
                                 transform=x->positive_or_nan(x; epsilon=1.e-30),

From 87b3f52081e7650972a0fa5725042adc836c0104 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 18:47:40 +0000
Subject: [PATCH 171/428] Fix "Make vz resolution consistent with vpa in
 wall+sheath-bc_kinetic_loworder.toml"

---
 .../kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index bc30e359d..49819db60 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -73,8 +73,8 @@ vpa_L = 12.0 #8.0
 vpa_bc = "zero"
 vpa_discretization = "chebyshev_pseudospectral"
 #vpa_discretization = "gausslegendre_pseudospectral"
-vz_ngrid = 80
-vz_nelement = 5
+vz_ngrid = 5
+vz_nelement = 80
 vz_L = 12.0
 vz_bc = "zero"
 vz_discretization = "chebyshev_pseudospectral"

From c9a6aacc6ef19d00740110b59d45c88ca8a6ecb1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 18:48:19 +0000
Subject: [PATCH 172/428] Fix typos in `neutral_1V` in `load_data`

---
 moment_kinetics/src/load_data.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 623a12e55..55d18bbda 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -608,8 +608,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
 
             neutral_1V = (vzeta.n_global == 1 && vr.n_global == 1)
             restart_neutral_1V = (restart_vzeta.n_global == 1 && restart_vr.n_global == 1)
-            if any(geometry.bzeta .!= 0.0) && ((neutral1V && !restart_neutral_1V) ||
-                                               (!neutral1V && restart_neutral_1V))
+            if any(geometry.bzeta .!= 0.0) && ((neutral_1V && !restart_neutral_1V) ||
+                                               (!neutral_1V && restart_neutral_1V))
                 # One but not the other of the run being restarted from and this run are
                 # 1V, but the interpolation below does not allow for vz and vpa being in
                 # different directions. Therefore interpolation between 1V and 3V cases
@@ -904,8 +904,8 @@ function reload_electron_data!(pdf, moments, restart_prefix_iblock, time_index,
 
             neutral_1V = (vzeta.n_global == 1 && vr.n_global == 1)
             restart_neutral_1V = (restart_vzeta.n_global == 1 && restart_vr.n_global == 1)
-            if geometry.bzeta != 0.0 && ((neutral1V && !restart_neutral_1V) ||
-                                         (!neutral1V && restart_neutral_1V))
+            if geometry.bzeta != 0.0 && ((neutral_1V && !restart_neutral_1V) ||
+                                         (!neutral_1V && restart_neutral_1V))
                 # One but not the other of the run being restarted from and this run are
                 # 1V, but the interpolation below does not allow for vz and vpa being in
                 # different directions. Therefore interpolation between 1V and 3V cases

From 3fe506849f655b55a8387ecea337057564a12d0d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 24 Mar 2024 13:50:18 +0000
Subject: [PATCH 173/428] Split update of t_params out of
 adaptive_timestep_update!()

...into a separate function `adaptive_timestep_update_t_params!()` that
will be able to be called from the electron version of
`adaptive_timestep_update!()` as well, to reduce code duplication.
`adaptive_timestep_update_t_params!()` is defined in the `runge_kutta`
module to allow it to be used both in time_advance and
electron_kinetic_equation modules.
---
 moment_kinetics/src/runge_kutta.jl  | 207 ++++++++++++++++++++++++++++
 moment_kinetics/src/time_advance.jl | 199 +-------------------------
 2 files changed, 210 insertions(+), 196 deletions(-)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index ed20a2b6c..c311bb30b 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -11,6 +11,9 @@ using ..array_allocation: allocate_float
 using ..looping
 using ..type_definitions: mk_float
 
+using MPI
+using StatsBase: mean
+
 """
 given the number of Runge Kutta stages that are requested,
 returns the needed Runge Kutta coefficients;
@@ -506,4 +509,208 @@ function rk_update_loop_neutrals!(rk_coefs,
     return nothing
 end
 
+"""
+    adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
+                                       total_points, current_dt, error_norm_method)
+
+Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_params`.
+"""
+function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
+                                            total_points, current_dt, error_norm_method)
+    # Get global minimum of CFL limits
+    CFL_limit = nothing
+    this_limit_caused_by = nothing
+    @serial_region begin
+        # Get maximum error over all blocks
+        CFL_limits = MPI.Allreduce(CFL_limits, min, comm_inter_block[])
+        CFL_limit_caused_by = argmin(CFL_limits)
+        CFL_limit = CFL_limits[CFL_limit_caused_by]
+        # Reserve first five entries of t_params.limit_caused_by for accuracy,
+        # max_increase_factor, max_increase_factor_near_fail, minimum_dt and maximum_dt
+        # limits.
+        this_limit_caused_by = CFL_limit_caused_by + 5
+    end
+
+    if error_norm_method == "Linf"
+        # Get overall maximum error on the shared-memory block
+        error_norms = MPI.Reduce(error_norms, max, comm_block[]; root=0)
+
+        error_norm = nothing
+        @serial_region begin
+            # Get maximum error over all blocks
+            error_norms = MPI.Allreduce(error_norms, max, comm_inter_block[])
+            error_norm = maximum(error_norms)
+        end
+        error_norm = MPI.bcast(error_norm, 0, comm_block[])
+    elseif error_norm_method == "L2"
+        # Get overall maximum error on the shared-memory block
+        error_norms = MPI.Reduce(error_norms, +, comm_block[]; root=0)
+
+        error_norm = nothing
+        @serial_region begin
+            # Get maximum error over all blocks
+            error_norms = MPI.Allreduce(error_norms, +, comm_inter_block[])
+
+            # So far `error_norms` is the sum of squares of the errors. Now that summation
+            # is finished, need to divide by total number of points and take square-root.
+            error_norms .= sqrt.(error_norms ./ total_points)
+            open("debug$(global_size[]).txt", "a") do io
+                for e in error_norms
+                    print(io, e, " ")
+                end
+                println(io, t_params.dt[], " ;")
+            end
+
+            # Weight the error from each variable equally by taking the mean, so the
+            # larger number of points in the distribution functions does not mean that
+            # error on the moments is ignored.
+            error_norm = mean(error_norms)
+        end
+
+        error_norm = MPI.bcast(error_norm, 0, comm_block[])
+    else
+        error("Unrecognized error_norm_method '$method'")
+    end
+
+    # Use current_dt instead of t_params.dt[] here because we are about to write to
+    # the shared-memory variable t_params.dt[] below, and we do not want to add an extra
+    # _block_synchronize() call after reading it here.
+    if error_norm > 1.0 && current_dt > t_params.minimum_dt
+        # Timestep failed, reduce timestep and re-try
+
+        # Set scratch[end] equal to scratch[1] to start the timestep over
+        scratch_temp = scratch[t_params.n_rk_stages+1]
+        scratch[t_params.n_rk_stages+1] = scratch[1]
+        scratch[1] = scratch_temp
+
+        @serial_region begin
+            t_params.failure_counter[] += 1
+
+            if t_params.previous_dt[] > 0.0
+                # If previous_dt=0, the previous step was also a failure so only update
+                # dt_before_last_fail when previous_dt>0
+                t_params.dt_before_last_fail[] = t_params.previous_dt[]
+            end
+
+            # If we were trying to take a step to the output timestep, dt will be smaller on
+            # the re-try, so will not reach the output time.
+            t_params.step_to_output[] = false
+
+            # Get new timestep estimate using same formula as for a successful step, but
+            # limit decrease to factor 1/2 - this factor should probably be settable!
+            t_params.dt[] = max(t_params.dt[] / 2.0,
+                                t_params.dt[] * t_params.step_update_prefactor * error_norm^(-1.0/t_params.rk_order))
+            t_params.dt[] = max(t_params.dt[], t_params.minimum_dt)
+
+            minimum_dt = 1.e-14
+            if t_params.dt[] < minimum_dt
+                println("Time advance failed: trying to set dt=$(t_params.dt[]) less than "
+                        * "$minimum_dt at t=$t. Ending run.")
+                # Set dt negative to signal an error
+                t_params.dt[] = -1.0
+            end
+
+            # Don't update the simulation time, as this step failed
+            t_params.previous_dt[] = 0.0
+
+            # Call the 'cause' of the timestep failure the variable that has the biggest
+            # error norm here
+            max_error_variable_index = argmax(error_norms)
+            t_params.failure_caused_by[max_error_variable_index] += 1
+
+            #println("t=$t, timestep failed, error_norm=$error_norm, error_norms=$error_norms, decreasing timestep to ", t_params.dt[])
+        end
+    else
+        @serial_region begin
+            # Save the timestep used to complete this step, this is used to update the
+            # simulation time.
+            t_params.previous_dt[] = t_params.dt[]
+
+            if t_params.step_to_output[]
+                # Completed an output step, reset dt to what it was before it was reduced to reach
+                # the output time
+                t_params.dt[] = t_params.dt_before_output[]
+                t_params.step_to_output[] = false
+
+                if t_params.dt[] > CFL_limit
+                    t_params.dt[] = CFL_limit
+                end
+            else
+                # Adjust timestep according to Fehlberg's suggestion
+                # (https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta%E2%80%93Fehlberg_method).
+                # `step_update_prefactor` is a constant numerical factor to make the estimate
+                # of a good value for the next timestep slightly conservative. It defaults to
+                # 0.9.
+                t_params.dt[] *= t_params.step_update_prefactor * error_norm^(-1.0/t_params.rk_order)
+
+                if t_params.dt[] > CFL_limit
+                    t_params.dt[] = CFL_limit
+                else
+                    this_limit_caused_by = 1
+                end
+
+                # Limit so timestep cannot increase by a large factor, which might lead to
+                # numerical instability in some cases.
+                max_cap_limit_caused_by = 2
+                if isinf(t_params.max_increase_factor_near_last_fail)
+                    # Not using special timestep limiting near last failed dt value
+                    max_cap = t_params.max_increase_factor * t_params.previous_dt[]
+                else
+                    max_cap = t_params.max_increase_factor * t_params.previous_dt[]
+                    slow_increase_threshold = t_params.dt_before_last_fail[] / t_params.last_fail_proximity_factor
+                    if t_params.previous_dt[] > t_params.dt_before_last_fail[] * t_params.last_fail_proximity_factor
+                        # dt has successfully exceeded the last failed value, so allow it
+                        # to increase more quickly again
+                        t_params.dt_before_last_fail[] = Inf
+                    elseif max_cap > slow_increase_threshold
+                        # dt is getting close to last failed value, so increase more
+                        # slowly
+                        max_cap = max(slow_increase_threshold,
+                                      t_params.max_increase_factor_near_last_fail *
+                                      t_params.previous_dt[])
+                        max_cap_limit_caused_by = 3
+                    end
+                end
+                if t_params.dt[] > max_cap
+                    t_params.dt[] = max_cap
+                    this_limit_caused_by = max_cap_limit_caused_by
+                end
+
+                # Prevent timestep from going below minimum_dt
+                if t_params.dt[] < t_params.minimum_dt
+                    t_params.dt[] = t_params.minimum_dt
+                    this_limit_caused_by = 4
+                end
+
+                # Prevent timestep from going above maximum_dt
+                if t_params.dt[] > t_params.maximum_dt
+                    t_params.dt[] = t_params.maximum_dt
+                    this_limit_caused_by = 5
+                end
+
+                t_params.limit_caused_by[this_limit_caused_by] += 1
+
+                if (t_params.step_counter[] % 1000 == 0) && global_rank[] == 0
+                    println("step ", t_params.step_counter[], ": t=",
+                            round(t, sigdigits=6), ", nfail=", t_params.failure_counter[],
+                            ", dt=", t_params.dt[])
+                end
+            end
+        end
+    end
+
+    @serial_region begin
+        if t + t_params.dt[] >= t_params.next_output_time[]
+            t_params.dt_before_output[] = t_params.dt[]
+            t_params.dt[] = t_params.next_output_time[] - t
+            t_params.step_to_output[] = true
+        end
+    end
+
+    # Shared-memory variables have been updated, so synchronize
+    _block_synchronize()
+
+    return nothing
+end
+
 end # runge_kutta
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index cef4f29a6..614128bfb 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -8,7 +8,6 @@ export allocate_advection_structs
 export setup_dummy_and_buffer_arrays
 
 using MPI
-using StatsBase: mean
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_shared_float, allocate_shared_bool
 using ..communication
@@ -61,7 +60,7 @@ using ..manufactured_solns: manufactured_sources
 using ..advection: advection_info
 using ..runge_kutta: rk_update_evolved_moments!, rk_update_evolved_moments_neutral!,
                      rk_update_variable!, rk_error_variable!,
-                     setup_runge_kutta_coefficients!
+                     setup_runge_kutta_coefficients!, adaptive_timestep_update_t_params!
 using ..utils: to_minutes, get_minimum_CFL_z, get_minimum_CFL_vpa,
                get_minimum_CFL_neutral_z, get_minimum_CFL_neutral_vz
 using ..electron_fluid_equations: calculate_electron_density!
@@ -2196,200 +2195,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
         end
     end
 
-    # Get global minimum of CFL limits
-    CFL_limit = nothing
-    this_limit_caused_by = nothing
-    @serial_region begin
-        # Get maximum error over all blocks
-        CFL_limits = MPI.Allreduce(CFL_limits, min, comm_inter_block[])
-        CFL_limit_caused_by = argmin(CFL_limits)
-        CFL_limit = CFL_limits[CFL_limit_caused_by]
-        # Reserve first two entries of t_params.limit_caused_by for accuracy limit and
-        # max_increase_factor limit.
-        this_limit_caused_by = CFL_limit_caused_by + 5
-    end
-
-    if error_norm_method == "Linf"
-        # Get overall maximum error on the shared-memory block
-        error_norms = MPI.Reduce(error_norms, max, comm_block[]; root=0)
-
-        error_norm = nothing
-        @serial_region begin
-            # Get maximum error over all blocks
-            error_norms = MPI.Allreduce(error_norms, max, comm_inter_block[])
-            error_norm = maximum(error_norms)
-        end
-        error_norm = MPI.bcast(error_norm, 0, comm_block[])
-    elseif error_norm_method == "L2"
-        # Get overall maximum error on the shared-memory block
-        error_norms = MPI.Reduce(error_norms, +, comm_block[]; root=0)
-
-        error_norm = nothing
-        @serial_region begin
-            # Get maximum error over all blocks
-            error_norms = MPI.Allreduce(error_norms, +, comm_inter_block[])
-
-            # So far `error_norms` is the sum of squares of the errors. Now that summation
-            # is finished, need to divide by total number of points and take square-root.
-            error_norms .= sqrt.(error_norms ./ total_points)
-            open("debug$(global_size[]).txt", "a") do io
-                for e in error_norms
-                    print(io, e, " ")
-                end
-                println(io, t_params.dt[], " ;")
-            end
-
-            # Weight the error from each variable equally by taking the mean, so the
-            # larger number of points in the distribution functions does not mean that
-            # error on the moments is ignored.
-            error_norm = mean(error_norms)
-        end
-
-        error_norm = MPI.bcast(error_norm, 0, comm_block[])
-    else
-        error("Unrecognized error_norm_method '$method'")
-    end
-
-    # Use current_dt instead of t_params.dt[] here because we are about to write to
-    # the shared-memory variable t_params.dt[] below, and we do not want to add an extra
-    # _block_synchronize() call after reading it here.
-    if error_norm > 1.0 && current_dt > t_params.minimum_dt
-        # Timestep failed, reduce timestep and re-try
-        success = false
-
-        # Set scratch[end] equal to scratch[1] to start the timestep over
-        scratch_temp = scratch[end]
-        scratch[end] = scratch[1]
-        scratch[1] = scratch_temp
-
-        @serial_region begin
-            t_params.failure_counter[] += 1
-
-            if t_params.previous_dt[] > 0.0
-                # If previous_dt=0, the previous step was also a failure so only update
-                # dt_before_last_fail when previous_dt>0
-                t_params.dt_before_last_fail[] = t_params.previous_dt[]
-            end
-
-            # If we were trying to take a step to the output timestep, dt will be smaller on
-            # the re-try, so will not reach the output time.
-            t_params.step_to_output[] = false
-
-            # Get new timestep estimate using same formula as for a successful step, but
-            # limit decrease to factor 1/2 - this factor should probably be settable!
-            t_params.dt[] = max(t_params.dt[] / 2.0,
-                                t_params.dt[] * t_params.step_update_prefactor * error_norm^(-1.0/t_params.rk_order))
-            t_params.dt[] = max(t_params.dt[], t_params.minimum_dt)
-
-            minimum_dt = 1.e-14
-            if t_params.dt[] < minimum_dt
-                println("Time advance failed: trying to set dt=$(t_params.dt[]) less than "
-                        * "$minimum_dt at t=$t. Ending run.")
-                # Set dt negative to signal an error
-                t_params.dt[] = -1.0
-            end
-
-            # Don't update the simulation time, as this step failed
-            t_params.previous_dt[] = 0.0
-
-            # Call the 'cause' of the timestep failure the variable that has the biggest
-            # error norm here
-            max_error_variable_index = argmax(error_norms)
-            t_params.failure_caused_by[max_error_variable_index] += 1
-
-            #println("t=$t, timestep failed, error_norm=$error_norm, error_norms=$error_norms, decreasing timestep to ", t_params.dt[])
-        end
-    else
-        success = true
-
-        @serial_region begin
-            # Save the timestep used to complete this step, this is used to update the
-            # simulation time.
-            t_params.previous_dt[] = t_params.dt[]
-
-            if t_params.step_to_output[]
-                # Completed an output step, reset dt to what it was before it was reduced to reach
-                # the output time
-                t_params.dt[] = t_params.dt_before_output[]
-                t_params.step_to_output[] = false
-
-                if t_params.dt[] > CFL_limit
-                    t_params.dt[] = CFL_limit
-                end
-            else
-                # Adjust timestep according to Fehlberg's suggestion
-                # (https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta%E2%80%93Fehlberg_method).
-                # `step_update_prefactor` is a constant numerical factor to make the estimate
-                # of a good value for the next timestep slightly conservative. It defaults to
-                # 0.9.
-                t_params.dt[] *= t_params.step_update_prefactor * error_norm^(-1.0/t_params.rk_order)
-
-                if t_params.dt[] > CFL_limit
-                    t_params.dt[] = CFL_limit
-                else
-                    this_limit_caused_by = 1
-                end
-
-                # Limit so timestep cannot increase by a large factor, which might lead to
-                # numerical instability in some cases.
-                max_cap_limit_caused_by = 2
-                if isinf(t_params.max_increase_factor_near_last_fail)
-                    # Not using special timestep limiting near last failed dt value
-                    max_cap = t_params.max_increase_factor * t_params.previous_dt[]
-                else
-                    max_cap = t_params.max_increase_factor * t_params.previous_dt[]
-                    slow_increase_threshold = t_params.dt_before_last_fail[] / t_params.last_fail_proximity_factor
-                    if t_params.previous_dt[] > t_params.dt_before_last_fail[] * t_params.last_fail_proximity_factor
-                        # dt has successfully exceeded the last failed value, so allow it
-                        # to increase more quickly again
-                        t_params.dt_before_last_fail[] = Inf
-                    elseif max_cap > slow_increase_threshold
-                        # dt is getting close to last failed value, so increase more
-                        # slowly
-                        max_cap = max(slow_increase_threshold,
-                                      t_params.max_increase_factor_near_last_fail *
-                                      t_params.previous_dt[])
-                        max_cap_limit_caused_by = 3
-                    end
-                end
-                if t_params.dt[] > max_cap
-                    t_params.dt[] = max_cap
-                    this_limit_caused_by = max_cap_limit_caused_by
-                end
-
-                # Prevent timestep from going below minimum_dt
-                if t_params.dt[] < t_params.minimum_dt
-                    t_params.dt[] = t_params.minimum_dt
-                    this_limit_caused_by = 4
-                end
-
-                # Prevent timestep from going above maximum_dt
-                if t_params.dt[] > t_params.maximum_dt
-                    t_params.dt[] = t_params.maximum_dt
-                    this_limit_caused_by = 5
-                end
-
-                t_params.limit_caused_by[this_limit_caused_by] += 1
-
-                if (t_params.step_counter[] % 1000 == 0) && global_rank[] == 0
-                    println("step ", t_params.step_counter[], ": t=",
-                            round(t, sigdigits=6), ", nfail=", t_params.failure_counter[],
-                            ", dt=", t_params.dt[])
-                end
-            end
-        end
-    end
-
-    @serial_region begin
-        if t + t_params.dt[] >= t_params.next_output_time[]
-            t_params.dt_before_output[] = t_params.dt[]
-            t_params.dt[] = t_params.next_output_time[] - t
-            t_params.step_to_output[] = true
-        end
-    end
-
-    # Shared-memory variables have been updated, so synchronize
-    _block_synchronize()
+    adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
+                                       total_points, current_dt, error_norm_method)
 
     return nothing
 end

From a01d47566e53de172b1869a1b384136d4e52c5dc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 08:49:52 +0000
Subject: [PATCH 174/428] Move local_error_norm() to runge_kutta module

Will allow it to be used by electron_kinetic_equation module as well as
time_advance.
---
 moment_kinetics/src/runge_kutta.jl  | 183 +++++++++++++++++++++++++++-
 moment_kinetics/src/time_advance.jl | 183 +---------------------------
 2 files changed, 184 insertions(+), 182 deletions(-)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index c311bb30b..7a55695fc 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -5,9 +5,10 @@ module runge_kutta
 
 export setup_runge_kutta_coefficients!, rk_update_evolved_moments!,
        rk_update_evolved_moments_electron!, rk_update_evolved_moments_neutral!,
-       rk_update_variable!, rk_error_variable!
+       rk_update_variable!, rk_error_variable!, local_error_norm
 
 using ..array_allocation: allocate_float
+using ..communication
 using ..looping
 using ..type_definitions: mk_float
 
@@ -509,6 +510,186 @@ function rk_update_loop_neutrals!(rk_coefs,
     return nothing
 end
 
+"""
+    local_error_norm(error, f, rtol, atol)
+    local_error_norm(error, f, rtol, atol, neutral=false; method="Linf",
+                     skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
+
+Maximum error norm in the range owned by this MPI process, given by
+```math
+\\max(\\frac{|\\mathtt{error}|}{\\mathtt{rtol}*|\\mathtt{f}| + \\mathtt{atol})
+```
+
+3 dimensional arrays (which represent moments) are treated as ion moments unless
+`neutral=true` is passed.
+
+`method` can be "Linf" (to take the maximum error) or "L2" to take the root-mean-square
+(RMS) error.
+
+`skip_r_inner` and `skip_z_lower` can be set to true to skip the contribution from the
+inner/lower boundaries, to avoid double-counting those points when using
+distributed-memory MPI.
+
+`error_sum_zero` should always have value 0.0, but is included so that different types can
+be used for L2sum. For testing, if we want consistency of results when using different
+numbers of processes (when the number of processes changes the order of operations in the
+sum is changed, which changes the rounding errors) then we have to use higher precision
+(i.e. use the Float128 type from the Quadmath package). The type of a 0.0 value can be set
+according to the `high_precision_error_sum` option in the `[timestepping]` section, and
+stored in a template-typed value in the `t_params` object - when that value is passed in
+as the argument to `error_sum_zero`, that type will be used for L2sum, and the type will
+be known at compile time, allowing this function to be efficient.
+"""
+function local_error_norm end
+
+function local_error_norm(error::MPISharedArray{mk_float,2},
+                          f::MPISharedArray{mk_float,2}, rtol, atol; method="Linf",
+                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
+    if method == "Linf"
+        f_max = -Inf
+        @loop_r_z ir iz begin
+            error_norm = abs(error[iz,ir]) / (rtol*abs(f[iz,ir]) + atol)
+            f_max = max(f_max, error_norm)
+        end
+        return f_max
+    elseif method == "L2"
+        L2sum = error_sum_zero
+        @loop_r_z ir iz begin
+            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                continue
+            end
+            error_norm = (error[iz,ir] / (rtol*abs(f[iz,ir]) + atol))^2
+            L2sum += error_norm
+        end
+        # Will sum results from different processes in shared memory block after returning
+        # from this function.
+        nz, nr = size(error)
+        if skip_r_inner
+            nr -= 1
+        end
+        if skip_z_lower
+            nz -= 1
+        end
+        return L2sum
+    else
+        error("Unrecognized method '$method'")
+    end
+end
+function local_error_norm(error::MPISharedArray{mk_float,3},
+                          f::MPISharedArray{mk_float,3}, rtol, atol, neutral=false;
+                          method="Linf", skip_r_inner=false, skip_z_lower=false,
+                          error_sum_zero=0.0)
+    if method == "Linf"
+        f_max = -Inf
+        if neutral
+            @loop_sn_r_z isn ir iz begin
+                error_norm = abs(error[iz,ir,isn]) / (rtol*abs(f[iz,ir,isn]) + atol)
+                f_max = max(f_max, error_norm)
+            end
+        else
+            @loop_s_r_z is ir iz begin
+                error_norm = abs(error[iz,ir,is]) / (rtol*abs(f[iz,ir,is]) + atol)
+                f_max = max(f_max, error_norm)
+            end
+        end
+        return f_max
+    elseif method == "L2"
+        L2sum = error_sum_zero
+        if neutral
+            @loop_sn_r_z isn ir iz begin
+                if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                    continue
+                end
+                error_norm = (error[iz,ir,isn] / (rtol*abs(f[iz,ir,isn]) + atol))^2
+                L2sum += error_norm
+            end
+        else
+            @loop_s_r_z is ir iz begin
+                if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                    continue
+                end
+                error_norm = (error[iz,ir,is] / (rtol*abs(f[iz,ir,is]) + atol))^2
+                L2sum += error_norm
+            end
+        end
+        # Will sum results from different processes in shared memory block after returning
+        # from this function.
+        nz, nr, nspecies = size(error)
+        if skip_r_inner
+            nr -= 1
+        end
+        if skip_z_lower
+            nz -= 1
+        end
+        return L2sum
+    else
+        error("Unrecognized method '$method'")
+    end
+end
+function local_error_norm(error::MPISharedArray{mk_float,5},
+                          f::MPISharedArray{mk_float,5}, rtol, atol; method="Linf",
+                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
+    if method == "Linf"
+        f_max = -Inf
+        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+            error_norm = abs(error[ivpa,ivperp,iz,ir,is]) /
+                         (rtol*abs(f[ivpa,ivperp,iz,ir,is]) + atol)
+            f_max = max(f_max, error_norm)
+        end
+        return f_max
+    elseif method == "L2"
+        L2sum = error_sum_zero
+        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                continue
+            end
+            error_norm = (error[ivpa,ivperp,iz,ir,is] /
+                          (rtol*abs(f[ivpa,ivperp,iz,ir,is]) + atol))^2
+            L2sum += error_norm
+        end
+        # Will sum results from different processes in shared memory block after returning
+        # from this function.
+        nvpa, nvperp, nz, nr, nspecies = size(error)
+        if skip_r_inner
+            nr -= 1
+        end
+        if skip_z_lower
+            nz -= 1
+        end
+        return L2sum
+    else
+        error("Unrecognized method '$method'")
+    end
+end
+function local_error_norm(error::MPISharedArray{mk_float,6},
+                          f::MPISharedArray{mk_float,6}, rtol, atol; method="Linf",
+                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
+    if method == "Linf"
+        f_max = -Inf
+        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
+            error_norm = abs(error[ivz,ivr,ivzeta,iz,ir,isn]) /
+                         (rtol*abs(f[ivz,ivr,ivzeta,iz,ir,isn]) + atol)
+            f_max = max(f_max, error_norm)
+        end
+        return f_max
+    elseif method == "L2"
+        L2sum = error_sum_zero
+        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
+            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                continue
+            end
+            error_norm = (error[ivz,ivr,ivzeta,iz,ir,isn] /
+                          (rtol*abs(f[ivz,ivr,ivzeta,iz,ir,isn]) + atol))^2
+            L2sum += error_norm
+        end
+        # Will sum results from different processes in shared memory block after returning
+        # from this function.
+        return L2sum
+    else
+        error("Unrecognized method '$method'")
+    end
+end
+
 """
     adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
                                        total_points, current_dt, error_norm_method)
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 614128bfb..aaf2e8437 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -60,7 +60,8 @@ using ..manufactured_solns: manufactured_sources
 using ..advection: advection_info
 using ..runge_kutta: rk_update_evolved_moments!, rk_update_evolved_moments_neutral!,
                      rk_update_variable!, rk_error_variable!,
-                     setup_runge_kutta_coefficients!, adaptive_timestep_update_t_params!
+                     setup_runge_kutta_coefficients!, local_error_norm,
+                     adaptive_timestep_update_t_params!
 using ..utils: to_minutes, get_minimum_CFL_z, get_minimum_CFL_vpa,
                get_minimum_CFL_neutral_z, get_minimum_CFL_neutral_vz
 using ..electron_fluid_equations: calculate_electron_density!
@@ -1814,186 +1815,6 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     end
 end
 
-"""
-    local_error_norm(error, f, rtol, atol)
-    local_error_norm(error, f, rtol, atol, neutral=false; method="Linf",
-                     skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
-
-Maximum error norm in the range owned by this MPI process, given by
-```math
-\\max(\\frac{|\\mathtt{error}|}{\\mathtt{rtol}*|\\mathtt{f}| + \\mathtt{atol})
-```
-
-3 dimensional arrays (which represent moments) are treated as ion moments unless
-`neutral=true` is passed.
-
-`method` can be "Linf" (to take the maximum error) or "L2" to take the root-mean-square
-(RMS) error.
-
-`skip_r_inner` and `skip_z_lower` can be set to true to skip the contribution from the
-inner/lower boundaries, to avoid double-counting those points when using
-distributed-memory MPI.
-
-`error_sum_zero` should always have value 0.0, but is included so that different types can
-be used for L2sum. For testing, if we want consistency of results when using different
-numbers of processes (when the number of processes changes the order of operations in the
-sum is changed, which changes the rounding errors) then we have to use higher precision
-(i.e. use the Float128 type from the Quadmath package). The type of a 0.0 value can be set
-according to the `high_precision_error_sum` option in the `[timestepping]` section, and
-stored in a template-typed value in the `t_params` object - when that value is passed in
-as the argument to `error_sum_zero`, that type will be used for L2sum, and the type will
-be known at compile time, allowing this function to be efficient.
-"""
-function local_error_norm end
-
-function local_error_norm(error::MPISharedArray{mk_float,2},
-                          f::MPISharedArray{mk_float,2}, rtol, atol; method="Linf",
-                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
-    if method == "Linf"
-        f_max = -Inf
-        @loop_r_z ir iz begin
-            error_norm = abs(error[iz,ir]) / (rtol*abs(f[iz,ir]) + atol)
-            f_max = max(f_max, error_norm)
-        end
-        return f_max
-    elseif method == "L2"
-        L2sum = error_sum_zero
-        @loop_r_z ir iz begin
-            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
-                continue
-            end
-            error_norm = (error[iz,ir] / (rtol*abs(f[iz,ir]) + atol))^2
-            L2sum += error_norm
-        end
-        # Will sum results from different processes in shared memory block after returning
-        # from this function.
-        nz, nr = size(error)
-        if skip_r_inner
-            nr -= 1
-        end
-        if skip_z_lower
-            nz -= 1
-        end
-        return L2sum
-    else
-        error("Unrecognized method '$method'")
-    end
-end
-function local_error_norm(error::MPISharedArray{mk_float,3},
-                          f::MPISharedArray{mk_float,3}, rtol, atol, neutral=false;
-                          method="Linf", skip_r_inner=false, skip_z_lower=false,
-                          error_sum_zero=0.0)
-    if method == "Linf"
-        f_max = -Inf
-        if neutral
-            @loop_sn_r_z isn ir iz begin
-                error_norm = abs(error[iz,ir,isn]) / (rtol*abs(f[iz,ir,isn]) + atol)
-                f_max = max(f_max, error_norm)
-            end
-        else
-            @loop_s_r_z is ir iz begin
-                error_norm = abs(error[iz,ir,is]) / (rtol*abs(f[iz,ir,is]) + atol)
-                f_max = max(f_max, error_norm)
-            end
-        end
-        return f_max
-    elseif method == "L2"
-        L2sum = error_sum_zero
-        if neutral
-            @loop_sn_r_z isn ir iz begin
-                if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
-                    continue
-                end
-                error_norm = (error[iz,ir,isn] / (rtol*abs(f[iz,ir,isn]) + atol))^2
-                L2sum += error_norm
-            end
-        else
-            @loop_s_r_z is ir iz begin
-                if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
-                    continue
-                end
-                error_norm = (error[iz,ir,is] / (rtol*abs(f[iz,ir,is]) + atol))^2
-                L2sum += error_norm
-            end
-        end
-        # Will sum results from different processes in shared memory block after returning
-        # from this function.
-        nz, nr, nspecies = size(error)
-        if skip_r_inner
-            nr -= 1
-        end
-        if skip_z_lower
-            nz -= 1
-        end
-        return L2sum
-    else
-        error("Unrecognized method '$method'")
-    end
-end
-function local_error_norm(error::MPISharedArray{mk_float,5},
-                          f::MPISharedArray{mk_float,5}, rtol, atol; method="Linf",
-                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
-    if method == "Linf"
-        f_max = -Inf
-        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-            error_norm = abs(error[ivpa,ivperp,iz,ir,is]) /
-                         (rtol*abs(f[ivpa,ivperp,iz,ir,is]) + atol)
-            f_max = max(f_max, error_norm)
-        end
-        return f_max
-    elseif method == "L2"
-        L2sum = error_sum_zero
-        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
-                continue
-            end
-            error_norm = (error[ivpa,ivperp,iz,ir,is] /
-                          (rtol*abs(f[ivpa,ivperp,iz,ir,is]) + atol))^2
-            L2sum += error_norm
-        end
-        # Will sum results from different processes in shared memory block after returning
-        # from this function.
-        nvpa, nvperp, nz, nr, nspecies = size(error)
-        if skip_r_inner
-            nr -= 1
-        end
-        if skip_z_lower
-            nz -= 1
-        end
-        return L2sum
-    else
-        error("Unrecognized method '$method'")
-    end
-end
-function local_error_norm(error::MPISharedArray{mk_float,6},
-                          f::MPISharedArray{mk_float,6}, rtol, atol; method="Linf",
-                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
-    if method == "Linf"
-        f_max = -Inf
-        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
-            error_norm = abs(error[ivz,ivr,ivzeta,iz,ir,isn]) /
-                         (rtol*abs(f[ivz,ivr,ivzeta,iz,ir,isn]) + atol)
-            f_max = max(f_max, error_norm)
-        end
-        return f_max
-    elseif method == "L2"
-        L2sum = error_sum_zero
-        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
-            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
-                continue
-            end
-            error_norm = (error[ivz,ivr,ivzeta,iz,ir,isn] /
-                          (rtol*abs(f[ivz,ivr,ivzeta,iz,ir,isn]) + atol))^2
-            L2sum += error_norm
-        end
-        # Will sum results from different processes in shared memory block after returning
-        # from this function.
-        return L2sum
-    else
-        error("Unrecognized method '$method'")
-    end
-end
-
 """
     adaptive_timestep_update!(scratch, t_params, rk_coefs, moments, n_neutral_species)
 

From 5827a42573315db54efea09c2b757be6d851f2c3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 08:55:17 +0000
Subject: [PATCH 175/428] Define local_error_norm() for 4d array

...so it can be used for electron distribution function.
---
 moment_kinetics/src/runge_kutta.jl | 35 ++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 7a55695fc..5f7329da1 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -626,6 +626,41 @@ function local_error_norm(error::MPISharedArray{mk_float,3},
         error("Unrecognized method '$method'")
     end
 end
+function local_error_norm(error::MPISharedArray{mk_float,4},
+                          f::MPISharedArray{mk_float,4}, rtol, atol; method="Linf",
+                          skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)
+    if method == "Linf"
+        f_max = -Inf
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            error_norm = abs(error[ivpa,ivperp,iz,ir]) /
+                         (rtol*abs(f[ivpa,ivperp,iz,ir]) + atol)
+            f_max = max(f_max, error_norm)
+        end
+        return f_max
+    elseif method == "L2"
+        L2sum = error_sum_zero
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            if (skip_r_inner && ir == 1) || (skip_z_lower && iz == 1)
+                continue
+            end
+            error_norm = (error[ivpa,ivperp,iz,ir] /
+                          (rtol*abs(f[ivpa,ivperp,iz,ir]) + atol))^2
+            L2sum += error_norm
+        end
+        # Will sum results from different processes in shared memory block after returning
+        # from this function.
+        nvpa, nvperp, nz, nr = size(error)
+        if skip_r_inner
+            nr -= 1
+        end
+        if skip_z_lower
+            nz -= 1
+        end
+        return L2sum
+    else
+        error("Unrecognized method '$method'")
+    end
+end
 function local_error_norm(error::MPISharedArray{mk_float,5},
                           f::MPISharedArray{mk_float,5}, rtol, atol; method="Linf",
                           skip_r_inner=false, skip_z_lower=false, error_sum_zero=0.0)

From 4ef5c5008820a8d0d7af4a97ef2477fc40f453d4 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 23 Mar 2024 21:15:46 +0000
Subject: [PATCH 176/428] Use Runge-Kutta timestepping for electrons

Implement pseudo-timestepping for electrons in the same way as
timestepping for ions/neutrals. Input for electrons is given in the
`[electron_timestepping]` section of the input file - defaults are
copied from the `[timestepping]` section, but with `dt`-type options
reduced by sqrt(me/mi).
---
 .../wall+sheath-bc_boltzmann_loworder.toml    |  15 +-
 .../wall+sheath-bc_kinetic.toml               |  13 +-
 ...wall+sheath-bc_kinetic_krook_loworder.toml |  13 +-
 .../wall+sheath-bc_kinetic_loworder.toml      |  25 +-
 .../src/electron_fluid_equations.jl           |  25 +-
 .../src/electron_kinetic_equation.jl          | 799 +++++++++++-------
 moment_kinetics/src/electron_vpa_advection.jl |  19 +-
 moment_kinetics/src/electron_z_advection.jl   |  11 +-
 moment_kinetics/src/file_io.jl                | 160 +++-
 moment_kinetics/src/initial_conditions.jl     |  29 +-
 moment_kinetics/src/input_structs.jl          |   3 +-
 moment_kinetics/src/load_data.jl              |  21 +-
 moment_kinetics/src/moment_kinetics.jl        |  18 +-
 moment_kinetics/src/moment_kinetics_input.jl  |  84 ++
 moment_kinetics/src/time_advance.jl           | 139 ++-
 15 files changed, 945 insertions(+), 429 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
index aced83ff7..73a0871fc 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
@@ -49,14 +49,6 @@ ionization_frequency = 2.0
 #electron_ionization_frequency = 2.0
 #ionization_energy = 1.0
 constant_ionization_rate = false
-nstep = 1000000
-#nstep = 1
-dt = 1.0e-5
-nwrite = 10000
-nwrite_dfns = 10000
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
 r_ngrid = 1
 r_nelement = 1
 z_ngrid = 5
@@ -90,3 +82,10 @@ vpa_dissipation_coefficient = 0.002
 #vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
 #vpa_dissipation_coefficient = 20.0
+
+[timestepping]
+nstep = 1000000
+#nstep = 1
+dt = 1.0e-5
+nwrite = 10000
+nwrite_dfns = 10000
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index cd65ce083..be6965b24 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -47,13 +47,6 @@ ionization_frequency = 2.0
 electron_ionization_frequency = 2.0
 ionization_energy = 1.0
 constant_ionization_rate = false
-nstep = 40000
-#nstep = 1
-dt = 0.0005
-nwrite = 200
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
 r_ngrid = 1
 r_nelement = 1
 z_ngrid = 9
@@ -86,3 +79,9 @@ ascii_output = true
 #vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
 vpa_dissipation_coefficient = 20.0
+
+[timestepping]
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
index f0ae0e9c2..b8e8bd9bd 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
@@ -48,13 +48,6 @@ ionization_frequency = 2.0
 electron_ionization_frequency = 2.0
 ionization_energy = 1.0
 constant_ionization_rate = false
-nstep = 40000
-#nstep = 1
-dt = 0.0005
-nwrite = 200
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
 r_ngrid = 1
 r_nelement = 1
 z_ngrid = 5
@@ -87,3 +80,9 @@ ascii_output = true
 #vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
 vpa_dissipation_coefficient = 20.0
+
+[timestepping]
+nstep = 40000
+#nstep = 1
+dt = 0.0005
+nwrite = 200
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index 49819db60..17903e50f 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -49,14 +49,6 @@ ionization_frequency = 2.0
 #electron_ionization_frequency = 2.0
 #ionization_energy = 1.0
 constant_ionization_rate = false
-nstep = 40000
-#nstep = 1
-dt = 5.0e-5
-nwrite = 200
-nwrite_dfns = 10000
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
 r_ngrid = 1
 r_nelement = 1
 z_ngrid = 5
@@ -89,3 +81,20 @@ ascii_output = true
 vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
 #vpa_dissipation_coefficient = 20.0
+
+[timestepping]
+nstep = 40000
+#nstep = 1
+dt = 5.0e-5
+nwrite = 200
+nwrite_dfns = 10000
+type = "SSPRK4"
+
+[electron_timestepping]
+nstep = 50000
+#nstep = 1
+dt = 2.0e-8
+nwrite = 1000
+nwrite_dfns = 1000
+type = "SSPRK4"
+#type = "Fekete4(3)"
diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 9febf93f9..968ddc923 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -118,7 +118,7 @@ an explicit time advance.
 NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
-function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt, composition,
+function electron_energy_equation!(ppar, fvec, moments, collisions, dt, composition,
                                    num_diss_params, z)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
@@ -127,21 +127,21 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     # calculate contribution to rhs of energy equation (formulated in terms of pressure)
     # arising from derivatives of ppar, qpar and upar
     @loop_r_z ir iz begin
-        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.electron.dppar_dz_upwind[iz,ir]
-                    + moments.electron.dqpar_dz[iz,ir]
-                    + 3*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
+        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz_upwind[iz,ir]
+                    + moments.dqpar_dz[iz,ir]
+                    + 3*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
     end
     # @loop_r_z ir iz begin
-    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.electron.dppar_dz_upwind[iz,ir]
-    #                 + (2/3)*moments.electron.dqpar_dz[iz,ir]
-    #                 + (5/3)*fvec.electron_ppar[iz,ir]*moments.electron.dupar_dz[iz,ir])
+    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz_upwind[iz,ir]
+    #                 + (2/3)*moments.dqpar_dz[iz,ir]
+    #                 + (5/3)*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
     # end
     # compute the contribution to the rhs of the energy equation
     # arising from artificial diffusion
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_r_z ir iz begin
-            ppar[iz,ir] += dt*diffusion_coefficient*moments.electron.d2ppar_dz2[iz,ir]
+            ppar[iz,ir] += dt*diffusion_coefficient*moments.d2ppar_dz2[iz,ir]
         end
     end
     # compute the contribution to the rhs of the energy equation
@@ -149,7 +149,7 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     if nu_ei > 0.0
         @loop_s_r_z is ir iz begin
             ppar[iz,ir] += dt * (2 * me_over_mi * nu_ei * (fvec.ppar[iz,ir,is] - fvec.electron_ppar[iz,ir]))
-            ppar[iz,ir] += dt * ((2/3) * moments.electron.parallel_friction[iz,ir]
+            ppar[iz,ir] += dt * ((2/3) * moments.parallel_friction[iz,ir]
                                 * (fvec.upar[iz,ir,is]-fvec.electron_upar[iz,ir]))
         end
     end
@@ -181,18 +181,15 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
         end
     end
     # calculate the external electron heat source, if any
-    calculate_electron_heat_source!(moments.electron.heat_source, fvec.electron_ppar, moments.electron.dupar_dz,
+    calculate_electron_heat_source!(moments.heat_source, fvec.electron_ppar, moments.dupar_dz,
                                     fvec.density_neutral, collisions.ionization, collisions.ionization_energy,
                                     fvec.electron_density, fvec.ppar, collisions.nu_ei, composition.me_over_mi,
                                     composition.T_wall, z)
     # add the contribution from the electron heat source                                
     begin_r_z_region()
     @loop_r_z ir iz begin
-        ppar[iz,ir] += dt * moments.electron.heat_source[iz,ir]
+        ppar[iz,ir] += dt * moments.heat_source[iz,ir]
     end
-    # enforce the parallel boundary condtion on the electron parallel pressure
-    #println("!!!NO PARALLEL BC IS BEING ENFORCED ON ELECTRON PRESSURE!!!")
-    #enforce_parallel_BC_on_electron_pressure!(ppar, dens_i, composition.T_wall, fvec.ppar)
     return nothing
 end
 
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index f7eebb389..9010ba294 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -7,21 +7,25 @@ using MPI
 export get_electron_critical_velocities
 
 using ..looping
+using ..analysis: steady_state_residuals
 using ..derivatives: derivative_z!
 using ..boundary_conditions: enforce_v_boundary_condition_local!,
                              enforce_vperp_boundary_condition!
 using ..calculus: derivative!, second_derivative!, integral
 using ..communication
 using ..interpolation: interpolate_to_grid_1d!
-using ..type_definitions: mk_float
+using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float
 using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
-using ..electron_z_advection: electron_z_advection!
-using ..electron_vpa_advection: electron_vpa_advection!
+using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
+using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
 using ..file_io: write_initial_electron_state, finish_initial_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
+using ..runge_kutta: rk_update_variable!, rk_error_variable!, local_error_norm,
+                     adaptive_timestep_update_t_params!
+using ..utils: get_minimum_CFL_z, get_minimum_CFL_vpa
 using ..velocity_moments: integrate_over_vspace
 
 """
@@ -32,6 +36,7 @@ The electron kinetic equation is:
     zdot * d(pdf)/dz + wpadot * d(pdf)/dwpa = pdf * pre_factor
 
     INPUTS:
+    scratch = `scratch_pdf` struct used to store Runge-Kutta stages
     pdf = modified electron pdf @ previous time level = (true electron pdf / dens_e) * vth_e
     dens = electron density
     vthe = electron thermal speed
@@ -50,11 +55,11 @@ The electron kinetic equation is:
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
-function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
+function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
         phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
-        vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt, collisions, composition,
-        num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing,
-        initial_time=0.0)
+        vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
+        collisions, composition, num_diss_params, max_electron_pdf_iterations;
+        io_initial_electron=nothing, initial_time=0.0, evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -62,11 +67,11 @@ function update_electron_pdf!(fvec, pdf, moments, dens, vthe, ppar, qpar, qpar_u
     #solution_method = "picard_iteration"
     # solve the electron kinetic equation using the specified method
     if solution_method == "artificial_time_derivative"
-        return update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
-            moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition, 
-            r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
-            num_diss_params, max_electron_pdf_iterations;
-            io_initial_electron=io_initial_electron, initial_time=initial_time)
+        return update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi,
+            collisions, composition, r, z, vperp, vpa, z_spectral, vperp_spectral,
+            vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params, num_diss_params,
+            max_electron_pdf_iterations; io_initial_electron=io_initial_electron,
+            initial_time=initial_time, evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
         return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
             qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
@@ -102,35 +107,22 @@ The electron kinetic equation is:
     z_spectral = struct containing spectral information for the z-coordinate
     vpa_spectral = struct containing spectral information for the vpa-coordinate
     scratch_dummy = dummy arrays to be used for temporary storage
-    dt = time step size
     max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
     io_initial_electron = info struct for binary file I/O
     initial_time = initial value for the (pseudo-)time
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
-function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated, 
-    moments, dens, vthe, ppar, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, phi, collisions, composition,
-    r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, dt,
-    num_diss_params, max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0)
+function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, collisions,
+        composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
+        vpa_advect, scratch_dummy, t_params, num_diss_params, max_electron_pdf_iterations;
+        io_initial_electron=nothing, initial_time=0.0, evolve_ppar=false)
 
     begin_r_z_region()
 
-    #println("TMP FOR TESTING: SETTING UPAR_I = UPAR_E = 0!!!")
-    #moments.electron.upar .= 0.0
-
-    # there will be a better way of doing this
-    # store the incoming ppar
-    #ppar_old = allocate_float(z.n,1)
-    #ppar_old .= ppar
-
-    upar = moments.electron.upar
-    println("WARNING: when time-evolving the ions, this should really be something like fvec_in.upar, so that it is defined at the correct time level...")
-    upar_ion = moments.ion.upar
-
     # create a (z,r) dimension dummy array for use in taking derivatives
     dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
-    # create several (z) dimension dummy arrays for use in taking derivatives
+    # create several (r) dimension dummy arrays for use in taking derivatives
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
     buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
@@ -138,336 +130,408 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
     buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
 
+    @loop_r_z ir iz begin
+        dummy_zr[iz,ir] = -moments.electron.upar[iz,ir]
+    end
     # compute the z-derivative of the input electron parallel flow, needed for the electron kinetic equation
-    @views derivative_z!(moments.electron.dupar_dz, upar, buffer_r_1, buffer_r_2, buffer_r_3,
-                         buffer_r_4, z_spectral, z)
+    @views derivative_z!(moments.electron.dupar_dz, moments.electron.upar, buffer_r_1,
+                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
     # compute the z-derivative of the input electron parallel pressure, needed for the electron kinetic equation
-    @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
-                         buffer_r_4, z_spectral, z)
-
-    #println("TMP FOR TESTING -- dens and ddens_dz artificially set!!!")
-    #@. dens = 1.0
-    #@. ddens_dz = (dppar_dz / ppar)*0.9
+    @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, buffer_r_1,
+                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
+    begin_r_z_region()
     @loop_r_z ir iz begin
         # update the electron thermal speed using the updated electron parallel pressure
-        vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
+        moments.electron.vth[iz,ir] = sqrt(abs(2.0 * moments.electron.ppar[iz,ir] /
+                                                (moments.electron.dens[iz,ir] *
+                                                 composition.me_over_mi)))
         # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
         # and parallel pressure
-        dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
-        fvec.electron_ppar[iz,ir] = ppar[iz,ir]
+        moments.electron.dvth_dz[iz,ir] =
+            0.5 * moments.electron.vth[iz,ir] *
+            (moments.electron.dppar_dz[iz,ir] / moments.electron.ppar[iz,ir] -
+             moments.electron.ddens_dz[iz,ir] / moments.electron.dens[iz,ir])
+        scratch[1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
     end
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
-    @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
-                         buffer_r_4, z_spectral, z)
-
-    #dt_electron = dt * sqrt(composition.me_over_mi)
-    #dt_max = 3.0e-8
-    dt_max = 2.0e-8
-    #dt_max = 2.5e-9
-
-    #dt_energy = 1.0e-7
-    dt_energy = 2.0e-8
-    #dt_energy = 2.5e-9
-    #n_ppar_subcycles = 1000
-    #n_ppar_subcycles = 200
-    n_ppar_subcycles = 1
+    @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
+                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+
     time = initial_time
+    if initial_time > 0.0
+        # Make sure that output times are set relative to this initial_time (before this they
+        # will be set relative to 0.0).
+        t_params.moments_output_times .+= initial_time
+        t_params.dfns_output_times .+= initial_time
+    end
+    if io_initial_electron !== nothing
+        t_params.next_output_time[] = t_params.dfns_output_times[1]
+    end
 
     #z_speedup_fac = 20.0
     #z_speedup_fac = 5.0
     z_speedup_fac = 1.0
 
-    # define residual to point to a dummy array;
-    # to be filled with the contributions to the electron kinetic equation; i.e., d(pdf)/dt = -residual
-    residual = scratch_dummy.buffer_vpavperpzr_6
-    max_term = scratch_dummy.buffer_vpavperpzr_5
-    single_term = scratch_dummy.buffer_vpavperpzr_4
+    text_output = false
+
+    epsilon = 1.e-11
     # initialise the number of iterations in the solution of the electron kinetic equation to be 1
-    iteration = 1
+    t_params.step_counter[] = 1
     # initialise the electron pdf convergence flag to false
     electron_pdf_converged = false
-    # calculate the residual of the electron kinetic equation for the initial guess of the electron pdf
-    dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, upar_ion,
-                                        ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
-                                        z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                        collisions, num_diss_params, dt_max)
-    speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
-    if n_blocks[] == 1
-        text_output_suffix = ""
-    else
-        text_output_suffix = "$(iblock_index[])"
-    end
-    begin_serial_region()
-    @serial_region begin
-        # open files to write the electron heat flux and pdf to file                                
-        io_upar = open("upar$text_output_suffix.txt", "w")
-        io_qpar = open("qpar$text_output_suffix.txt", "w")
-        io_ppar = open("ppar$text_output_suffix.txt", "w")
-        io_pdf = open("pdf$text_output_suffix.txt", "w")
-        io_vth = open("vth$text_output_suffix.txt", "w")
-        if !electron_pdf_converged
-            # need to exit or handle this appropriately
-            @loop_vpa ivpa begin
+
+    if text_output
+        if n_blocks[] == 1
+            text_output_suffix = ""
+        else
+            text_output_suffix = "$(iblock_index[])"
+        end
+        begin_serial_region()
+        @serial_region begin
+            # open files to write the electron heat flux and pdf to file                                
+            io_upar = open("upar$text_output_suffix.txt", "w")
+            io_qpar = open("qpar$text_output_suffix.txt", "w")
+            io_ppar = open("ppar$text_output_suffix.txt", "w")
+            io_pdf = open("pdf$text_output_suffix.txt", "w")
+            io_vth = open("vth$text_output_suffix.txt", "w")
+            if !electron_pdf_converged
+                # need to exit or handle this appropriately
+                @loop_vpa ivpa begin
+                    @loop_z iz begin
+                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", scratch[1].pdf_electron[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                    end
+                    println(io_pdf,"")
+                end
                 @loop_z iz begin
-                    println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                    println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                    println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[], " dens: ", dens[iz,1])
                 end
-                println(io_pdf,"")
-            end
-            @loop_z iz begin
-                println(io_upar, "z: ", z.grid[iz], " upar: ", upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+                println(io_upar,"")
+                println(io_qpar,"")
+                println(io_ppar,"")
+                println(io_vth,"")
             end
-            println(io_upar,"")
-            println(io_qpar,"")
-            println(io_ppar,"")
-            println(io_vth,"")
+            io_pdf_stages = open("pdf_zright$text_output_suffix.txt", "w")
         end
-        io_pdf_stages = open("pdf_zright$text_output_suffix.txt", "w")
     end
 
-    # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
-    #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-    average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
-
-    output_interval = 1000
     output_counter = 0
     begin_serial_region()
+    output_counter += 1
     @serial_region begin
-        output_counter += 1
         if io_initial_electron !== nothing
-            write_initial_electron_state(pdf, moments, time, io_initial_electron,
-                                         output_counter, r, z, vperp, vpa)
+            write_initial_electron_state(scratch[1].pdf_electron, moments, t_params, time,
+                                         io_initial_electron, output_counter, r, z, vperp,
+                                         vpa)
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
-    while !electron_pdf_converged && (iteration <= max_electron_pdf_iterations)
-        #dt_energy = dt_electron * 10.0
-
-        # get an updated iterate of the electron parallel pressure
-        begin_r_z_region()
-        #ppar .= ppar_old
-        wpa3_moment = @view scratch_dummy.buffer_zrs_1[:,:,1]
-        @loop_r_z ir iz begin
-            wpa3_moment[iz,ir] = qpar[iz,ir] / vthe[iz,ir]^3
-        end
-        for i in 1:n_ppar_subcycles
-            @loop_r_z ir iz begin
-                qpar[iz,ir] = vthe[iz,ir]^3 * wpa3_moment[iz,ir]
-                dummy_zr[iz,ir] = -upar[iz,ir]
+    while !electron_pdf_converged && (t_params.step_counter[] < max_electron_pdf_iterations) && t_params.dt[] > 0.0
+        for istage ∈ 1:t_params.n_rk_stages
+            # Set the initial values for this stage to the final values from the previous
+            # stage
+            begin_r_z_vperp_vpa_region()
+            new_pdf = scratch[istage+1].pdf_electron
+            old_pdf = scratch[istage].pdf_electron
+            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+                new_pdf[ivpa,ivperp,iz,ir] = old_pdf[ivpa,ivperp,iz,ir]
             end
-            @views derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3,
-                                 buffer_r_4, z_spectral, z)
-            # Compute the upwinded z-derivative of the electron parallel pressure for the
-            # electron energy equation
-            @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
-                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
-                             buffer_r_5, buffer_r_6, z_spectral, z)
-            # centred second derivative for dissipation
-            if num_diss_params.moment_dissipation_coefficient > 0.0
-                @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
-                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-                begin_serial_region()
-                @serial_region begin
-                    if z.irank == 0
-                        moments.electron.d2ppar_dz2[1,:] .= 0.0
-                    end
-                    if z.irank == z.nrank - 1
-                        moments.electron.d2ppar_dz2[end,:] .= 0.0
-                    end
+            if evolve_ppar
+                begin_r_z_region()
+                new_ppar = scratch[istage+1].electron_ppar
+                old_ppar = scratch[istage].electron_ppar
+                @loop_r_z ir iz begin
+                    new_ppar[iz,ir] = old_ppar[iz,ir]
+                end
+            end
+            # Do a forward-Euler update of the electron pdf, and (if evove_ppar=true) the
+            # electron parallel pressure.
+            electron_kinetic_equation_euler_update!(scratch[istage+1], scratch[istage],
+                                                    moments, z, vperp, vpa, z_spectral,
+                                                    vpa_spectral, z_advect, vpa_advect,
+                                                    scratch_dummy, collisions,
+                                                    composition, num_diss_params,
+                                                    t_params.dt[];
+                                                    evolve_ppar=evolve_ppar)
+            speedup_hack!(scratch[istage+1], scratch[istage], z_speedup_fac, z, vpa;
+                          evolve_ppar=evolve_ppar)
+
+            rk_update_variable!(scratch, :pdf_electron, t_params, istage)
+            if evolve_ppar
+                rk_update_variable!(scratch, :electron_ppar, t_params, istage)
+                moments_struct_ppar = moments.electron.ppar
+                scratch_ppar = scratch[istage+1].electron_ppar
+                @loop_r_z ir iz begin
+                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
                 end
             end
 
-            # TMP FOR TESTING -- MAB
-            #dt_energy = dt_electron
-            electron_energy_equation!(ppar, dens, fvec, moments, collisions, dt_energy, composition, num_diss_params, z)
-
-            # Apply same 'speed up' hack to ppar that we do to the distribution function,
-            # but without the wpa dependence.
-            @loop_r_z ir iz begin
-                zval = z.grid[iz]
-                znorm = 2.0*zval/z.L
-                ppar[iz,ir] = fvec.electron_ppar[iz,ir] +
-                              (ppar[iz,ir] - fvec.electron_ppar[iz,ir]) *
-                              (1.0 + z_speedup_fac*(1.0 - znorm^2))
+            latest_pdf = scratch[istage+1].pdf_electron
+            begin_r_z_vperp_vpa_region()
+            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+                latest_pdf[ivpa,ivperp,iz,ir] = max(latest_pdf[ivpa,ivperp,iz,ir], 0.0)
             end
+
+            # enforce the boundary condition(s) on the electron pdf
+            enforce_boundary_condition_on_electron_pdf!(scratch[istage+1].pdf_electron, phi,
+                                                        moments.electron.vth,
+                                                        moments.electron.upar, vperp, vpa,
+                                                        vperp_spectral, vpa_spectral,
+                                                        vpa_advect, moments,
+                                                        num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                        composition.me_over_mi)
+
             begin_r_z_region()
+            A = moments.electron.constraints_A_coefficient
+            B = moments.electron.constraints_B_coefficient
+            C = moments.electron.constraints_C_coefficient
             @loop_r_z ir iz begin
-                fvec.electron_ppar[iz,ir] = ppar[iz,ir]
+                if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
+                    continue
+                end
+                (A[iz,ir], B[iz,ir], C[iz,ir]) =
+                    @views hard_force_moment_constraints!(latest_pdf[:,:,iz,ir],
+                                                          (evolve_density=true,
+                                                           evolve_upar=true,
+                                                           evolve_ppar=true), vpa)
             end
+            
+            function update_derived_moments_and_derivatives()
+                # update the electron heat flux
+                moments.electron.qpar_updated[] = false
+                calculate_electron_qpar_from_pdf!(moments.electron.qpar,
+                                                  moments.electron.ppar, moments.electron.vth,
+                                                  latest_pdf, vpa)
+
+                # compute the z-derivative of the parallel electron heat flux
+                @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar,
+                                     buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                                     z_spectral, z)
+
+                if evolve_ppar
+                    # get an updated iterate of the electron parallel pressure
+                    begin_r_z_region()
+                    # Compute the upwinded z-derivative of the electron parallel pressure for the
+                    # electron energy equation
+                    @views derivative_z!(moments.electron.dppar_dz_upwind,
+                                         scratch[istage+1].electron_ppar, dummy_zr, buffer_r_1,
+                                         buffer_r_2, buffer_r_3, buffer_r_4, buffer_r_5,
+                                         buffer_r_6, z_spectral, z)
+                    # compute the z-derivative of the updated electron parallel pressure
+                    @views derivative_z!(moments.electron.dppar_dz,
+                                         scratch[istage+1].electron_ppar, buffer_r_1, buffer_r_2,
+                                         buffer_r_3, buffer_r_4, z_spectral, z)
+
+                    this_ppar = scratch[istage+1].electron_ppar
+                    this_dppar_dz = moments.electron.dppar_dz
+                    this_ddens_dz = moments.electron.ddens_dz
+                    this_dens = moments.electron.dens
+                    this_vth = moments.electron.vth
+                    this_dvth_dz = moments.electron.dvth_dz
+                    @loop_r_z ir iz begin
+                        # update the electron thermal speed using the updated electron
+                        # parallel pressure
+                        this_vth[iz,ir] = sqrt(abs(2.0 * this_ppar[iz,ir] /
+                                                   (this_dens[iz,ir] *
+                                                    composition.me_over_mi)))
+                        # update the z-derivative of the electron thermal speed from the
+                        # z-derivatives of the electron density and parallel pressure
+                        this_dvth_dz[iz,ir] = 0.5 * this_vth[iz,ir] *
+                                              (this_dppar_dz[iz,ir] / this_ppar[iz,ir] -
+                                               this_ddens_dz[iz,ir] / this_dens[iz,ir])
+                    end
 
-            # compute the z-derivative of the updated electron parallel pressure
-            @views derivative_z!(dppar_dz, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
-                                 buffer_r_4, z_spectral, z)
-            begin_r_z_region()
-            @loop_r_z ir iz begin
-                # update the electron thermal speed using the updated electron parallel pressure
-                vthe[iz,ir] = sqrt(abs(2.0 * ppar[iz,ir] / (dens[iz,ir] * composition.me_over_mi)))
-                # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
-                # and parallel pressure
-                dvth_dz[iz,ir] = 0.5 * vthe[iz,ir] * (dppar_dz[iz,ir] / ppar[iz,ir] - ddens_dz[iz,ir] / dens[iz,ir])
+                    # centred second derivative for dissipation
+                    if num_diss_params.moment_dissipation_coefficient > 0.0
+                        @views derivative_z!(moments.electron.d2ppar_dz2,
+                                             moments.electron.dppar_dz, buffer_r_1, buffer_r_2,
+                                             buffer_r_3, buffer_r_4, z_spectral, z)
+                        begin_serial_region()
+                        @serial_region begin
+                            if z.irank == 0
+                                moments.electron.d2ppar_dz2[1,:] .= 0.0
+                            end
+                            if z.irank == z.nrank - 1
+                                moments.electron.d2ppar_dz2[end,:] .= 0.0
+                            end
+                        end
+                    end
+                end
+            end
+            update_derived_moments_and_derivatives()
+
+            if t_params.adaptive && istage == t_params.n_rk_stages
+                electron_adaptive_timestep_update!(scratch, time, t_params, moments,
+                                                   z_advect, vpa_advect, r, z, vperp, vpa;
+                                                   evolve_ppar=evolve_ppar)
+                # Re-do this in case electron_adaptive_timestep_update!() re-arranged the
+                # `scratch` vector
+                new_scratch = scratch[istage+1]
+                old_scratch = scratch[istage]
+
+                if t_params.previous_dt[] == 0.0
+                    # Re-calculate moments and moment derivatives as the timstep needs to
+                    # be re-done with a smaller dt, so scratch[t_params.n_rk_stages+1] has
+                    # been reset to the values from the beginning of the timestep here.
+                    update_derived_moments_and_derivatives()
+                end
             end
         end
 
-        begin_r_z_region()
-        # d(pdf)/dt = -kinetic_eqn_terms, so pdf_new = pdf - dt * kinetic_eqn_terms
-        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-            pdf[ivpa,ivperp,iz,ir] -= dt_electron * residual[ivpa,ivperp,iz,ir]
+        # update the time following the pdf update
+        time += t_params.previous_dt[]
+
+        residual = -1.0
+        if t_params.previous_dt[] > 0.0
+            # Calculate residuals to decide if iteration is converged.
+            # Might want an option to calculate the residual only after a certain number
+            # of iterations (especially during initialization when there are likely to be
+            # a large number of iterations required) to avoid the expense, and especially
+            # the global MPI.Bcast()?
+            begin_r_z_vperp_vpa_region()
+            residual = steady_state_residuals(scratch[t_params.n_rk_stages+1].pdf_electron,
+                                              scratch[1].pdf_electron, t_params.previous_dt[];
+                                              use_mpi=true, only_max_abs=true)
+            if global_rank[] == 0
+                residual = first(values(residual))[1]
+            end
+            if evolve_ppar
+                ppar_residual =
+                    steady_state_residuals(scratch[t_params.n_rk_stages+1].electron_ppar,
+                                           scratch[1].electron_ppar, t_params.previous_dt[];
+                                           use_mpi=true, only_max_abs=true)
+                if global_rank[] == 0
+                    ppar_residual = first(values(ppar_residual))[1]
+                    residual = max(residual, ppar_residual)
+                end
+            end
+            if global_rank[] == 0
+                electron_pdf_converged = abs(residual) < t_params.converged_residual_value
+            end
+            electron_pdf_converged = MPI.Bcast(electron_pdf_converged, 0, comm_world)
         end
 
+        # Set the initial values for the next step to the final values from the previous
+        # step
+        begin_r_z_vperp_vpa_region()
+        new_pdf = scratch[1].pdf_electron
+        old_pdf = scratch[t_params.n_rk_stages+1].pdf_electron
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-            pdf[ivpa,ivperp,iz,ir] = max(pdf[ivpa,ivperp,iz,ir], 0.0)
+            new_pdf[ivpa,ivperp,iz,ir] = old_pdf[ivpa,ivperp,iz,ir]
         end
-
-        # enforce the boundary condition(s) on the electron pdf
-        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
-                                                    vperp_spectral, vpa_spectral,
-                                                    vpa_advect, moments,
-                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
-                                                    composition.me_over_mi)
-        #println("A pdf 1 ", pdf[:,1,1,1])
-        #println("A pdf end ", pdf[:,1,end,1])
-
-        begin_r_z_region()
-        A = moments.electron.constraints_A_coefficient
-        B = moments.electron.constraints_B_coefficient
-        C = moments.electron.constraints_C_coefficient
-        @loop_r_z ir iz begin
-            if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
-                continue
+        if evolve_ppar
+            begin_r_z_region()
+            new_ppar = scratch[1].electron_ppar
+            old_ppar = scratch[t_params.n_rk_stages+1].electron_ppar
+            @loop_r_z ir iz begin
+                new_ppar[iz,ir] = old_ppar[iz,ir]
             end
-            (A[iz,ir], B[iz,ir], C[iz,ir]) =
-                #@views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=false, evolve_ppar=true), vpa)
-                @views hard_force_moment_constraints!(pdf[:,:,iz,ir], (evolve_density=true, evolve_upar=true, evolve_ppar=true), vpa)
         end
-        #println("B pdf 1 ", pdf[:,1,1,1])
-        #println("B pdf end ", pdf[:,1,end,1])
-        #error("foo")
-        
-        if (mod(iteration,output_interval)==1)
-            begin_serial_region()
-            @serial_region begin
-                @loop_vpa ivpa begin
-                    println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " iteration: ", iteration, " flag: ", 1)
+
+        if text_output
+            if (mod(t_params.step_counter[],t_params.nwrite)==1)
+                begin_serial_region()
+                @serial_region begin
+                    @loop_vpa ivpa begin
+                        println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa,1,end,1], " iteration: ", t_params.step_counter[], " flag: ", 1)
+                    end
+                    println(io_pdf_stages,"")
                 end
-                println(io_pdf_stages,"")
             end
         end
 
-        # update the time following the pdf update
-        time += dt_electron
-
-        qpar_updated = false
-        # update the electron heat flux
-        calculate_electron_qpar_from_pdf!(qpar, ppar, vthe, pdf, vpa)
-        qpar_updated = true
-        
-        # compute the z-derivative of the parallel electron heat flux
-        @views derivative_z!(dqpar_dz, qpar, buffer_r_1,
-            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-
-        # Compute the upwinded z-derivative of the electron parallel pressure for the 
-        # electron energy equation
-        @loop_r_z ir iz begin
-            dummy_zr[iz,ir] = -upar[iz,ir]
-        end
-        @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
-                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
-                             buffer_r_5, buffer_r_6, z_spectral, z)
-         # centred second derivative for dissipation
-        if num_diss_params.moment_dissipation_coefficient > 0.0
-            @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
-                                 buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-        end
-
-        if (mod(iteration,100) == 0)
+        if (mod(t_params.step_counter[],100) == 0)
             begin_serial_region()
-            @serial_region begin
-                println("iteration: ", iteration, " time: ", time, " dt_electron: ", dt_electron, " phi_boundary: ", phi[[1,end],1], " average_residual: ", average_residual)
+            if global_rank[] == 0
+                println("iteration: ", t_params.step_counter[], " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
             end
         end
-        if (mod(iteration,output_interval) == 0)
+        if (time ≥ t_params.dfns_output_times[output_counter] - epsilon)
             begin_serial_region()
             @serial_region begin
-                if (mod(iteration,100*output_interval) == 0)
-                    @loop_vpa ivpa begin
-                        @loop_z iz begin
-                            println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                if text_output
+                    if (mod(output_counter, 100) == 0)
+                        @loop_vpa ivpa begin
+                            @loop_z iz begin
+                                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                            end
+                            println(io_pdf,"")
                         end
                         println(io_pdf,"")
                     end
-                    println(io_pdf,"")
+                    @loop_z iz begin
+                        println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                        println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
+                        println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[], " dens: ", dens[iz,1])
+                    end
+                    println(io_upar,"")
+                    println(io_qpar,"")
+                    println(io_ppar,"")
+                    println(io_vth,"")
                 end
-                @loop_z iz begin
-                    println(io_upar, "z: ", z.grid[iz], " upar: ", upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", qpar[iz,1], " dqpar_dz: ", dqpar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", ppar[iz,1], " dppar_dz: ", dppar_dz[iz,1], " time: ", time, " iteration: ", iteration)
-                    println(io_vth, "z: ", z.grid[iz], " vthe: ", vthe[iz,1], " dvth_dz: ", dvth_dz[iz,1], " time: ", time, " iteration: ", iteration, " dens: ", dens[iz,1])
+            end
+            output_counter += 1
+            if output_counter ≤ length(t_params.dfns_output_times)
+                @serial_region begin
+                    t_params.next_output_time[] =
+                        t_params.dfns_output_times[output_counter]
                 end
-                println(io_upar,"")
-                println(io_qpar,"")
-                println(io_ppar,"")
-                println(io_vth,"")
-                output_counter += 1
+            end
+            @serial_region begin
                 if io_initial_electron !== nothing
-                    write_initial_electron_state(pdf, moments, time, io_initial_electron,
-                                                 output_counter, r, z, vperp, vpa)
+                    write_initial_electron_state(scratch[t_params.n_rk_stages+1].pdf_electron,
+                                                 moments, t_params, time,
+                                                 io_initial_electron, output_counter, r,
+                                                 z, vperp, vpa)
                 end
             end
         end
 
-        #@views derivative_z!(dqpar_dz, qpar, 
-        #        scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
-        #        scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-        # TMP FOR TESTING
-        #dqpar_dz .= 0.0
-        # calculate the residual of the electron kinetic equation for the updated electron pdf
-        dt_electron = electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vthe, ppar, upar_ion, ddens_dz,
-                                            dppar_dz, dqpar_dz, dvth_dz, 
-                                            z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                            collisions, num_diss_params, dt_max)
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
-        #average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, max_term)
-        average_residual, electron_pdf_converged = check_electron_pdf_convergence(residual, pdf, upar, vthe, z, vpa)
 
-        speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
         if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
             break
         end
-        iteration += 1
+        t_params.step_counter[] += 1
     end
-    # Update the 'scratch' arrays with the final result
+    # Update the 'pdf' arrays with the final result
     begin_r_z_vperp_vpa_region()
+    final_scratch_pdf = scratch[t_params.n_rk_stages+1].pdf_electron
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-        fvec.pdf_electron[ivpa,ivperp,iz,ir] = pdf[ivpa,ivperp,iz,ir]
+        pdf[ivpa,ivperp,iz,ir] = final_scratch_pdf[ivpa,ivperp,iz,ir]
     end
     begin_serial_region()
     @serial_region begin
-        if !electron_pdf_converged
-            @loop_vpa ivpa begin
-                @loop_z iz begin
-                    println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+        if text_output
+            if !electron_pdf_converged
+                @loop_vpa ivpa begin
+                    @loop_z iz begin
+                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                    end
+                    println(io_pdf,"")
                 end
-                println(io_pdf,"")
             end
+            close(io_upar)
+            close(io_qpar)
+            close(io_ppar)
+            close(io_vth)
+            close(io_pdf)
+            close(io_pdf_stages)
         end
-        close(io_upar)
-        close(io_qpar)
-        close(io_ppar)
-        close(io_vth)
-        close(io_pdf)
-        close(io_pdf_stages)
         if !electron_pdf_converged
             # need to exit or handle this appropriately
 
             if io_initial_electron !== nothing
                 output_counter += 1
-                write_initial_electron_state(pdf, moments, time, io_initial_electron,
-                                             output_counter, r, z, vperp, vpa)
+                write_initial_electron_state(final_scratch_pdf, moments, t_params, time,
+                                             io_initial_electron, output_counter, r, z,
+                                             vperp, vpa)
                 finish_initial_electron_io(io_initial_electron)
             end
 
@@ -480,7 +544,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
     return time, output_counter
 end
 
-function speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
+function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=false)
     # Divide by wpa to relax CFL condition at large wpa - only looking for steady
     # state here, so does not matter that this makes time evolution incorrect.
     # Also increase the effective timestep for z-values far from the sheath boundary -
@@ -494,12 +558,30 @@ function speedup_hack_residual!(residual, z_speedup_fac, z, vpa)
     vpa_fudge_factor = 0.0
 
     Lz = z.L
+
+    if evolve_ppar
+        begin_r_z_region()
+        ppar_out = fvec_out.electron_ppar
+        ppar_in = fvec_in.electron_ppar
+        @loop_r_z ir iz begin
+            zval = z.grid[iz]
+            znorm = 2.0*zval/Lz
+            ppar_out[iz,ir] = ppar_in[iz,ir] +
+                (1.0 + z_speedup_fac*(1.0 - znorm^2)) *
+                (ppar_out[iz,ir] - ppar_in[iz,ir])
+        end
+    end
+
+    begin_r_z_vperp_vpa_region()
+    pdf_out = fvec_out.pdf_electron
+    pdf_in = fvec_in.pdf_electron
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
         zval = z.grid[iz]
         znorm = 2.0*zval/Lz
-        residual[ivpa,ivperp,iz,ir] *=
+        pdf_out[ivpa,ivperp,iz,ir] = pdf_in[ivpa,ivperp,iz,ir] +
             (1.0 + z_speedup_fac*(1.0 - znorm^2)) /
-            sqrt(1.0 + vpa_fudge_factor * vpa.grid[ivpa]^2)
+            sqrt(1.0 + vpa_fudge_factor * vpa.grid[ivpa]^2) *
+            (pdf_out[ivpa,ivperp,iz,ir] - pdf_in[ivpa,ivperp,iz,ir])
     end
     return nothing
 end
@@ -1065,6 +1147,96 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
     # end
 end
 
+"""
+    electron_adaptive_timestep_update!(scratch, t_params, rk_coefs, moments)
+
+Check the error estimate for the embedded RK method and adjust the timestep if
+appropriate.
+"""
+function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_advect,
+                                            vpa_advect, r, z, vperp, vpa;
+                                            evolve_ppar=false)
+    #error_norm_method = "Linf"
+    error_norm_method = "L2"
+
+    error_coeffs = t_params.rk_coefs[:,end]
+    if t_params.n_rk_stages < 3
+        # This should never happen as an adaptive RK scheme needs at least 2 RHS evals so
+        # (with the pre-timestep data) there must be at least 3 entries in `scratch`.
+        error("adaptive timestep needs a buffer scratch array")
+    end
+
+    CFL_limits = mk_float[]
+    error_norm_type = typeof(t_params.error_sum_zero)
+    error_norms = error_norm_type[]
+    total_points = mk_int[]
+
+    # Read the current dt here, so we only need one _block_synchronize() call for this and
+    # the begin_s_r_z_vperp_vpa_region()
+    current_dt = t_params.dt[]
+    _block_synchronize()
+
+    # Test CFL conditions for advection in electron kinetic equation to give stability
+    # limit for timestep
+    #
+    # z-advection
+    # No need to synchronize here, as we just called _block_synchronize()
+    begin_r_z_vperp_vpa_region(; no_synchronize=true)
+    update_electron_speed_z!(z_advect[1], moments.electron.upar, moments.electron.vth,
+                             vpa.grid)
+    z_CFL = get_minimum_CFL_z(z_advect[1].speed, z)
+    if block_rank[] == 0
+        push!(CFL_limits, t_params.CFL_prefactor * z_CFL)
+    else
+        push!(CFL_limits, Inf)
+    end
+
+    # vpa-advection
+    update_electron_speed_vpa!(vpa_advect[1], moments.electron.ppar, moments.electron.vth,
+                               moments.electron.dppar_dz, moments.electron.dqpar_dz,
+                               moments.electron.dvth_dz, vpa.grid)
+    vpa_CFL = get_minimum_CFL_vpa(vpa_advect[1].speed, vpa)
+    if block_rank[] == 0
+        push!(CFL_limits, t_params.CFL_prefactor * vpa_CFL)
+    else
+        push!(CFL_limits, Inf)
+    end
+
+    # To avoid double counting points when we use distributed-memory MPI, skip the
+    # inner/lower point in r and z if this process is not the first block in that
+    # dimension.
+    skip_r_inner = r.irank != 0
+    skip_z_lower = z.irank != 0
+
+    # Calculate error ion distribution functions
+    # Note rk_error_variable!() stores the calculated error in `scratch[2]`.
+    rk_error_variable!(scratch, :pdf_electron, t_params)
+    pdf_error = local_error_norm(scratch[2].pdf_electron,
+                                 scratch[t_params.n_rk_stages+1].pdf_electron,
+                                 t_params.rtol, t_params.atol; method=error_norm_method,
+                                 skip_r_inner=skip_r_inner, skip_z_lower=skip_z_lower,
+                                 error_sum_zero=t_params.error_sum_zero)
+    push!(error_norms, pdf_error)
+    push!(total_points, vpa.n_global * vperp.n_global * z.n_global * r.n_global)
+
+    # Calculate error for moments, if necessary
+    if evolve_ppar
+        begin_r_z_region()
+        rk_error_variable!(scratch, :electron_ppar, t_params)
+        p_err = local_error_norm(scratch[2].electron_ppar,
+                                 scratch[t_params.n_rk_stages+1].electron_ppar,
+                                 t_params.rtol, t_params.atol; method=error_norm_method,
+                                 skip_r_inner=skip_r_inner, skip_z_lower=skip_z_lower,
+                                 error_sum_zero=t_params.error_sum_zero)
+        push!(error_norms, p_err)
+        push!(total_points, z.n_global * r.n_global)
+    end
+
+    adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
+                                       total_points, current_dt, error_norm_method)
+
+    return nothing
+end
 
 """
 update_electron_pdf_with_shooting_method is a function that using a shooting method to solve for the electron pdf
@@ -1264,6 +1436,61 @@ function update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens
     return nothing
 end
 
+"""
+    electron_kinetic_equation_euler_update!(fvec, pdf, moments, z, vperp, vpa,
+                                            z_spectral, vpa_spectral, z_advect,
+                                            vpa_advect, scratch_dummy, collisions,
+                                            num_diss_params, dt; evolve_ppar=false)
+
+Do a forward-Euler update of the electron kinetic equation.
+
+When `evolve_ppar=true` is passed, also updates the electron parallel pressure.
+"""
+function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z, vperp,
+                                                 vpa, z_spectral, vpa_spectral, z_advect,
+                                                 vpa_advect, scratch_dummy, collisions,
+                                                 composition, num_diss_params, dt;
+                                                 evolve_ppar=false)
+    # add the contribution from the z advection term
+    electron_z_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                          moments.electron.upar, moments.electron.vth, z_advect, z,
+                          vpa.grid, z_spectral, scratch_dummy, dt)
+
+    # add the contribution from the wpa advection term
+    electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                            fvec_in.electron_ppar, moments.electron.vth,
+                            moments.electron.dppar_dz, moments.electron.dqpar_dz,
+                            moments.electron.dvth_dz, vpa_advect, vpa, vpa_spectral,
+                            scratch_dummy, dt)
+
+    # add in the contribution to the residual from the term proportional to the pdf
+    add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                                    fvec_in.electron_ppar, moments.electron.vth,
+                                    moments.electron.dens, moments.electron.ddens_dz,
+                                    moments.electron.dvth_dz, moments.electron.dqpar_dz,
+                                    vpa.grid, z, dt)
+
+    # add in numerical dissipation terms
+    add_dissipation_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, scratch_dummy,
+                          z_spectral, z, vpa, vpa_spectral, num_diss_params, dt)
+
+    if collisions.krook_collision_frequency_prefactor_ee > 0.0
+        # Add a Krook collision operator
+        # Set dt=-1 as we update the residual here rather than adding an update to
+        # 'fvec_out'.
+        electron_krook_collisions!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                                   moments.electron.dens, moments.electron.upar,
+                                   moments.ion.upar, moments.electron.vth, collisions,
+                                   vperp, vpa, dt)
+    end
+
+    if evolve_ppar
+        electron_energy_equation!(fvec_out.electron_ppar, fvec_in, moments.electron,
+                                  collisions, dt, composition, num_diss_params, z)
+    end
+
+    return nothing
+end
 
 """
 electron_kinetic_equation_residual! calculates the residual of the (time-independent) electron kinetic equation
@@ -1283,7 +1510,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
         residual[ivpa,ivperp,iz,ir] = 0.0
     end
     # calculate the contribution to the residual from the z advection term
-    electron_z_advection!(residual, pdf, upar, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy)
+    electron_z_advection!(residual, pdf, upar, vth, z_advect, z, vpa.grid, z_spectral, scratch_dummy, -1.0)
     #dt_max_zadv = simple_z_advection!(residual, pdf, vth, z, vpa.grid, dt_electron)
     #single_term .= residual
     #max_term .= abs.(residual)
@@ -1292,7 +1519,7 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
     # add in the contribution to the residual from the wpa advection term
     electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
-                            vpa_advect, vpa, vpa_spectral, scratch_dummy)#, z)
+                            vpa_advect, vpa, vpa_spectral, scratch_dummy, -1.0)
     #dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
     #@. single_term = residual - single_term
     #max_term .= max.(max_term, abs.(single_term))
@@ -1300,7 +1527,8 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #println("v_adv residual = ", maximum(abs.(single_term)))
     #add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
     # add in the contribution to the residual from the term proportional to the pdf
-    add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+    add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz,
+                                    dqpar_dz, vpa.grid, z, -1.0)
     #@. single_term = residual - single_term
     #max_term .= max.(max_term, abs.(single_term))
     #@. single_term = residual
@@ -1313,7 +1541,8 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     # end
     # println("")
     # add in numerical dissipation terms
-    add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
+    add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral,
+                          num_diss_params, -1.0)
     #@. single_term = residual - single_term
     #println("dissipation residual = ", maximum(abs.(single_term)))
     #max_term .= max.(max_term, abs.(single_term))
@@ -1429,7 +1658,8 @@ function add_source_term!(source_term, vpa, z, dvth_dz)
     return nothing
 end
 
-function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa, vpa_spectral, num_diss_params)
+function add_dissipation_term!(pdf_out, pdf, scratch_dummy, z_spectral, z, vpa,
+                               vpa_spectral, num_diss_params, dt)
     dummy_zr1 = @view scratch_dummy.dummy_zrs[:,:,1]
     dummy_zr2 = @view scratch_dummy.buffer_vpavperpzr_1[1,1,:,:]
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
@@ -1450,7 +1680,7 @@ function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa,
         #@views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
         #@. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
         @views second_derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, vpa_spectral)
-        @. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch
+        @. pdf_out[:,ivperp,iz,ir] += dt * num_diss_params.vpa_dissipation_coefficient * vpa.scratch
     end
     #stop()
     return nothing
@@ -1654,7 +1884,7 @@ function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_
 end
 
 # add contribution to the residual coming from the term proporational to the pdf
-function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z)
+function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z, dt)
     begin_r_z_vperp_vpa_region()
     @loop_r_z ir iz begin
         this_dqpar_dz = dqpar_dz[iz,ir]
@@ -1665,8 +1895,11 @@ function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_d
         this_dvth_dz = dvth_dz[iz,ir]
         this_vth = vth[iz,ir]
         @loop_vperp_vpa ivperp ivpa begin
-            residual[ivpa,ivperp,iz,ir] -= (-0.5 * this_dqpar_dz / this_ppar - vpa[ivpa] * this_vth *
-                                    (this_ddens_dz / this_dens - this_dvth_dz / this_vth)) * pdf[ivpa,ivperp,iz,ir]
+            residual[ivpa,ivperp,iz,ir] += dt * (-0.5 * this_dqpar_dz / this_ppar -
+                                                 vpa[ivpa] * this_vth *
+                                                 (this_ddens_dz / this_dens
+                                                  - this_dvth_dz / this_vth)) *
+                                           pdf[ivpa,ivperp,iz,ir]
             #residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf[ivpa, ivperp, :, :]
         end
     end
diff --git a/moment_kinetics/src/electron_vpa_advection.jl b/moment_kinetics/src/electron_vpa_advection.jl
index 8aa8c5de1..c41e884e4 100644
--- a/moment_kinetics/src/electron_vpa_advection.jl
+++ b/moment_kinetics/src/electron_vpa_advection.jl
@@ -12,11 +12,13 @@ using ..calculus: derivative!, second_derivative!
 calculate the wpa-advection term for the electron kinetic equation 
 = (vthe / 2 ppare * dppare/dz + wpa / 2 ppare * dqpare/dz - wpa^2 * dvthe/dz) * df/dwpa
 """
-function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
-                                 advect, vpa, spectral, scratch_dummy)
+function electron_vpa_advection!(pdf_out, pdf_in, ppar, vth, dppar_dz, dqpar_dz, dvth_dz,
+                                 advect, vpa, spectral, scratch_dummy, dt)
+    begin_r_z_vperp_region()
+
     # create a reference to a scratch_dummy array to store the wpa-derivative of the electron pdf
     dpdf_dvpa = scratch_dummy.buffer_vpavperpzr_1
-    d2pdf_dvpa2 = scratch_dummy.buffer_vpavperpzr_2
+    #d2pdf_dvpa2 = scratch_dummy.buffer_vpavperpzr_2
     begin_r_z_vperp_region()
     # get the updated speed along the wpa direction using the current pdf
     @views update_electron_speed_vpa!(advect[1], ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa.grid)
@@ -27,18 +29,19 @@ function electron_vpa_advection!(advection_term, pdf, ppar, vth, dppar_dz, dqpar
     end
     #calculate the upwind derivative of the electron pdf w.r.t. wpa
     @loop_r_z_vperp ir iz ivperp begin
-        @views derivative!(dpdf_dvpa[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, advect[1].adv_fac[:,ivperp,iz,ir], spectral)
+        @views derivative!(dpdf_dvpa[:,ivperp,iz,ir], pdf_in[:,ivperp,iz,ir], vpa,
+                           advect[1].adv_fac[:,ivperp,iz,ir], spectral)
     end
     #@loop_r_z_vperp ir iz ivperp begin
-    #    @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf[:,ivperp,iz,ir], vpa, spectral)
+    #    @views second_derivative!(d2pdf_dvpa2[:,ivperp,iz,ir], pdf_in[:,ivperp,iz,ir], vpa, spectral)
     #end
     # calculate the advection term
     @loop_r_z_vperp ir iz ivperp begin
-        @. advection_term[:,ivperp,iz,ir] -= advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir]
-        #@. advection_term[:,ivperp,iz,ir] -= advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir] + 0.0001*d2pdf_dvpa2[:,ivperp,iz,ir]
+        @. pdf_out[:,ivperp,iz,ir] += dt * advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir]
+        #@. pdf_out[:,ivperp,iz,ir] -= advect[1].adv_fac[:,ivperp,iz,ir] * dpdf_dvpa[:,ivperp,iz,ir] + 0.0001*d2pdf_dvpa2[:,ivperp,iz,ir]
     end
     #@loop_vpa ivpa begin
-    #    println("electron_vpa_advection: ", advection_term[ivpa,1,10,1], " vpa: ", vpa.grid[ivpa], " dpdf_dvpa: ", dpdf_dvpa[ivpa,1,10,1],
+    #    println("electron_vpa_advection: ", pdf_out[ivpa,1,10,1], " vpa: ", vpa.grid[ivpa], " dpdf_dvpa: ", dpdf_dvpa[ivpa,1,10,1],
     #        " pdf: ", pdf[ivpa,1,10,1])
     #end
     #exit()
diff --git a/moment_kinetics/src/electron_z_advection.jl b/moment_kinetics/src/electron_z_advection.jl
index 10f1e8b2a..b0b8f26aa 100644
--- a/moment_kinetics/src/electron_z_advection.jl
+++ b/moment_kinetics/src/electron_z_advection.jl
@@ -14,7 +14,10 @@ using ..calculus: second_derivative!, derivative!
 """
 calculate the z-advection term for the electron kinetic equation = wpa * vthe * df/dz
 """
-function electron_z_advection!(advection_term, pdf, upar, vth, advect, z, vpa, spectral, scratch_dummy)
+function electron_z_advection!(pdf_out, pdf_in, upar, vth, advect, z, vpa, spectral,
+                               scratch_dummy, dt)
+    begin_r_vperp_vpa_region()
+
     # create a pointer to a scratch_dummy array to store the z-derivative of the electron pdf
     dpdf_dz = scratch_dummy.buffer_vpavperpzr_1
     d2pdf_dz2 = scratch_dummy.buffer_vpavperpzr_2
@@ -27,7 +30,7 @@ function electron_z_advection!(advection_term, pdf, upar, vth, advect, z, vpa, s
         @views advect[1].adv_fac[:,ivpa,ivperp,ir] = -advect[1].speed[:,ivpa,ivperp,ir]
     end
     #calculate the upwind derivative
-    derivative_z!(dpdf_dz, pdf,
+    derivative_z!(dpdf_dz, pdf_in,
                   advect, scratch_dummy.buffer_vpavperpr_1,
                   scratch_dummy.buffer_vpavperpr_2, scratch_dummy.buffer_vpavperpr_3,
                   scratch_dummy.buffer_vpavperpr_4, scratch_dummy.buffer_vpavperpr_5,
@@ -37,8 +40,8 @@ function electron_z_advection!(advection_term, pdf, upar, vth, advect, z, vpa, s
     #end
     # calculate the advection term
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-        advection_term[ivpa,ivperp,iz,ir] -= advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir]
-        #advection_term[ivpa,ivperp,iz,ir] -= advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir] + 0.0001*d2pdf_dz2[ivpa,ivperp,iz,ir]
+        pdf_out[ivpa,ivperp,iz,ir] += dt * advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir]
+        #pdf_out[ivpa,ivperp,iz,ir] += dt * advect[1].adv_fac[iz,ivpa,ivperp,ir] * dpdf_dz[ivpa,ivperp,iz,ir] + 0.0001*d2pdf_dz2[ivpa,ivperp,iz,ir]
     end
     return nothing
 end
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index bd102bf58..08877aa76 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -56,7 +56,8 @@ moments & fields only
 struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
                        Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
                        Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Tconstri, Tconstrn,
-                       Tconstre, Tint, Tfailcause}
+                       Tconstre, Tint, Tfailcause, Telectrontime, Telectronint,
+                       Telectronfailcause}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the time variable
@@ -141,6 +142,19 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     # Last successful timestep before most recent timestep failure, used by adaptve
     # timestepping algorithm
     dt_before_last_fail::Ttime
+    # cumulative number of electron pseudo-timesteps taken
+    electron_step_counter::Telectronint
+    # current electron pseudo-timestep size
+    electron_dt::Telectrontime
+    # cumulative number of electron pseudo-timestep failures
+    electron_failure_counter::Telectronint
+    # cumulative count of which variable caused a electron pseudo-timstep failure
+    electron_failure_caused_by::Telectronfailcause
+    # cumulative count of which factors limited the electron pseudo-timestep at each step
+    electron_limit_caused_by::Telectronfailcause
+    # Last successful timestep before most recent electron pseudo-timestep failure, used
+    # by adaptve timestepping algorithm
+    electron_dt_before_last_fail::Telectrontime
 
     # Use parallel I/O?
     parallel_io::Bool
@@ -171,7 +185,8 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr}
+struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr, Telectrontime,
+                                Telectronint, Telectronfailcause}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the pseudotime variable
@@ -192,6 +207,19 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr}
     electron_constraints_A_coefficient::Tconstr
     electron_constraints_B_coefficient::Tconstr
     electron_constraints_C_coefficient::Tconstr
+    # cumulative number of electron pseudo-timesteps taken
+    electron_step_counter::Telectronint
+    # current electron pseudo-timestep size
+    electron_dt::Telectrontime
+    # cumulative number of electron pseudo-timestep failures
+    electron_failure_counter::Telectronint
+    # cumulative count of which variable caused a electron pseudo-timstep failure
+    electron_failure_caused_by::Telectronfailcause
+    # cumulative count of which factors limited the electron pseudo-timestep at each step
+    electron_limit_caused_by::Telectronfailcause
+    # Last successful timestep before most recent electron pseudo-timestep failure, used
+    # by adaptve timestepping algorithm
+    electron_dt_before_last_fail::Telectrontime
 
     # Use parallel I/O?
     parallel_io::Bool
@@ -339,11 +367,13 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
                                                  description="electron distribution function")
 
         io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-        io_electron_vth =
+        io_electron_vth, io_electron_step_counter, io_electron_dt,
+        io_electron_failure_counter, io_electron_failure_caused_by,
+        io_electron_limit_caused_by, io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar)
+                                                      evolve_ppar, true)
 
         close(fid)
 
@@ -379,6 +409,12 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_constraints_A_coefficient"),
                                         getvar("electron_constraints_B_coefficient"),
                                         getvar("electron_constraints_C_coefficient"),
+                                        getvar("electron_step_counter"),
+                                        getvar("electron_dt"),
+                                        getvar("electron_failure_counter"),
+                                        getvar("electron_failure_caused_by"),
+                                        getvar("electron_limit_caused_by"),
+                                        getvar("electron_dt_before_last_fail"),
                                         parallel_io)
     end
 
@@ -789,7 +825,7 @@ define dynamic (time-evolving) moment variables for writing to the hdf5 file
 function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                           r::coordinate, z::coordinate, parallel_io,
                                           external_source_settings, evolve_density,
-                                          evolve_upar, evolve_ppar)
+                                          evolve_upar, evolve_ppar, kinetic_electrons)
     @serial_region begin
         dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
 
@@ -811,11 +847,14 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
 
         io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
         io_electron_vth, electron_constraints_A_coefficient,
-        electron_constraints_B_coefficient, electron_constraints_C_coefficient =
+        electron_constraints_B_coefficient, electron_constraints_C_coefficient,
+        io_electron_step_counter, io_electron_dt, io_electron_failure_counter,
+        io_electron_failure_caused_by, io_electron_limit_caused_by,
+        io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar)
+                                                      evolve_ppar, kinetic_electrons)
 
         io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
         io_thermal_speed_neutral, external_source_neutral_amplitude,
@@ -897,7 +936,11 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                electron_constraints_C_coefficient,
                                io_time_for_run, io_step_counter, io_dt,
                                io_failure_counter, io_failure_caused_by,
-                               io_limit_caused_by, io_dt_before_last_fail, parallel_io)
+                               io_limit_caused_by, io_dt_before_last_fail,
+                               io_electron_step_counter, io_electron_dt,
+                               io_electron_failure_counter, io_electron_failure_caused_by,
+                               io_electron_limit_caused_by,
+                               io_electron_dt_before_last_fail, parallel_io)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1099,7 +1142,8 @@ end
 define dynamic (time-evolving) electron moment variables for writing to the hdf5 file
 """
 function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordinate,
-        parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar)
+        parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar,
+        kinetic_electrons)
 
     dynamic = get_group(fid, "dynamic_data")
 
@@ -1146,9 +1190,52 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                parallel_io=parallel_io,
                                description="'C' coefficient enforcing pressure constraint for electrons")
 
+    if kinetic_electrons
+        io_electron_step_counter = create_dynamic_variable!(
+            dynamic, "electron_step_counter", mk_int; parallel_io=parallel_io,
+            description="cumulative number of electron pseudo-timesteps for the run")
+
+        io_electron_dt = create_dynamic_variable!(
+            dynamic, "electron_dt", mk_float; parallel_io=parallel_io,
+            description="current electron pseudo-timestep size")
+
+        io_electron_failure_counter = create_dynamic_variable!(
+            dynamic, "electron_failure_counter", mk_int; parallel_io=parallel_io,
+            description="cumulative number of electron pseudo-timestep failures for the run")
+
+        n_failure_vars = 1 + 1
+        io_electron_failure_caused_by = create_dynamic_variable!(
+            dynamic, "electron_failure_caused_by", mk_int;
+            diagnostic_var_size=n_failure_vars, parallel_io=parallel_io,
+            description="cumulative count of how many times each variable caused an "
+                        * "electron pseudo-timestep failure for the run")
+
+        n_limit_vars = 5 + 2
+        io_electron_limit_caused_by = create_dynamic_variable!(
+            dynamic, "electron_limit_caused_by", mk_int; diagnostic_var_size=n_limit_vars,
+            parallel_io=parallel_io,
+            description="cumulative count of how many times each factor limited the "
+                        * "electron pseudo-timestep for the run")
+
+        io_electron_dt_before_last_fail = create_dynamic_variable!(
+            dynamic, "electron_dt_before_last_fail", mk_float; parallel_io=parallel_io,
+            description="Last successful electron pseudo-timestep before most recent "
+                        * "electron pseudo-timestep failure, used by adaptve "
+                        * "timestepping algorithm")
+    else
+        io_electron_step_counter = nothing
+        io_electron_dt = nothing
+        io_electron_failure_counter = nothing
+        io_electron_failure_caused_by = nothing
+        io_electron_limit_caused_by = nothing
+        io_electron_dt_before_last_fail = nothing
+    end
+
     return io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
            io_electron_vth, electron_constraints_A_coefficient, electron_constraints_B_coefficient,
-           electron_constraints_C_coefficient
+           electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
+           io_electron_failure_counter, io_electron_failure_caused_by,
+           io_electron_limit_caused_by, io_electron_dt_before_last_fail
 end
 
 """
@@ -1294,7 +1381,8 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
                                                       parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar)
+                                                      evolve_ppar,
+                                                      composition.electron_physics == kinetic_electrons)
 
         dynamic = get_group(fid, "dynamic_data")
 
@@ -1405,7 +1493,7 @@ function setup_moments_io(prefix, binary_format, vz, vr, vzeta, vpa, vperp, r, z
         io_moments = define_dynamic_moment_variables!(
             fid, composition.n_ion_species, composition.n_neutral_species, r, z,
             parallel_io, external_source_settings, evolve_density, evolve_upar,
-            evolve_ppar)
+            evolve_ppar, composition.electron_physics == kinetic_electrons)
 
         close(fid)
 
@@ -1470,7 +1558,11 @@ function reopen_moments_io(file_info)
                                getvar("time_for_run"), getvar("step_counter"),
                                getvar("dt"), getvar("failure_counter"),
                                getvar("failure_caused_by"), getvar("limit_caused_by"),
-                               getvar("dt_before_last_fail"), parallel_io)
+                               getvar("dt_before_last_fail"),getvar("electron_step_counter"),
+                               getvar("electron_dt"), getvar("electron_failure_counter"),
+                               getvar("electron_failure_caused_by"),
+                               getvar("electron_limit_caused_by"),
+                               getvar("electron_dt_before_last_fail"), parallel_io)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1586,7 +1678,13 @@ function reopen_dfns_io(file_info)
                                      getvar("dt"), getvar("failure_counter"),
                                      getvar("failure_caused_by"),
                                      getvar("limit_caused_by"),
-                                     getvar("dt_before_last_fail"), parallel_io)
+                                     getvar("dt_before_last_fail"),
+                                     getvar("electron_step_counter"),
+                                     getvar("electron_dt"),
+                                     getvar("electron_failure_counter"),
+                                     getvar("electron_failure_caused_by"),
+                                     getvar("electron_limit_caused_by"),
+                                     getvar("electron_dt_before_last_fail"), parallel_io)
 
         return io_dfns_info(fid, getvar("f"), getvar("f_electron"), getvar("f_neutral"),
                             parallel_io, io_moments)
@@ -1644,7 +1742,8 @@ function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
 
         write_ion_moments_data_to_binary(moments, n_ion_species, io_moments, t_idx, r, z)
 
-        write_electron_moments_data_to_binary(moments, io_moments, t_idx, r, z)
+        write_electron_moments_data_to_binary(moments, t_params.electron, io_moments,
+                                              t_idx, r, z)
 
         write_neutral_moments_data_to_binary(moments, n_neutral_species, io_moments,
                                              t_idx, r, z)
@@ -1785,7 +1884,7 @@ write time-dependent moments data for electrons to the binary output file
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_electron_moments_data_to_binary(moments,
+function write_electron_moments_data_to_binary(moments, t_params,
                                                io_moments::Union{io_moments_info,io_initial_electron_info},
                                                t_idx, r, z)
     @serial_region begin
@@ -1812,6 +1911,21 @@ function write_electron_moments_data_to_binary(moments,
         append_to_dynamic_var(io_moments.electron_constraints_C_coefficient,
                               moments.electron.constraints_C_coefficient, t_idx,
                               parallel_io, z, r)
+
+        if t_params !== nothing
+            # Save timestepping info
+            append_to_dynamic_var(io_moments.electron_step_counter, t_params.step_counter[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_dt, t_params.dt_before_output[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_failure_counter, t_params.failure_counter[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_failure_caused_by, t_params.failure_caused_by,
+                                  t_idx, parallel_io, length(t_params.failure_caused_by);
+                                  only_root=true)
+            append_to_dynamic_var(io_moments.electron_limit_caused_by, t_params.limit_caused_by, t_idx,
+                                  parallel_io, length(t_params.limit_caused_by);
+                                  only_root=true)
+            append_to_dynamic_var(io_moments.electron_dt_before_last_fail,
+                                  t_params.dt_before_last_fail[], t_idx, parallel_io)
+        end
     end
 
     return nothing
@@ -1991,13 +2105,15 @@ function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
 end
 
 """
-    write_initial_electron_state(pdf, moments, io_initial_electron, t_idx, r, z,
+    write_initial_electron_state(pdf, moments, t_params, t,
+                                 io_initial_electron, t_idx, r, z,
                                  vperp, vpa)
 
 Write the electron state to an output file.
 """
-function write_initial_electron_state(pdf, moments, t, io_or_file_info_initial_electron,
-                                      t_idx, r, z, vperp, vpa)
+function write_initial_electron_state(pdf, moments, t_params, t,
+                                      io_or_file_info_initial_electron, t_idx, r, z,
+                                      vperp, vpa)
 
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -2018,7 +2134,8 @@ function write_initial_electron_state(pdf, moments, t, io_or_file_info_initial_e
         write_electron_dfns_data_to_binary(pdf, io_initial_electron, t_idx, r, z, vperp,
                                            vpa)
 
-        write_electron_moments_data_to_binary(moments, io_initial_electron, t_idx, r, z)
+        write_electron_moments_data_to_binary(moments, t_params, io_initial_electron,
+                                              t_idx, r, z)
 
         closefile && close(io_initial_electron.fid)
     end
@@ -2311,7 +2428,8 @@ function debug_dump(vz::coordinate, vr::coordinate, vzeta::coordinate, vpa::coor
                                                           r, z, false,
                                                           external_source_settings,
                                                           evolve_density, evolve_upar,
-                                                          evolve_ppar)
+                                                          evolve_ppar,
+                                                          composition.electron_physics == kinetic_electrons)
             io_dfns = define_dynamic_dfn_variables!(
                 fid, r, z, vperp, vpa, vzeta, vr, vz, composition.n_ion_species,
                 composition.n_neutral_species, false, external_source_settings,
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index b5df32128..6d3880156 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -225,7 +225,7 @@ end
 function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                                vperp_spectral, vpa_spectral, collisions,
-                               external_source_settings, scratch_dummy, scratch, t_input,
+                               external_source_settings, scratch_dummy, scratch, t_params,
                                num_diss_params, advection_structs, io_input, input_dict;
                                restart_from_Boltzmann_electrons=false)
     
@@ -329,12 +329,12 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     end
 
     # initialize the electron pdf that satisfies the electron kinetic equation
-    initialize_electron_pdf!(scratch[1], pdf, moments, fields.phi, r, z, vpa, vperp,
+    initialize_electron_pdf!(scratch, pdf, moments, fields.phi, r, z, vpa, vperp,
                              vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
-                             num_diss_params, t_input.dt, io_input, input_dict)
+                             num_diss_params, t_params.electron, io_input, input_dict)
 
     return nothing
 end
@@ -425,11 +425,11 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
     return nothing
 end
 
-function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
+function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
                                   vz, z_spectral, vperp_spectral, vpa_spectral, z_advect,
                                   vpa_advect, scratch_dummy, collisions, composition,
-                                  geometry, external_source_settings, num_diss_params, dt,
-                                  io_input, input_dict)
+                                  geometry, external_source_settings, num_diss_params,
+                                  t_params, io_input, input_dict)
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
@@ -454,7 +454,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
 
             # Reload pdf and moments from an existing output file
             code_time, previous_runs_info, restart_time_index =
-                reload_electron_data!(pdf, moments, backup_prefix_iblock, -1,
+                reload_electron_data!(pdf, moments, t_params, backup_prefix_iblock, -1,
                                       geometry, r, z, vpa, vperp, vzeta, vr, vz)
 
             # Broadcast code_time from the root process of each shared-memory block (on which it
@@ -475,8 +475,8 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                                             restart_time_index,
                                                             previous_runs_info)
 
-            # update the electron pdf in the fvec struct
-            fvec.pdf_electron .= pdf.electron.norm
+            # update the electron pdf in the first scratch
+            scratch[1].pdf_electron .= pdf.electron.norm
         end
 
         @loop_r_z ir iz begin
@@ -501,7 +501,7 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
         electron_pseudotime, n_debug_outputs =
-            @views update_electron_pdf!(fvec, pdf.electron.norm, moments,
+            @views update_electron_pdf!(scratch, pdf.electron.norm, moments,
                                         moments.electron.dens, moments.electron.vth,
                                         moments.electron.ppar, moments.electron.qpar,
                                         moments.electron.qpar_updated, phi,
@@ -510,17 +510,18 @@ function initialize_electron_pdf!(fvec, pdf, moments, phi, r, z, vpa, vperp, vze
                                         moments.electron.dqpar_dz,
                                         moments.electron.dvth_dz, r, z, vperp, vpa,
                                         z_spectral, vperp_spectral, vpa_spectral,
-                                        z_advect, vpa_advect, scratch_dummy, dt,
+                                        z_advect, vpa_advect, scratch_dummy, t_params,
                                         collisions, composition, num_diss_params,
                                         max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
-                                        initial_time=code_time)
+                                        initial_time=code_time, evolve_ppar=true)
 
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
         t_idx = n_debug_outputs+1
-        write_initial_electron_state(pdf.electron.norm, moments, electron_pseudotime,
-                                     io_initial_electron, t_idx, r, z, vperp, vpa)
+        write_initial_electron_state(pdf.electron.norm, moments, t_params,
+                                     electron_pseudotime, io_initial_electron, t_idx, r,
+                                     z, vperp, vpa)
         finish_initial_electron_io(io_initial_electron)
 
     end
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 982575bfd..75c05697f 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -38,7 +38,7 @@ end
 an option but known at compile time when a `time_info` struct is passed as a function
 argument.
 """
-struct time_info{Terrorsum <: Real}
+struct time_info{Terrorsum <: Real, T_electron}
     nstep::mk_int
     end_time::mk_float
     dt::MPISharedArray{mk_float,1}
@@ -75,6 +75,7 @@ struct time_info{Terrorsum <: Real}
     converged_residual_value::mk_float
     use_manufactured_solns_for_advance::Bool
     stopfile::String
+    electron::T_electron
 end
 
 """
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 55d18bbda..9e3c18eca 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -573,6 +573,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
     code_time = 0.0
     dt = nothing
     dt_before_last_fail = nothing
+    electron_dt = nothing
+    electron_dt_before_last_fail = nothing
     previous_runs_info = nothing
     restart_had_kinetic_electrons = false
     begin_serial_region()
@@ -854,19 +856,27 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                 dt_before_last_fail = load_slice(dynamic, "dt_before_last_fail",
                                                 time_index)
             end
+            if "electron_dt" ∈ keys(dynamic)
+                electron_dt = load_slice(dynamic, "electron_dt", time_index)
+            end
+            if "electron_dt_before_last_fail" ∈ keys(dynamic)
+                electron_dt_before_last_fail =
+                    load_slice(dynamic, "electron_dt_before_last_fail", time_index)
+            end
         finally
             close(fid)
         end
     end
 
-    return code_time, dt, dt_before_last_fail, previous_runs_info, time_index, restart_had_kinetic_electrons
+    return code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
+           previous_runs_info, time_index, restart_had_kinetic_electrons
 end
 
 """
 Reload electron pdf and moments from an existing output file.
 """
-function reload_electron_data!(pdf, moments, restart_prefix_iblock, time_index, geometry,
-                               r, z, vpa, vperp, vzeta, vr, vz)
+function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, time_index,
+                               geometry, r, z, vpa, vperp, vzeta, vr, vz)
     code_time = 0.0
     previous_runs_info = nothing
     begin_serial_region()
@@ -959,6 +969,11 @@ function reload_electron_data!(pdf, moments, restart_prefix_iblock, time_index,
                                     restart_evolve_density, restart_evolve_upar,
                                     restart_evolve_ppar)
 
+            t_params.dt[] = load_slice(dynamic, "electron_dt", time_index)
+            t_params.previous_dt[] = t_params.dt[]
+            t_params.dt_before_output[] = t_params.dt[]
+            t_params.dt_before_last_fail[] =
+                load_slice(dynamic, "electron_dt_before_last_fail", time_index)
         finally
             close(fid)
         end
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 8bba7e0ff..4abe967eb 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -59,7 +59,6 @@ include("boundary_conditions.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("krook_collisions.jl")
-include("electron_kinetic_equation.jl")
 include("continuity.jl")
 include("energy_equation.jl")
 include("force_balance.jl")
@@ -68,9 +67,10 @@ include("numerical_dissipation.jl")
 include("moment_kinetics_input.jl")
 include("utils.jl")
 include("load_data.jl")
+include("analysis.jl")
+include("electron_kinetic_equation.jl")
 include("initial_conditions.jl")
 include("parameter_scans.jl")
-include("analysis.jl")
 include("time_advance.jl")
 
 using TimerOutputs
@@ -286,6 +286,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         code_time = 0.
         dt = nothing
         dt_before_last_fail = nothing
+        electron_dt = nothing
+        electron_dt_before_last_fail = nothing
         previous_runs_info = nothing
         restart_had_kinetic_electrons = false
     else
@@ -301,8 +303,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
                                                                         io_input.output_dir)
 
         # Reload pdf and moments from an existing output file
-        code_time, dt, dt_before_last_fail, previous_runs_info, restart_time_index,
-        restart_had_kinetic_electrons =
+        code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
+        previous_runs_info, restart_time_index, restart_had_kinetic_electrons =
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
@@ -328,10 +330,10 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral, z_spectral,
             r_spectral, composition, drive_input, moments, t_input, code_time, dt,
-            dt_before_last_fail, collisions, species, geometry, boundary_distributions,
-            external_source_settings, num_diss_params, manufactured_solns_input,
-            advection_structs, scratch_dummy, io_input, restarting,
-            restart_had_kinetic_electrons, input_dict)
+            dt_before_last_fail, electron_dt, electron_dt_before_last_fail, collisions,
+            species, geometry, boundary_distributions, external_source_settings,
+            num_diss_params, manufactured_solns_input, advection_structs, scratch_dummy,
+            io_input, restarting, restart_had_kinetic_electrons, input_dict)
 
     # This is the closest we can get to the end time of the setup before writing it to the
     # output file
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 40444a12e..246d3cced 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -223,6 +223,90 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     if timestepping_section["atol_upar"] === nothing
         timestepping_section["atol_upar"] = 1.0e-2 * timestepping_section["rtol"]
     end
+
+    # parameters related to electron time stepping
+    electron_timestepping_section = set_defaults_and_check_section!(
+        scan_input, "electron_timestepping";
+        nstep=50000,
+        dt=timestepping_section["dt"] * sqrt(composition.me_over_mi),
+        CFL_prefactor=timestepping_section["CFL_prefactor"],
+        nwrite=nothing,
+        nwrite_dfns=nothing,
+        type=timestepping_section["type"],
+        split_operators=false,
+        converged_residual_value=1.0e-3,
+        rtol=timestepping_section["rtol"],
+        atol=timestepping_section["atol"],
+        step_update_prefactor=timestepping_section["step_update_prefactor"],
+        max_increase_factor=timestepping_section["max_increase_factor"],
+        max_increase_factor_near_last_fail=timestepping_section["max_increase_factor_near_last_fail"],
+        last_fail_proximity_factor=timestepping_section["last_fail_proximity_factor"],
+        minimum_dt=timestepping_section["minimum_dt"] * sqrt(composition.me_over_mi),
+        maximum_dt=timestepping_section["maximum_dt"] * sqrt(composition.me_over_mi),
+        high_precision_error_sum=timestepping_section["high_precision_error_sum"],
+       )
+    if electron_timestepping_section["nwrite"] === nothing
+        electron_timestepping_section["nwrite"] = electron_timestepping_section["nstep"]
+    elseif electron_timestepping_section["nwrite"] > electron_timestepping_section["nstep"]
+        electron_timestepping_section["nwrite"] = electron_timestepping_section["nstep"]
+    end
+    if electron_timestepping_section["nwrite_dfns"] === nothing
+        electron_timestepping_section["nwrite_dfns"] = electron_timestepping_section["nstep"]
+    elseif electron_timestepping_section["nwrite_dfns"] > electron_timestepping_section["nstep"]
+        electron_timestepping_section["nwrite_dfns"] = electron_timestepping_section["nstep"]
+    end
+    # Make a copy because "stopfile_name" is not a separate input for the electrons, so we
+    # do not want to add a value to the `input_dict`. We also add a few dummy inputs that
+    # are not actually used for electrons.
+    electron_timestepping_section = copy(electron_timestepping_section)
+    electron_timestepping_section["stopfile_name"] = timestepping_section["stopfile_name"]
+    electron_timestepping_section["atol_upar"] = NaN
+    electron_timestepping_section["steady_state_residual"] = true
+    electron_timestepping_input = Dict_to_NamedTuple(electron_timestepping_section)
+    if !(0.0 < electron_timestepping_input.step_update_prefactor < 1.0)
+        error("[electron_timestepping] step_update_prefactor="
+              * "$(electron_timestepping_input.step_update_prefactor) must be between "
+              * "0.0 and 1.0.")
+    end
+    if electron_timestepping_input.max_increase_factor ≤ 1.0
+        error("[electron_timestepping] max_increase_factor="
+              * "$(electron_timestepping_input.max_increase_factor) must be greater than "
+              * "1.0.")
+    end
+    if electron_timestepping_input.max_increase_factor_near_last_fail ≤ 1.0
+        error("[electron_timestepping] max_increase_factor_near_last_fail="
+              * "$(electron_timestepping_input.max_increase_factor_near_last_fail) must "
+              * "be greater than 1.0.")
+    end
+    if !isinf(electron_timestepping_input.max_increase_factor_near_last_fail) &&
+            electron_timestepping_input.max_increase_factor_near_last_fail > electron_timestepping_input.max_increase_factor
+        error("[electron_timestepping] max_increase_factor_near_last_fail="
+              * "$(electron_timestepping_input.max_increase_factor_near_last_fail) should be "
+              * "less than max_increase_factor="
+              * "$(electron_timestepping_input.max_increase_factor).")
+    end
+    if electron_timestepping_input.last_fail_proximity_factor ≤ 1.0
+        error("[electron_timestepping] last_fail_proximity_factor="
+              * "$(electron_timestepping_input.last_fail_proximity_factor) must be "
+              * "greater than 1.0.")
+    end
+    if electron_timestepping_input.minimum_dt > electron_timestepping_input.maximum_dt
+        error("[electron_timestepping] minimum_dt="
+              * "$(electron_timestepping_input.minimum_dt) must be less than "
+              * "maximum_dt=$(electron_timestepping_input.maximum_dt)")
+    end
+    if electron_timestepping_input.maximum_dt ≤ 0.0
+        error("[electron_timestepping] maximum_dt="
+              * "$(electron_timestepping_input.maximum_dt) must be positive")
+    end
+
+    # Make a copy of `timestepping_section` here as we do not want to add
+    # `electron_timestepping_input` to the `input_dict` because there is already an
+    # "electron_timestepping" section containing the input info - we only want to put
+    # `electron_timestepping_input` into the Dict that is used to make
+    # `timestepping_input`, so that it becomes part of `timestepping_input`.
+    timestepping_section = copy(timestepping_section)
+    timestepping_section["electron_t_input"] = electron_timestepping_input
     timestepping_input = Dict_to_NamedTuple(timestepping_section)
     if !(0.0 < timestepping_input.step_update_prefactor < 1.0)
         error("step_update_prefactor=$(timestepping_input.step_update_prefactor) must "
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index aaf2e8437..e3e8fc5c8 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -264,25 +264,15 @@ function allocate_advection_structs(composition, z, r, vpa, vperp, vz, vr, vzeta
 end
 
 """
-create arrays and do other work needed to setup
-the main time advance loop.
-this includes creating and populating structs
-for Chebyshev transforms, velocity space moments,
-EM fields, and advection terms
-"""
-function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
-                             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
-                             z_spectral, r_spectral, composition, drive_input, moments,
-                             t_input, code_time, dt_reload, dt_before_last_fail_reload,
-                             collisions, species, geometry, boundary_distributions,
-                             external_source_settings, num_diss_params,
-                             manufactured_solns_input, advection_structs, scratch_dummy,
-                             io_input, restarting, restart_had_kinetic_electrons,
-                             input_dict)
-    # define some local variables for convenience/tidiness
-    n_ion_species = composition.n_ion_species
-    n_neutral_species = composition.n_neutral_species
+    setup_time_info(t_input; electrons=nothing)
 
+Create a [`input_structs.time_info`](@ref) struct using the settings in `t_input`.
+
+If something is passed in `electron`, it is stored in the `electron_t_params` member of
+the returned `time_info`.
+"""
+function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_reload,
+                         manufactured_solns_input; electron=nothing)
     dt_shared = allocate_shared_float(1)
     previous_dt_shared = allocate_shared_float(1)
     next_output_time = allocate_shared_float(1)
@@ -320,7 +310,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
     else
         error_sum_zero = 0.0
     end
-    t_params = time_info(t_input.nstep, end_time, dt_shared, previous_dt_shared, next_output_time,
+    return time_info(t_input.nstep, end_time, dt_shared, previous_dt_shared, next_output_time,
                          dt_before_output, dt_before_last_fail, CFL_prefactor,
                          step_to_output, Ref(0), Ref(0), mk_int[], mk_int[],
                          moments_output_times, dfns_output_times, t_input.type, rk_coefs,
@@ -331,7 +321,56 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                          t_input.last_fail_proximity_factor, t_input.minimum_dt,
                          t_input.maximum_dt, error_sum_zero, t_input.split_operators,
                          t_input.steady_state_residual, t_input.converged_residual_value,
-                         manufactured_solns_input.use_for_advance, t_input.stopfile_name)
+                         manufactured_solns_input.use_for_advance, t_input.stopfile_name,
+                         electron)
+end
+
+"""
+create arrays and do other work needed to setup
+the main time advance loop.
+this includes creating and populating structs
+for Chebyshev transforms, velocity space moments,
+EM fields, and advection terms
+"""
+function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_spectral,
+                             vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
+                             z_spectral, r_spectral, composition, drive_input, moments,
+                             t_input, code_time, dt_reload, dt_before_last_fail_reload,
+                             electron_dt_reload, electron_dt_before_last_fail_reload,
+                             collisions, species, geometry, boundary_distributions,
+                             external_source_settings, num_diss_params,
+                             manufactured_solns_input, advection_structs, scratch_dummy,
+                             io_input, restarting, restart_had_kinetic_electrons,
+                             input_dict)
+    # define some local variables for convenience/tidiness
+    n_ion_species = composition.n_ion_species
+    n_neutral_species = composition.n_neutral_species
+
+    if composition.electron_physics == kinetic_electrons
+        electron_t_params = setup_time_info(t_input.electron_t_input, 0.0,
+                                            electron_dt_reload,
+                                            electron_dt_before_last_fail_reload,
+                                            manufactured_solns_input)
+        # Make Vectors that count which variable caused timestep limits and timestep failures
+        # the right length. Do this setup even when not using adaptive timestepping, because
+        # it is easier than modifying the file I/O according to whether we are using adaptive
+        # timestepping.
+        #
+        # Entries for limit by accuracy (which is an average over all variables),
+        # max_increase_factor, minimum_dt and maximum_dt
+        push!(electron_t_params.limit_caused_by, 0, 0, 0, 0, 0)
+
+        # electron pdf
+        push!(electron_t_params.limit_caused_by, 0, 0)
+        push!(electron_t_params.failure_caused_by, 0)
+
+        # electron ppar
+        push!(electron_t_params.failure_caused_by, 0)
+    else
+        electron_t_params = nothing
+    end
+    t_params = setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_reload,
+                               manufactured_solns_input; electron=electron_t_params)
 
     # Make Vectors that count which variable caused timestep limits and timestep failures
     # the right length. Do this setup even when not using adaptive timestepping, because
@@ -388,6 +427,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
     # that may be evolved separately via fluid equations
+    n_rk_stages = t_params.n_rk_stages
+    if t_params.electron !== nothing
+        n_rk_stages = max(n_rk_stages, t_params.electron.n_rk_stages)
+    end
     scratch = setup_scratch_arrays(moments, pdf, n_rk_stages)
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
@@ -406,7 +449,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions,
                               external_source_settings, scratch_dummy, scratch, t_params,
-                              num_diss_params, advection_structs, io_input, input_dict)
+                              num_diss_params, advection_structs, io_input, input_dict;
+                              restart_from_Boltzmann_electrons=restarting)
     end
 
     # update the derivatives of the electron moments as these may be needed when
@@ -1177,7 +1221,7 @@ function time_advance!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyr
 
         if write_moments || write_dfns || finish_now
             # update the diagnostic chodura condition
-            update_chodura!(moments,scratch[end].pdf,vpa,vperp,z,r,spectral_objects.r_spectral,composition,geometry,scratch_dummy,advect_objects.z_advect)
+            update_chodura!(moments,scratch[t_params.n_rk_stages+1].pdf,vpa,vperp,z,r,spectral_objects.r_spectral,composition,geometry,scratch_dummy,advect_objects.z_advect)
 
             # Always synchronise here, regardless of if we changed region or not
             begin_serial_region(no_synchronize=true)
@@ -1246,7 +1290,7 @@ function time_advance!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyr
                 all_residuals = Vector{mk_float}()
                 @loop_s is begin
                     @views residual_ni =
-                        steady_state_residuals(scratch[end].density[:,:,is],
+                        steady_state_residuals(scratch[t_params.n_rk_stages+1].density[:,:,is],
                                                scratch[1].density[:,:,is], t_params.previous_dt[];
                                                use_mpi=true, only_max_abs=true)
                     if global_rank[] == 0
@@ -1259,7 +1303,7 @@ function time_advance!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyr
                 if composition.n_neutral_species > 0
                     @loop_sn isn begin
                         residual_nn =
-                            steady_state_residuals(scratch[end].density_neutral[:,:,isn],
+                            steady_state_residuals(scratch[t_params.n_rk_stages+1].density_neutral[:,:,isn],
                                                    scratch[1].density_neutral[:,:,isn],
                                                    t_params.previous_dt[]; use_mpi=true,
                                                    only_max_abs=true)
@@ -1727,7 +1771,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     update_phi!(fields, scratch[istage+1], z, r, composition, collisions, moments,
                 z_spectral, r_spectral, scratch_dummy)
     if !(( moments.evolve_upar || moments.evolve_ppar) &&
-              istage == length(scratch)-1)
+              istage == t_params.n_rk_stages)
         # _block_synchronize() here because phi needs to be read on different ranks than
         # it was written on, even though the loop-type does not change here. However,
         # after the final RK stage can skip if:
@@ -1803,7 +1847,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
             update_phi!(fields, scratch[istage+1], z, r, composition, collisions, moments,
                         z_spectral, r_spectral, scratch_dummy)
             if !(( moments.evolve_upar || moments.evolve_ppar) &&
-                      istage == length(scratch)-1)
+                      istage == t_params.n_rk_stages)
                 # _block_synchronize() here because phi needs to be read on different ranks than
                 # it was written on, even though the loop-type does not change here. However,
                 # after the final RK stage can skip if:
@@ -1828,7 +1872,7 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
     error_norm_method = "L2"
 
     error_coeffs = t_params.rk_coefs[:,end]
-    if length(scratch) < 3
+    if t_params.n_rk_stages < 3
         # This should never happen as an adaptive RK scheme needs at least 2 RHS evals so
         # (with the pre-timestep data) there must be at least 3 entries in `scratch`.
         error("adaptive timestep needs a buffer scratch array")
@@ -1872,7 +1916,7 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
 
     # ion vpa-advection
     ion_vpa_CFL = Inf
-    update_speed_vpa!(vpa_advect, fields, scratch[end], moments, vpa, vperp, z, r,
+    update_speed_vpa!(vpa_advect, fields, scratch[t_params.n_rk_stages+1], moments, vpa, vperp, z, r,
                       composition, collisions, external_source_settings.ion, t,
                       geometry)
     @loop_s is begin
@@ -1892,9 +1936,10 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
     # Calculate error for ion distribution functions
     # Note rk_error_variable!() stores the calculated error in `scratch[2]`.
     rk_error_variable!(scratch, :pdf, t_params)
-    ion_pdf_error = local_error_norm(scratch[2].pdf, scratch[end].pdf, t_params.rtol,
-                                     t_params.atol; method=error_norm_method,
-                                     skip_r_inner=skip_r_inner, skip_z_lower=skip_z_lower,
+    ion_pdf_error = local_error_norm(scratch[2].pdf, scratch[t_params.n_rk_stages+1].pdf,
+                                     t_params.rtol, t_params.atol;
+                                     method=error_norm_method, skip_r_inner=skip_r_inner,
+                                     skip_z_lower=skip_z_lower,
                                      error_sum_zero=t_params.error_sum_zero)
     push!(error_norms, ion_pdf_error)
     push!(total_points,
@@ -1904,7 +1949,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
     if moments.evolve_density
         begin_s_r_z_region()
         rk_error_variable!(scratch, :density, t_params)
-        ion_n_err = local_error_norm(scratch[2].density, scratch[end].density,
+        ion_n_err = local_error_norm(scratch[2].density,
+                                     scratch[t_params.n_rk_stages+1].density,
                                      t_params.rtol, t_params.atol;
                                      method=error_norm_method, skip_r_inner=skip_r_inner,
                                      skip_z_lower=skip_z_lower,
@@ -1915,7 +1961,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
     if moments.evolve_upar
         begin_s_r_z_region()
         rk_error_variable!(scratch, :upar, t_params)
-        ion_u_err = local_error_norm(scratch[2].upar, scratch[end].upar, t_params.rtol,
+        ion_u_err = local_error_norm(scratch[2].upar,
+                                     scratch[t_params.n_rk_stages+1].upar, t_params.rtol,
                                      t_params.atol; method=error_norm_method,
                                      skip_r_inner=skip_r_inner, skip_z_lower=skip_z_lower,
                                      error_sum_zero=t_params.error_sum_zero)
@@ -1925,7 +1972,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
     if moments.evolve_ppar
         begin_s_r_z_region()
         rk_error_variable!(scratch, :ppar, t_params)
-        ion_p_err = local_error_norm(scratch[2].ppar, scratch[end].ppar, t_params.rtol,
+        ion_p_err = local_error_norm(scratch[2].ppar,
+                                     scratch[t_params.n_rk_stages+1].ppar, t_params.rtol,
                                      t_params.atol; method=error_norm_method,
                                      skip_r_inner=skip_r_inner, skip_z_lower=skip_z_lower,
                                      error_sum_zero=t_params.error_sum_zero)
@@ -1953,9 +2001,10 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
 
         # neutral vz-advection
         neutral_vz_CFL = Inf
-        update_speed_neutral_vz!(neutral_vz_advect, fields, scratch[end],
-                                 moments, vz, vr, vzeta, z, r, composition,
-                                 collisions, external_source_settings.neutral)
+        update_speed_neutral_vz!(neutral_vz_advect, fields,
+                                 scratch[t_params.n_rk_stages+1], moments, vz, vr, vzeta,
+                                 z, r, composition, collisions,
+                                 external_source_settings.neutral)
         @loop_sn isn begin
             this_minimum = get_minimum_CFL_neutral_vz(neutral_vz_advect[isn].speed, vz)
             @serial_region begin
@@ -1967,8 +2016,9 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
         # Calculate error for neutral distribution functions
         rk_error_variable!(scratch, :pdf_neutral, t_params; neutrals=true)
         neut_pdf_error = local_error_norm(scratch[2].pdf_neutral,
-                                          scratch[end].pdf_neutral, t_params.rtol,
-                                          t_params.atol; method=error_norm_method,
+                                          scratch[t_params.n_rk_stages+1].pdf_neutral,
+                                          t_params.rtol, t_params.atol;
+                                          method=error_norm_method,
                                           skip_r_inner=skip_r_inner,
                                           skip_z_lower=skip_z_lower,
                                           error_sum_zero=t_params.error_sum_zero)
@@ -1982,8 +2032,9 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
             begin_sn_r_z_region()
             rk_error_variable!(scratch, :density_neutral, t_params; neutrals=true)
             neut_n_err = local_error_norm(scratch[2].density_neutral,
-                                          scratch[end].density, t_params.rtol,
-                                          t_params.atol; method=error_norm_method,
+                                          scratch[t_params.n_rk_stages+1].density,
+                                          t_params.rtol, t_params.atol;
+                                          method=error_norm_method,
                                           skip_r_inner=skip_r_inner,
                                           skip_z_lower=skip_z_lower,
                                           error_sum_zero=t_params.error_sum_zero)
@@ -1993,7 +2044,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
         if moments.evolve_upar
             begin_s_r_z_region()
             rk_error_variable!(scratch, :uz_neutral, t_params; neutrals=true)
-            neut_u_err = local_error_norm(scratch[2].uz_neutral, scratch[end].uz_neutral,
+            neut_u_err = local_error_norm(scratch[2].uz_neutral,
+                                          scratch[t_params.n_rk_stages+1].uz_neutral,
                                           t_params.rtol, t_params.atol;
                                           method=error_norm_method,
                                           skip_r_inner=skip_r_inner,
@@ -2005,7 +2057,8 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
         if moments.evolve_ppar
             begin_s_r_z_region()
             rk_error_variable!(scratch, :pz_neutral, t_params; neutrals=true)
-            neut_p_err = local_error_norm(scratch[2].pz_neutral, scratch[end].pz_neutral,
+            neut_p_err = local_error_norm(scratch[2].pz_neutral,
+                                          scratch[t_params.n_rk_stages+1].pz_neutral,
                                           t_params.rtol, t_params.atol;
                                           method=error_norm_method,
                                           skip_r_inner=skip_r_inner,

From 2702ab27c3164aca51fd7dabf38880eb073537c6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 27 Mar 2024 13:48:48 +0000
Subject: [PATCH 177/428] Don't use upwinded derivative in
 electron_energy_equation!()

Parallel flow speeds are always small compared to the electron thermal
speed, so it seems more sensible not to use an upwinded derivative
depending on the sign of electron upar in the electron energy equation.
---
 moment_kinetics/src/electron_fluid_equations.jl  |  4 ++--
 moment_kinetics/src/electron_kinetic_equation.jl | 11 -----------
 2 files changed, 2 insertions(+), 13 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 968ddc923..a3cc159c4 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -127,12 +127,12 @@ function electron_energy_equation!(ppar, fvec, moments, collisions, dt, composit
     # calculate contribution to rhs of energy equation (formulated in terms of pressure)
     # arising from derivatives of ppar, qpar and upar
     @loop_r_z ir iz begin
-        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz_upwind[iz,ir]
+        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
                     + moments.dqpar_dz[iz,ir]
                     + 3*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
     end
     # @loop_r_z ir iz begin
-    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz_upwind[iz,ir]
+    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
     #                 + (2/3)*moments.dqpar_dz[iz,ir]
     #                 + (5/3)*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
     # end
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 9010ba294..02af52d2e 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -120,8 +120,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
     begin_r_z_region()
 
-    # create a (z,r) dimension dummy array for use in taking derivatives
-    dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
     # create several (r) dimension dummy arrays for use in taking derivatives
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
@@ -130,9 +128,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
     buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
 
-    @loop_r_z ir iz begin
-        dummy_zr[iz,ir] = -moments.electron.upar[iz,ir]
-    end
     # compute the z-derivative of the input electron parallel flow, needed for the electron kinetic equation
     @views derivative_z!(moments.electron.dupar_dz, moments.electron.upar, buffer_r_1,
                          buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
@@ -316,12 +311,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 if evolve_ppar
                     # get an updated iterate of the electron parallel pressure
                     begin_r_z_region()
-                    # Compute the upwinded z-derivative of the electron parallel pressure for the
-                    # electron energy equation
-                    @views derivative_z!(moments.electron.dppar_dz_upwind,
-                                         scratch[istage+1].electron_ppar, dummy_zr, buffer_r_1,
-                                         buffer_r_2, buffer_r_3, buffer_r_4, buffer_r_5,
-                                         buffer_r_6, z_spectral, z)
                     # compute the z-derivative of the updated electron parallel pressure
                     @views derivative_z!(moments.electron.dppar_dz,
                                          scratch[istage+1].electron_ppar, buffer_r_1, buffer_r_2,

From 13db8b1610b71141a9cb0cc6c4856667ad4f287f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 24 Mar 2024 22:07:30 +0000
Subject: [PATCH 178/428] Pass all moments to electron_energy_equation!()
 explicitly as arguments

When initializing the electron distribution function and parallel
pressure, want to evolve the energy equation using ion and neutral
moments from the `moments` struct, because `scratch` is being used for
the electron time advance, but when advancing the electron parallel
pressure along with the ion/neutral moments and kinetic equation, need
to get some of the ion and neutral moments from `scratch`. This commit
makes electron_energy_equation!() take all the moments as arguments, so
that the calling site decides which struct they come from.
---
 .../src/electron_fluid_equations.jl           | 56 ++++++++++---------
 .../src/electron_kinetic_equation.jl          |  8 ++-
 moment_kinetics/src/time_advance.jl           |  8 ++-
 3 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index a3cc159c4..7be4441d6 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -118,7 +118,9 @@ an explicit time advance.
 NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
-function electron_energy_equation!(ppar, fvec, moments, collisions, dt, composition,
+function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron_upar,
+                                   ion_upar, ion_ppar, density_neutral, uz_neutral,
+                                   pz_neutral, moments, collisions, dt, composition,
                                    num_diss_params, z)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
@@ -127,68 +129,68 @@ function electron_energy_equation!(ppar, fvec, moments, collisions, dt, composit
     # calculate contribution to rhs of energy equation (formulated in terms of pressure)
     # arising from derivatives of ppar, qpar and upar
     @loop_r_z ir iz begin
-        ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
-                    + moments.dqpar_dz[iz,ir]
-                    + 3*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
+        ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
+                               + moments.dqpar_dz[iz,ir]
+                               + 3*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
     end
     # @loop_r_z ir iz begin
-    #     ppar[iz,ir] -= dt*(fvec.electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
+    #     ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
     #                 + (2/3)*moments.dqpar_dz[iz,ir]
-    #                 + (5/3)*fvec.electron_ppar[iz,ir]*moments.dupar_dz[iz,ir])
+    #                 + (5/3)*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
     # end
     # compute the contribution to the rhs of the energy equation
     # arising from artificial diffusion
     diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_r_z ir iz begin
-            ppar[iz,ir] += dt*diffusion_coefficient*moments.d2ppar_dz2[iz,ir]
+            ppar_out[iz,ir] += dt*diffusion_coefficient*moments.d2ppar_dz2[iz,ir]
         end
     end
     # compute the contribution to the rhs of the energy equation
     # arising from electron-ion collisions
     if nu_ei > 0.0
         @loop_s_r_z is ir iz begin
-            ppar[iz,ir] += dt * (2 * me_over_mi * nu_ei * (fvec.ppar[iz,ir,is] - fvec.electron_ppar[iz,ir]))
-            ppar[iz,ir] += dt * ((2/3) * moments.parallel_friction[iz,ir]
-                                * (fvec.upar[iz,ir,is]-fvec.electron_upar[iz,ir]))
+            ppar_out[iz,ir] += dt * (2 * me_over_mi * nu_ei * (ion_ppar[iz,ir,is] - ppar_in[iz,ir]))
+            ppar_out[iz,ir] += dt * ((2/3) * moments.parallel_friction[iz,ir]
+                                     * (ion_upar[iz,ir,is]-electron_upar[iz,ir]))
         end
     end
     # add in contributions due to charge exchange/ionization collisions
     if composition.n_neutral_species > 0
         if abs(collisions.charge_exchange_electron) > 0.0
             @loop_sn_r_z isn ir iz begin
-                ppar[iz,ir] +=
+                ppar_out[iz,ir] +=
                     dt * me_over_mi * collisions.charge_exchange_electron * (
-                    2*(fvec.electron_density[iz,ir]*fvec.pz_neutral[iz,ir,isn] -
-                    fvec.density_neutral[iz,ir,isn]*fvec.electron_ppar[iz,ir]) +
-                    (2/3)*fvec.electron_density[iz,ir]*fvec.density_neutral[iz,ir,isn] *
-                    (fvec.uz_neutral[iz,ir,isn] - fvec.electron_upar[iz,ir])^2)
+                    2*(electron_density[iz,ir]*pz_neutral[iz,ir,isn] -
+                    density_neutral[iz,ir,isn]*ppar_in[iz,ir]) +
+                    (2/3)*electron_density[iz,ir]*density_neutral[iz,ir,isn] *
+                    (uz_neutral[iz,ir,isn] - electron_upar[iz,ir])^2)
             end
         end
         if abs(collisions.ionization_electron) > 0.0
             # @loop_s_r_z is ir iz begin
-            #     ppar[iz,ir] +=
-            #         dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,is] * (
-            #         fvec.electron_ppar[iz,ir] -
-            #         (2/3)*fvec.electron_density[iz,ir] * collisions.ionization_energy)
+            #     ppar_out[iz,ir] +=
+            #         dt * collisions.ionization_electron * density_neutral[iz,ir,is] * (
+            #         ppar_in[iz,ir] -
+            #         (2/3)*electron_density[iz,ir] * collisions.ionization_energy)
             # end
             @loop_sn_r_z isn ir iz begin
-                ppar[iz,ir] +=
-                    dt * collisions.ionization_electron * fvec.density_neutral[iz,ir,isn] * (
-                    fvec.electron_ppar[iz,ir] -
-                    fvec.electron_density[iz,ir] * collisions.ionization_energy)
+                ppar_out[iz,ir] +=
+                    dt * collisions.ionization_electron * density_neutral[iz,ir,isn] * (
+                    ppar_in[iz,ir] -
+                    electron_density[iz,ir] * collisions.ionization_energy)
             end
         end
     end
     # calculate the external electron heat source, if any
-    calculate_electron_heat_source!(moments.heat_source, fvec.electron_ppar, moments.dupar_dz,
-                                    fvec.density_neutral, collisions.ionization, collisions.ionization_energy,
-                                    fvec.electron_density, fvec.ppar, collisions.nu_ei, composition.me_over_mi,
+    calculate_electron_heat_source!(moments.heat_source, ppar_in, moments.dupar_dz,
+                                    density_neutral, collisions.ionization, collisions.ionization_energy,
+                                    electron_density, ion_ppar, collisions.nu_ei, composition.me_over_mi,
                                     composition.T_wall, z)
     # add the contribution from the electron heat source                                
     begin_r_z_region()
     @loop_r_z ir iz begin
-        ppar[iz,ir] += dt * moments.heat_source[iz,ir]
+        ppar_out[iz,ir] += dt * moments.heat_source[iz,ir]
     end
     return nothing
 end
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 02af52d2e..d559bc66e 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1474,8 +1474,12 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
     end
 
     if evolve_ppar
-        electron_energy_equation!(fvec_out.electron_ppar, fvec_in, moments.electron,
-                                  collisions, dt, composition, num_diss_params, z)
+        electron_energy_equation!(fvec_out.electron_ppar, fvec_in.electron_ppar,
+                                  moments.electron.dens, moments.electron.upar,
+                                  moments.ion.upar, moments.ion.ppar,
+                                  moments.neutral.dens, moments.neutral.uz,
+                                  moments.neutral.pz, moments.electron, collisions, dt,
+                                  composition, num_diss_params, z)
     end
 
     return nothing
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index e3e8fc5c8..3ff3888ce 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2473,8 +2473,12 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                  external_source_settings.neutral, num_diss_params)
     end
     if advance.electron_energy
-        electron_energy_equation!(fvec_out.electron_ppar, fvec_out.density, fvec_in, moments, collisions, dt,
-                                  composition, num_diss_params, z)
+        electron_energy_equation!(fvec_out.electron_ppar, fvec_in.electron_ppar,
+                                  fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
+                                  fvec_in.ppar, fvec_in.density_neutral,
+                                  fvec_in.uz_neutral, fvec_in.pz_neutral,
+                                  moments.electron, collisions, dt, composition,
+                                  num_diss_params, z)
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not

From 68fba4aeb7be9c72fde51ed53026e8ebb31c4f62 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 24 Mar 2024 22:10:37 +0000
Subject: [PATCH 179/428] Remove checks for NaN in electron pdf and ppar

These checks were not parallised, and hopefully are less useful when we
have adaptive timestepping.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index d559bc66e..c377678f6 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -484,7 +484,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
 
-        if electron_pdf_converged || any(isnan.(ppar)) || any(isnan.(pdf))
+        if electron_pdf_converged
             break
         end
         t_params.step_counter[] += 1

From 5cb4f242997c91b2f3b20cb1386f16110f9d9ec8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 11:48:10 +0000
Subject: [PATCH 180/428] Plot electron timestep diagnostics for kinetic
 electron simulations

...in `makie_post_process()`.
---
 .../makie_post_processing/src/makie_post_processing.jl        | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 01aaf558d..204b46ffe 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -252,6 +252,10 @@ function makie_post_process(run_dir::Union{String,Tuple},
     end
 
     timestep_diagnostics(run_info; plot_prefix=plot_prefix)
+    if any(ri.composition.electron_physics ==
+           moment_kinetics.input_structs.kinetic_electrons for ri ∈ run_info)
+        timestep_diagnostics(run_info; plot_prefix=plot_prefix, electron=true)
+    end
 
     do_steady_state_residuals = any(input_dict[v]["steady_state_residual"]
                                     for v ∈ moment_variable_list)

From bbf3cb7b1d8711662093c5259fc9b493b964c97b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 24 Mar 2024 15:59:39 +0000
Subject: [PATCH 181/428] Make copies of `test_input` Dicts in tests

Making a copy before using the `test_input` Dict ensures that the inputs
defined in the test scripts are never modified.
---
 moment_kinetics/test/Krook_collisions_tests.jl             | 4 ++++
 moment_kinetics/test/fokker_planck_time_evolution_tests.jl | 4 ++++
 moment_kinetics/test/harrisonthompson.jl                   | 4 ++++
 moment_kinetics/test/nonlinear_sound_wave_tests.jl         | 4 ++++
 moment_kinetics/test/recycling_fraction_tests.jl           | 4 ++++
 moment_kinetics/test/restart_interpolation_tests.jl        | 4 ++++
 moment_kinetics/test/sound_wave_tests.jl                   | 4 ++++
 moment_kinetics/test/wall_bc_tests.jl                      | 4 ++++
 8 files changed, 32 insertions(+)

diff --git a/moment_kinetics/test/Krook_collisions_tests.jl b/moment_kinetics/test/Krook_collisions_tests.jl
index cadde203c..6552686f6 100644
--- a/moment_kinetics/test/Krook_collisions_tests.jl
+++ b/moment_kinetics/test/Krook_collisions_tests.jl
@@ -164,6 +164,10 @@ function run_test(test_input, rtol, atol; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     # Convert keyword arguments to a unique name
     name = test_input["run_name"]
     if length(args) > 0
diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index 708f41939..039d8a05b 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -179,6 +179,10 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     if upar_rtol === nothing
         upar_rtol = rtol
     end
diff --git a/moment_kinetics/test/harrisonthompson.jl b/moment_kinetics/test/harrisonthompson.jl
index 69e9db05d..310f5e84e 100644
--- a/moment_kinetics/test/harrisonthompson.jl
+++ b/moment_kinetics/test/harrisonthompson.jl
@@ -154,6 +154,10 @@ function run_test(test_input, analytic_rtol, analytic_atol, expected_phi,
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     # Convert keyword arguments to a unique name
     name = test_input["run_name"]
     if length(args) > 0
diff --git a/moment_kinetics/test/nonlinear_sound_wave_tests.jl b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
index 52ef17409..75a8e80a1 100644
--- a/moment_kinetics/test/nonlinear_sound_wave_tests.jl
+++ b/moment_kinetics/test/nonlinear_sound_wave_tests.jl
@@ -27,6 +27,10 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     if upar_rtol === nothing
         upar_rtol = rtol
     end
diff --git a/moment_kinetics/test/recycling_fraction_tests.jl b/moment_kinetics/test/recycling_fraction_tests.jl
index 37e911353..c7fee4666 100644
--- a/moment_kinetics/test/recycling_fraction_tests.jl
+++ b/moment_kinetics/test/recycling_fraction_tests.jl
@@ -167,6 +167,10 @@ function run_test(test_input, expected_phi; rtol=4.e-14, atol=1.e-15, args...)
     # by passing keyword arguments to run_test, args becomes a Tuple of Pairs which can be
     # used to update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     # Convert keyword arguments to a unique name
     name = test_input["run_name"]
     if length(args) > 0
diff --git a/moment_kinetics/test/restart_interpolation_tests.jl b/moment_kinetics/test/restart_interpolation_tests.jl
index 1380524f8..e6d161aeb 100644
--- a/moment_kinetics/test/restart_interpolation_tests.jl
+++ b/moment_kinetics/test/restart_interpolation_tests.jl
@@ -70,6 +70,10 @@ function run_test(test_input, base, message, rtol, atol; tol_3V, kwargs...)
     # by passing keyword arguments to run_test, kwargs becomes a Tuple of Pairs which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     if tol_3V === nothing
         atol_3V = atol
         rtol_3V = rtol
diff --git a/moment_kinetics/test/sound_wave_tests.jl b/moment_kinetics/test/sound_wave_tests.jl
index ea6275e63..f72cdf1c7 100644
--- a/moment_kinetics/test/sound_wave_tests.jl
+++ b/moment_kinetics/test/sound_wave_tests.jl
@@ -136,6 +136,10 @@ function run_test(test_input, analytic_frequency, analytic_growth_rate,
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     # Convert keyword arguments to a unique name
     name = test_input["run_name"]
     shortname = name
diff --git a/moment_kinetics/test/wall_bc_tests.jl b/moment_kinetics/test/wall_bc_tests.jl
index 175251802..6c46e294e 100644
--- a/moment_kinetics/test/wall_bc_tests.jl
+++ b/moment_kinetics/test/wall_bc_tests.jl
@@ -140,6 +140,10 @@ function run_test(test_input, expected_phi, tolerance; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
     # Convert keyword arguments to a unique name
     name = test_input["run_name"] * ", with element spacing: " * test_input["z_element_spacing_option"]
     if length(args) > 0

From fc6293fc0b8a458e6c99de7d8ede6d562401c795 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 16:35:53 +0000
Subject: [PATCH 182/428] Fix typo in external_sources docs

---
 docs/src/external_sources_notes.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/external_sources_notes.md b/docs/src/external_sources_notes.md
index f805d1552..bdaadeffc 100644
--- a/docs/src/external_sources_notes.md
+++ b/docs/src/external_sources_notes.md
@@ -14,8 +14,8 @@ S_n &= A_n(r,z) \frac{1}{(\pi)^{3/2} (2 T_{\mathrm{source},n} / m_n)^{3/2}} \exp
 or in 1V simulations that do not include $v_\perp$, $v_\zeta$, $v_r$ dimensions
 ```math
 \begin{align}
-S_i &= A_i(r,z) \frac{1}{sqrt{\pi} \sqrt{2 T_{\mathrm{source},i} / m_i}} \exp\left( -\frac{v_\perp^2}{T_{\mathrm{source},i}} \right) \\
-S_n &= A_n(r,z) \frac{1}{sqrt{\pi} \sqrt{2 T_{\mathrm{source},n} / m_n}} \exp\left( -\frac{v_z^2}{T_{\mathrm{source},n}} \right)
+S_i &= A_i(r,z) \frac{1}{\sqrt{\pi} \sqrt{2 T_{\mathrm{source},i} / m_i}} \exp\left( -\frac{v_\perp^2}{T_{\mathrm{source},i}} \right) \\
+S_n &= A_n(r,z) \frac{1}{\sqrt{\pi} \sqrt{2 T_{\mathrm{source},n} / m_n}} \exp\left( -\frac{v_z^2}{T_{\mathrm{source},n}} \right)
 \end{align}
 ```
 

From fd6ec3df0e56d68775b12d229460994ae9736e11 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 16:59:12 +0000
Subject: [PATCH 183/428] Only use ppar from fvec_in in electron KE when
 evolving ppar

...otherwise use `moments.electron.ppar` as `electron_ppar` is not
updated in the `scratch` objects by the electron kinetic equation (KE)
iteration.
---
 .../src/electron_kinetic_equation.jl          | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index c377678f6..233463cd7 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1440,24 +1440,27 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
                                                  vpa_advect, scratch_dummy, collisions,
                                                  composition, num_diss_params, dt;
                                                  evolve_ppar=false)
+    if evolve_ppar
+        ppar = fvec_in.electron_ppar
+    else
+        ppar = moments.electron.ppar
+    end
     # add the contribution from the z advection term
     electron_z_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
                           moments.electron.upar, moments.electron.vth, z_advect, z,
                           vpa.grid, z_spectral, scratch_dummy, dt)
 
     # add the contribution from the wpa advection term
-    electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
-                            fvec_in.electron_ppar, moments.electron.vth,
-                            moments.electron.dppar_dz, moments.electron.dqpar_dz,
-                            moments.electron.dvth_dz, vpa_advect, vpa, vpa_spectral,
-                            scratch_dummy, dt)
+    electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron, ppar,
+                            moments.electron.vth, moments.electron.dppar_dz,
+                            moments.electron.dqpar_dz, moments.electron.dvth_dz,
+                            vpa_advect, vpa, vpa_spectral, scratch_dummy, dt)
 
     # add in the contribution to the residual from the term proportional to the pdf
-    add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron,
-                                    fvec_in.electron_ppar, moments.electron.vth,
-                                    moments.electron.dens, moments.electron.ddens_dz,
-                                    moments.electron.dvth_dz, moments.electron.dqpar_dz,
-                                    vpa.grid, z, dt)
+    add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, ppar,
+                                    moments.electron.vth, moments.electron.dens,
+                                    moments.electron.ddens_dz, moments.electron.dvth_dz,
+                                    moments.electron.dqpar_dz, vpa.grid, z, dt)
 
     # add in numerical dissipation terms
     add_dissipation_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, scratch_dummy,

From 3dcfc6acb84ce08d4f1a85a056c3d5afc829d982 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 17:20:09 +0000
Subject: [PATCH 184/428] Relabel 'residual' to 'pdf_out' in
 add_contribution_from_pdf_term!()

---
 moment_kinetics/src/electron_kinetic_equation.jl | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 233463cd7..fd36f0234 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1880,7 +1880,7 @@ function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_
 end
 
 # add contribution to the residual coming from the term proporational to the pdf
-function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z, dt)
+function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z, dt)
     begin_r_z_vperp_vpa_region()
     @loop_r_z ir iz begin
         this_dqpar_dz = dqpar_dz[iz,ir]
@@ -1891,12 +1891,11 @@ function add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_d
         this_dvth_dz = dvth_dz[iz,ir]
         this_vth = vth[iz,ir]
         @loop_vperp_vpa ivperp ivpa begin
-            residual[ivpa,ivperp,iz,ir] += dt * (-0.5 * this_dqpar_dz / this_ppar -
-                                                 vpa[ivpa] * this_vth *
-                                                 (this_ddens_dz / this_dens
-                                                  - this_dvth_dz / this_vth)) *
-                                           pdf[ivpa,ivperp,iz,ir]
-            #residual[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf[ivpa, ivperp, :, :]
+            pdf_out[ivpa,ivperp,iz,ir] +=
+                dt * (-0.5 * this_dqpar_dz / this_ppar - vpa[ivpa] * this_vth *
+                      (this_ddens_dz / this_dens - this_dvth_dz / this_vth)) *
+                pdf_in[ivpa,ivperp,iz,ir]
+            #pdf_out[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf_in[ivpa, ivperp, :, :]
         end
     end
     return nothing

From 115194f0efb82a6cea46f122a5fba65ceb6f0d6d Mon Sep 17 00:00:00 2001
From: Lucas McConnell <89389250+LucasMontoya4@users.noreply.github.com>
Date: Thu, 28 Mar 2024 11:55:48 +0100
Subject: [PATCH 185/428] Separate dissipation parameters for ions, electrons,
 neutrals (#192)

---
 .gitignore                                    |   1 +
 ...lanck-1D2V-even_nz-shorttest-nstep200.toml |   2 +-
 examples/geometry/1D-mirror.toml              |   2 +-
 .../wall+sheath-bc_boltzmann_loworder.toml    |  10 +-
 .../wall+sheath-bc_kinetic.toml               |   2 +-
 ...wall+sheath-bc_kinetic_krook_loworder.toml |   2 +-
 .../wall+sheath-bc_kinetic_loworder.toml      |   2 +-
 .../num-diss-relaxation.toml                  |   2 +-
 examples/wall-bc/wall-bc_cheb.toml            |   6 +-
 examples/wall-bc/wall-bc_cheb_split1.toml     |   6 +-
 examples/wall-bc/wall-bc_cheb_split2.toml     |   6 +-
 examples/wall-bc/wall-bc_cheb_split3.toml     |   6 +-
 examples/wall-bc/wall-bc_no-neutrals.toml     |   2 +-
 .../wall-bc/wall-bc_no-neutrals_split1.toml   |   2 +-
 .../wall-bc/wall-bc_no-neutrals_split2.toml   |   2 +-
 .../wall-bc/wall-bc_no-neutrals_split3.toml   |   2 +-
 examples/wall-bc/wall-bc_volumerecycle.toml   |   8 +-
 .../wall-bc/wall-bc_volumerecycle_split1.toml |   8 +-
 .../wall-bc/wall-bc_volumerecycle_split2.toml |   8 +-
 .../wall-bc/wall-bc_volumerecycle_split3.toml |   8 +-
 .../restart_interpolation_inputs.jl           |   3 +-
 moment_kinetics/ext/manufactured_solns_ext.jl |  28 ++--
 moment_kinetics/src/continuity.jl             |   4 +-
 .../src/electron_fluid_equations.jl           |   2 +-
 .../src/electron_kinetic_equation.jl          |  12 +-
 moment_kinetics/src/energy_equation.jl        |   4 +-
 moment_kinetics/src/force_balance.jl          |   4 +-
 moment_kinetics/src/initial_conditions.jl     |  10 +-
 moment_kinetics/src/moment_kinetics_input.jl  |  13 +-
 moment_kinetics/src/numerical_dissipation.jl  | 140 ++++++++++++------
 moment_kinetics/src/time_advance.jl           |  68 +++++----
 moment_kinetics/src/velocity_moments.jl       |  44 +++---
 moment_kinetics/test/harrisonthompson.jl      |   2 +-
 .../test/recycling_fraction_tests.jl          |   3 +-
 .../wall-bc/wall-bc_recyclefraction0.5.toml   |   8 +-
 .../wall-bc_recyclefraction0.5_split1.toml    |   8 +-
 .../wall-bc_recyclefraction0.5_split2.toml    |   8 +-
 .../wall-bc_recyclefraction0.5_split3.toml    |   8 +-
 ..._new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml |   2 +-
 ...MS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml |   2 +-
 ...id_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml |   2 +-
 ...d_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml |   2 +-
 ...grid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml |   2 +-
 ...d_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml |   2 +-
 ...grid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml |   2 +-
 ..._new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml |   2 +-
 ...new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml |   2 +-
 ..._new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml |   2 +-
 ...d_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml |   2 +-
 ..._5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml |   2 +-
 ...grid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml |   2 +-
 ...grid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml |   2 +-
 ...MMS_nel_r_16_z_16_vpa_16_vperp_1_diss.toml |   2 +-
 ...l_MMS_nel_r_2_z_2_vpa_16_vperp_1_diss.toml |   2 +-
 ...MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml |   2 +-
 ...MS_nel_r_32_z_32_vpa_16_vperp_1_diss5.toml |   2 +-
 ...l_MMS_nel_r_4_z_4_vpa_16_vperp_1_diss.toml |   2 +-
 ...l_MMS_nel_r_8_z_8_vpa_16_vperp_1_diss.toml |   2 +-
 ...3_nel_r_1_z_1_vpa_6_vperp_3_fkpl_test.toml |   2 +-
 runs/wall+sheath-bc_boltzmann.toml            |  12 +-
 runs/wall+sheath-bc_braginskii.toml           |  13 +-
 ...l+sheath-bc_braginskii_boltzmann_test.toml |  10 +-
 runs/wall+sheath-bc_braginskii_colls.toml     |  10 +-
 util/compare_collision_frequencies.jl         |  12 +-
 64 files changed, 354 insertions(+), 201 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4d3aedd9f..a6eaf3e33 100644
--- a/.gitignore
+++ b/.gitignore
@@ -26,3 +26,4 @@ makie_postproc.so
 plots_postproc.so
 precompile-temp
 post_processing_input.toml
+.DS_Store
diff --git a/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml b/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
index 5163bfc67..01cdf06b1 100644
--- a/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
+++ b/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
@@ -88,7 +88,7 @@ vperp_discretization = "gausslegendre_pseudospectral"
 #vzeta_bc = "periodic"
 #vzeta_discretization = "chebyshev_pseudospectral"
 
-#[numerical_dissipation]
+#[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.0
 #vperp_dissipation_coefficient = 0.0
 #z_dissipation_coefficient = 0.1
diff --git a/examples/geometry/1D-mirror.toml b/examples/geometry/1D-mirror.toml
index 66c36fd5b..eb484f49d 100644
--- a/examples/geometry/1D-mirror.toml
+++ b/examples/geometry/1D-mirror.toml
@@ -71,7 +71,7 @@ vz_L = 18.0
 vz_bc = "both_zero"
 vz_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 vperp_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 force_minimum_pdf_value = 0.0
diff --git a/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
index aced83ff7..39d4f7043 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_boltzmann_loworder.toml
@@ -82,11 +82,15 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
-vpa_dissipation_coefficient = 0.002
-#vpa_dissipation_coefficient = 0.1
+#vpa_dissipation_coefficient = 0.002
+vpa_dissipation_coefficient = 0.1
 #vpa_dissipation_coefficient = 0.2
 #vpa_dissipation_coefficient = 2.0
 #vpa_dissipation_coefficient = 20.0
+
+[electron_numerical_dissipation]
+
+[neutral_numerical_dissipation]
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index cd65ce083..eab83014e 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -79,7 +79,7 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
index f0ae0e9c2..e81f88b8e 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
@@ -80,7 +80,7 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index bc30e359d..ea45ebf7f 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -82,7 +82,7 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
diff --git a/examples/numerical-dissipation/num-diss-relaxation.toml b/examples/numerical-dissipation/num-diss-relaxation.toml
index abd22e2da..ab953c0a2 100644
--- a/examples/numerical-dissipation/num-diss-relaxation.toml
+++ b/examples/numerical-dissipation/num-diss-relaxation.toml
@@ -58,7 +58,7 @@ vperp_L = 3.0
 vperp_bc = "zero"
 vperp_discretization = "gausslegendre_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.1
 vperp_dissipation_coefficient = 0.1
 z_dissipation_coefficient = -1.0
diff --git a/examples/wall-bc/wall-bc_cheb.toml b/examples/wall-bc/wall-bc_cheb.toml
index 118770700..0119de45c 100644
--- a/examples/wall-bc/wall-bc_cheb.toml
+++ b/examples/wall-bc/wall-bc_cheb.toml
@@ -65,6 +65,10 @@ vz_L = 18.0
 vz_bc = "both_zero"
 vz_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_cheb_split1.toml b/examples/wall-bc/wall-bc_cheb_split1.toml
index e47b10969..ea31a9017 100644
--- a/examples/wall-bc/wall-bc_cheb_split1.toml
+++ b/examples/wall-bc/wall-bc_cheb_split1.toml
@@ -66,6 +66,10 @@ vz_L = 18.0
 vz_bc = "both_zero"
 vz_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_cheb_split2.toml b/examples/wall-bc/wall-bc_cheb_split2.toml
index 173e9857b..897607283 100644
--- a/examples/wall-bc/wall-bc_cheb_split2.toml
+++ b/examples/wall-bc/wall-bc_cheb_split2.toml
@@ -66,6 +66,10 @@ vz_L = 18.0
 vz_bc = "both_zero"
 vz_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_cheb_split3.toml b/examples/wall-bc/wall-bc_cheb_split3.toml
index 999f0e39a..2e3c902aa 100644
--- a/examples/wall-bc/wall-bc_cheb_split3.toml
+++ b/examples/wall-bc/wall-bc_cheb_split3.toml
@@ -66,6 +66,10 @@ vz_L = 18.0
 vz_bc = "both_zero"
 vz_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-3 #1.0e-2 #1.0e-1
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_no-neutrals.toml b/examples/wall-bc/wall-bc_no-neutrals.toml
index 43af50cc9..ee79b7713 100644
--- a/examples/wall-bc/wall-bc_no-neutrals.toml
+++ b/examples/wall-bc/wall-bc_no-neutrals.toml
@@ -78,7 +78,7 @@ z_width = 0.125
 source_strength = 8.0
 source_T = 1.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
diff --git a/examples/wall-bc/wall-bc_no-neutrals_split1.toml b/examples/wall-bc/wall-bc_no-neutrals_split1.toml
index ab5c4c1d5..41c28fecf 100644
--- a/examples/wall-bc/wall-bc_no-neutrals_split1.toml
+++ b/examples/wall-bc/wall-bc_no-neutrals_split1.toml
@@ -78,7 +78,7 @@ z_width = 0.125
 source_strength = 8.0
 source_T = 1.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
diff --git a/examples/wall-bc/wall-bc_no-neutrals_split2.toml b/examples/wall-bc/wall-bc_no-neutrals_split2.toml
index 12732f289..346149f20 100644
--- a/examples/wall-bc/wall-bc_no-neutrals_split2.toml
+++ b/examples/wall-bc/wall-bc_no-neutrals_split2.toml
@@ -78,7 +78,7 @@ z_width = 0.125
 source_strength = 8.0
 source_T = 1.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
diff --git a/examples/wall-bc/wall-bc_no-neutrals_split3.toml b/examples/wall-bc/wall-bc_no-neutrals_split3.toml
index c966d47b2..179798a95 100644
--- a/examples/wall-bc/wall-bc_no-neutrals_split3.toml
+++ b/examples/wall-bc/wall-bc_no-neutrals_split3.toml
@@ -78,7 +78,7 @@ z_width = 0.125
 source_strength = 8.0
 source_T = 1.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
diff --git a/examples/wall-bc/wall-bc_volumerecycle.toml b/examples/wall-bc/wall-bc_volumerecycle.toml
index 86254f3a1..b61e0e062 100644
--- a/examples/wall-bc/wall-bc_volumerecycle.toml
+++ b/examples/wall-bc/wall-bc_volumerecycle.toml
@@ -85,8 +85,14 @@ active = true
 source_type = "recycling"
 recycling_controller_fraction = 0.25
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 #force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_volumerecycle_split1.toml b/examples/wall-bc/wall-bc_volumerecycle_split1.toml
index a1c5ad590..d3c1c1cbb 100644
--- a/examples/wall-bc/wall-bc_volumerecycle_split1.toml
+++ b/examples/wall-bc/wall-bc_volumerecycle_split1.toml
@@ -85,8 +85,14 @@ active = true
 source_type = "recycling"
 recycling_controller_fraction = 0.25
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 #force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+#force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_volumerecycle_split2.toml b/examples/wall-bc/wall-bc_volumerecycle_split2.toml
index 61c0dbf48..798a9f9cf 100644
--- a/examples/wall-bc/wall-bc_volumerecycle_split2.toml
+++ b/examples/wall-bc/wall-bc_volumerecycle_split2.toml
@@ -85,8 +85,14 @@ active = true
 source_type = "recycling"
 recycling_controller_fraction = 0.25
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 #force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+#force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/examples/wall-bc/wall-bc_volumerecycle_split3.toml b/examples/wall-bc/wall-bc_volumerecycle_split3.toml
index acb7c304c..2be7a0a18 100644
--- a/examples/wall-bc/wall-bc_volumerecycle_split3.toml
+++ b/examples/wall-bc/wall-bc_volumerecycle_split3.toml
@@ -85,8 +85,14 @@ active = true
 source_type = "recycling"
 recycling_controller_fraction = 0.25
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
\ No newline at end of file
diff --git a/moment_kinetics/debug_test/restart_interpolation_inputs.jl b/moment_kinetics/debug_test/restart_interpolation_inputs.jl
index a36c79042..a991f816c 100644
--- a/moment_kinetics/debug_test/restart_interpolation_inputs.jl
+++ b/moment_kinetics/debug_test/restart_interpolation_inputs.jl
@@ -60,7 +60,8 @@ base_input = Dict(
      "vzeta_nelement" => 1,
      "vr_ngrid" => 1,
      "vr_nelement" => 1,
-     "numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0))
+     "ion_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0),
+     "neutral_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0))
 
 test_input =
     merge(base_input,
diff --git a/moment_kinetics/ext/manufactured_solns_ext.jl b/moment_kinetics/ext/manufactured_solns_ext.jl
index b2b9f471f..cda63a819 100644
--- a/moment_kinetics/ext/manufactured_solns_ext.jl
+++ b/moment_kinetics/ext/manufactured_solns_ext.jl
@@ -638,31 +638,31 @@ using IfElse
             Si += - nuii_krook*(FMaxwellian - dfni)
         end
         include_num_diss_in_MMS = true
-        if num_diss_params.vpa_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Si += - num_diss_params.vpa_dissipation_coefficient*Dvpa(Dvpa(dfni))
+        if num_diss_params.ion.vpa_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Si += - num_diss_params.ion.vpa_dissipation_coefficient*Dvpa(Dvpa(dfni))
         end
-        if num_diss_params.vperp_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Si += - num_diss_params.vperp_dissipation_coefficient*Dvperp(Dvperp(dfni))
+        if num_diss_params.ion.vperp_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Si += - num_diss_params.ion.vperp_dissipation_coefficient*Dvperp(Dvperp(dfni))
         end
-        if num_diss_params.r_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Si += - rfac*num_diss_params.r_dissipation_coefficient*Dr(Dr(dfni))
+        if num_diss_params.ion.r_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Si += - rfac*num_diss_params.ion.r_dissipation_coefficient*Dr(Dr(dfni))
         end
-        if num_diss_params.z_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Si += - num_diss_params.z_dissipation_coefficient*Dz(Dz(dfni))
+        if num_diss_params.ion.z_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Si += - num_diss_params.ion.z_dissipation_coefficient*Dz(Dz(dfni))
         end
 
         Source_i = expand_derivatives(Si)
         
         # the neutral source to maintain the manufactured solution
         Sn = Dt(dfnn) + vz * Dz(dfnn) + rfac*vr * Dr(dfnn) + cx_frequency* (densi*dfnn - densn*vrvzvzeta_dfni) + ionization_frequency*dense*dfnn
-        if num_diss_params.vz_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Sn += - num_diss_params.vz_dissipation_coefficient*Dvz(Dvz(dfnn))
+        if num_diss_params.neutral.vz_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Sn += - num_diss_params.neutral.vz_dissipation_coefficient*Dvz(Dvz(dfnn))
         end
-        if num_diss_params.r_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Sn += - rfac*num_diss_params.r_dissipation_coefficient*Dr(Dr(dfnn))
+        if num_diss_params.neutral.r_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Sn += - rfac*num_diss_params.neutral.r_dissipation_coefficient*Dr(Dr(dfnn))
         end
-        if num_diss_params.z_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
-            Sn += - num_diss_params.z_dissipation_coefficient*Dz(Dz(dfnn))
+        if num_diss_params.neutral.z_dissipation_coefficient > 0.0 && include_num_diss_in_MMS
+            Sn += - num_diss_params.neutral.z_dissipation_coefficient*Dz(Dz(dfnn))
         end
         
         Source_n = expand_derivatives(Sn)
diff --git a/moment_kinetics/src/continuity.jl b/moment_kinetics/src/continuity.jl
index 1c3600f07..26e9e1314 100644
--- a/moment_kinetics/src/continuity.jl
+++ b/moment_kinetics/src/continuity.jl
@@ -39,7 +39,7 @@ function continuity_equation!(dens_out, fvec_in, moments, composition, dt, spect
     end
 
     # Ad-hoc diffusion to stabilise numerics...
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.ion.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
             dens_out[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2dens_dz2[iz,ir,is]
@@ -80,7 +80,7 @@ function neutral_continuity_equation!(dens_out, fvec_in, moments, composition, d
     end
 
     # Ad-hoc diffusion to stabilise numerics...
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.neutral.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_sn_r_z isn ir iz begin
             dens_out[iz,ir,isn] += dt*diffusion_coefficient*moments.neutral.d2dens_dz2[iz,ir,isn]
diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 9febf93f9..a65f15e7e 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -138,7 +138,7 @@ function electron_energy_equation!(ppar, dens_i, fvec, moments, collisions, dt,
     # end
     # compute the contribution to the rhs of the energy equation
     # arising from artificial diffusion
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.electron.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_r_z ir iz begin
             ppar[iz,ir] += dt*diffusion_coefficient*moments.electron.d2ppar_dz2[iz,ir]
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index f7eebb389..16fa21c0f 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -268,7 +268,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
             # centred second derivative for dissipation
-            if num_diss_params.moment_dissipation_coefficient > 0.0
+            if num_diss_params.electron.moment_dissipation_coefficient > 0.0
                 @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
                 begin_serial_region()
@@ -327,7 +327,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
                                                     vperp_spectral, vpa_spectral,
                                                     vpa_advect, moments,
-                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                    num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
                                                     composition.me_over_mi)
         #println("A pdf 1 ", pdf[:,1,1,1])
         #println("A pdf end ", pdf[:,1,end,1])
@@ -379,7 +379,7 @@ function update_electron_pdf_with_time_advance!(fvec, pdf, qpar, qpar_updated,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
          # centred second derivative for dissipation
-        if num_diss_params.moment_dissipation_coefficient > 0.0
+        if num_diss_params.electron.moment_dissipation_coefficient > 0.0
             @views derivative_z!(moments.electron.d2ppar_dz2, dppar_dz, buffer_r_1,
                                  buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         end
@@ -1442,15 +1442,15 @@ function add_dissipation_term!(residual, pdf, scratch_dummy, z_spectral, z, vpa,
     #                         buffer_r_4, z_spectral, z)
     #    @views derivative_z!(dummy_zr2, dummy_zr1, buffer_r_1, buffer_r_2, buffer_r_3,
     #                         buffer_r_4, z_spectral, z)
-    #    @. residual[ivpa,ivperp,:,:] -= num_diss_params.z_dissipation_coefficient * dummy_zr2
+    #    @. residual[ivpa,ivperp,:,:] -= num_diss_params.electron.z_dissipation_coefficient * dummy_zr2
     #end
     begin_r_z_vperp_region()
     @loop_r_z_vperp ir iz ivperp begin
         #@views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
         #@views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
-        #@. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
+        #@. residual[:,ivperp,iz,ir] -= num_diss_params.electron.vpa_dissipation_coefficient * vpa.scratch2
         @views second_derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, vpa_spectral)
-        @. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch
+        @. residual[:,ivperp,iz,ir] -= num_diss_params.electron.vpa_dissipation_coefficient * vpa.scratch
     end
     #stop()
     return nothing
diff --git a/moment_kinetics/src/energy_equation.jl b/moment_kinetics/src/energy_equation.jl
index 8bc058dc1..8890be3b6 100644
--- a/moment_kinetics/src/energy_equation.jl
+++ b/moment_kinetics/src/energy_equation.jl
@@ -30,7 +30,7 @@ function energy_equation!(ppar, fvec, moments, collisions, dt, spectral, composi
         end
     end
 
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.ion.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
             ppar[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2ppar_dz2[iz,ir,is]
@@ -82,7 +82,7 @@ function neutral_energy_equation!(pz, fvec, moments, collisions, dt, spectral,
         end
     end
 
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.neutral.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_sn_r_z isn ir iz begin
             pz[iz,ir,isn] += dt*diffusion_coefficient*moments.neutral.d2pz_dz2[iz,ir,isn]
diff --git a/moment_kinetics/src/force_balance.jl b/moment_kinetics/src/force_balance.jl
index 3040be002..902c3901e 100644
--- a/moment_kinetics/src/force_balance.jl
+++ b/moment_kinetics/src/force_balance.jl
@@ -38,7 +38,7 @@ function force_balance!(pflx, density_out, fvec, moments, fields, collisions, dt
     end
 
     # Ad-hoc diffusion to stabilise numerics...
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.ion.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_s_r_z is ir iz begin
             pflx[iz,ir,is] += dt*diffusion_coefficient*moments.ion.d2upar_dz2[iz,ir,is]*density[iz,ir,is]
@@ -92,7 +92,7 @@ function neutral_force_balance!(pflx, density_out, fvec, moments, fields, collis
     end
 
     # Ad-hoc diffusion to stabilise numerics...
-    diffusion_coefficient = num_diss_params.moment_dissipation_coefficient
+    diffusion_coefficient = num_diss_params.neutral.moment_dissipation_coefficient
     if diffusion_coefficient > 0.0
         @loop_sn_r_z isn ir iz begin
             pflx[iz,ir,isn] += dt*diffusion_coefficient*moments.neutral.d2uz_dz2[iz,ir,isn]*density[iz,ir,isn]
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index dfdb069c6..c7c378598 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -52,7 +52,7 @@ Creates the structs for the pdf and the velocity-space moments
 """
 function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
                                   evolve_moments, collisions, external_source_settings,
-                                  numerical_dissipation, t_input)
+                                  num_diss_params, t_input)
     pdf = create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
 
     # create the 'moments' struct that contains various v-space moments and other
@@ -63,13 +63,13 @@ function allocate_pdf_and_moments(composition, r, z, vperp, vpa, vzeta, vr, vz,
     ion = create_moments_ion(z.n, r.n, composition.n_ion_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, external_source_settings.ion,
-        numerical_dissipation)
+        num_diss_params)
     electron = create_moments_electron(z.n, r.n,
-        composition.electron_physics, numerical_dissipation)
+        composition.electron_physics, num_diss_params)
     neutral = create_moments_neutral(z.n, r.n, composition.n_neutral_species,
         evolve_moments.density, evolve_moments.parallel_flow,
         evolve_moments.parallel_pressure, external_source_settings.neutral,
-        numerical_dissipation)
+        num_diss_params)
 
     if abs(collisions.ionization) > 0.0 || z.bc == "wall"
         # if ionization collisions are included or wall BCs are enforced, then particle
@@ -1365,7 +1365,7 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
         enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
                                                     vperp_spectral, vpa_spectral,
                                                     vpa_advect, moments,
-                                                    num_diss_params.vpa_dissipation_coefficient > 0.0,
+                                                    num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
                                                     me_over_mi)
         begin_r_region()
         @loop_r ir begin
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 90ec6f310..edb7054b8 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -370,14 +370,19 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
 
     is_1V = (vperp.ngrid == vperp.nelement_global == 1 && vzeta.ngrid ==
              vzeta.nelement_global == 1 && vr.ngrid == vr.nelement_global == 1)
-    num_diss_params = setup_numerical_dissipation(
-        get(scan_input, "numerical_dissipation", Dict{String,Any}()), is_1V)
+    
+    ion_num_diss_param_dict = get(scan_input, "ion_numerical_dissipation", Dict{String,Any}())
+    electron_num_diss_param_dict = get(scan_input, "electron_numerical_dissipation", Dict{String,Any}())
+    neutral_num_diss_param_dict = get(scan_input, "neutral_numerical_dissipation", Dict{String,Any}())
+    num_diss_params = setup_numerical_dissipation(ion_num_diss_param_dict,
+                                                  electron_num_diss_param_dict,
+                                                  neutral_num_diss_param_dict, is_1V)
 
     # vperp.bc is set here (a bit out of place) so that we can use
-    # num_diss_params.vperp_dissipation_coefficient to set the default.
+    # num_diss_params.ion.vperp_dissipation_coefficient to set the default.
     vperp.bc = get(scan_input, "vperp_bc",
                    (collisions.nuii > 0.0 ||
-                    num_diss_params.vperp_dissipation_coefficient > 0.0) ?
+                    num_diss_params.ion.vperp_dissipation_coefficient > 0.0) ?
                     "zero" : "none")
     
     #########################################################################
diff --git a/moment_kinetics/src/numerical_dissipation.jl b/moment_kinetics/src/numerical_dissipation.jl
index 64eb28138..054357c29 100644
--- a/moment_kinetics/src/numerical_dissipation.jl
+++ b/moment_kinetics/src/numerical_dissipation.jl
@@ -16,11 +16,36 @@ using ..calculus: derivative!, second_derivative!
 using ..derivatives: derivative_r!, derivative_z!
 using ..type_definitions: mk_float
 
-Base.@kwdef struct numerical_dissipation_parameters
+
+#############################################################
+########### Numerical Dissipation Parameter setup ###########
+"""
+Define the dissipation parameters for each species, which means
+there need to be three sections in each input file that specify
+the parameters required of each species, as follows:
+
+```
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient
+...
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient
+...
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient
+...
+```
+
+There will still be the -1.0 default parameters.
+"""
+
+# define individual structs for each species with their particular parameters
+Base.@kwdef struct ion_num_diss_params
     vpa_boundary_buffer_damping_rate::mk_float = -1.0
     vpa_boundary_buffer_diffusion_coefficient::mk_float = -1.0
     vpa_dissipation_coefficient::mk_float = -1.0
-    vz_dissipation_coefficient::mk_float = -1.0
     vperp_dissipation_coefficient::mk_float = -1.0
     z_dissipation_coefficient::mk_float = -1.0
     r_dissipation_coefficient::mk_float = -1.0
@@ -28,18 +53,52 @@ Base.@kwdef struct numerical_dissipation_parameters
     force_minimum_pdf_value::Union{Nothing,mk_float} = nothing
 end
 
-function setup_numerical_dissipation(input_section::Dict, is_1V)
-    if is_1V && "vpa_dissipation_coefficient" ∈ keys(input_section)
+Base.@kwdef struct electron_num_diss_params
+    vpa_boundary_buffer_damping_rate::mk_float = -1.0
+    vpa_boundary_buffer_diffusion_coefficient::mk_float = -1.0
+    vpa_dissipation_coefficient::mk_float = -1.0
+    vperp_dissipation_coefficient::mk_float = -1.0
+    z_dissipation_coefficient::mk_float = -1.0
+    r_dissipation_coefficient::mk_float = -1.0
+    moment_dissipation_coefficient::mk_float = -1.0
+    force_minimum_pdf_value::Union{Nothing,mk_float} = nothing
+end
+
+Base.@kwdef struct neutral_num_diss_params
+    vz_dissipation_coefficient::mk_float = -1.0
+    z_dissipation_coefficient::mk_float = -1.0
+    r_dissipation_coefficient::mk_float = -1.0
+    moment_dissipation_coefficient::mk_float = -1.0
+    force_minimum_pdf_value::Union{Nothing,mk_float} = nothing
+end
+
+struct numerical_dissipation_parameters
+    ion::ion_num_diss_params
+    electron::electron_num_diss_params
+    neutral::neutral_num_diss_params
+end
+
+######### End Of Numerical Dissipation Parameter setup #########
+################################################################
+
+function setup_numerical_dissipation(ion_input::Dict, electron_input::Dict, 
+                                     neutral_input::Dict, is_1V)
+    if is_1V && "vpa_dissipation_coefficient" ∈ keys(ion_input)
         # Set default for vz_dissipation_coefficient the same as
-        # vpa_dissipation_coefficient for 1V case
-        input_section["vz_dissipation_coefficient"] =
-            get(input_section, "vz_dissipation_coefficient",
-                input_section["vpa_dissipation_coefficient"])
+        # ion_vpa_dissipation_coefficient for 1V case
+        neutral_input["vz_dissipation_coefficient"] =
+            get(neutral_input, "vz_dissipation_coefficient",
+                ion_input["vpa_dissipation_coefficient"])
     end
 
-    input = Dict(Symbol(k)=>v for (k,v) in input_section)
+    ion_input_dict = Dict(Symbol(k)=>v for (k,v) in ion_input)
+    ion_params = ion_num_diss_params(; ion_input_dict...)
+    electron_input_dict = Dict(Symbol(k)=>v for (k,v) in electron_input)
+    electron_params = electron_num_diss_params(; electron_input_dict...)
+    neutral_input_dict = Dict(Symbol(k)=>v for (k,v) in neutral_input)
+    neutral_params = neutral_num_diss_params(; neutral_input_dict...)
 
-    return numerical_dissipation_parameters(; input...)
+    return numerical_dissipation_parameters(ion_params, electron_params, neutral_params)
 end
 
 """
@@ -50,13 +109,11 @@ Disabled by default.
 
 The damping rate is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_boundary_buffer_damping_rate = 0.1
 ```
 """
-function vpa_boundary_buffer_decay!(f_out, fvec_in, moments, vpa, dt,
-                                    num_diss_params::numerical_dissipation_parameters)
-    damping_rate_prefactor = num_diss_params.vpa_boundary_buffer_damping_rate
+function vpa_boundary_buffer_decay!(f_out, fvec_in, moments, vpa, dt, damping_rate_prefactor)
 
     if damping_rate_prefactor <= 0.0
         return nothing
@@ -149,13 +206,11 @@ Disabled by default.
 
 The maximum diffusion rate in the buffer is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_boundary_buffer_diffusion_coefficient = 0.1
 ```
 """
-function vpa_boundary_buffer_diffusion!(f_out, fvec_in, vpa, vpa_spectral, dt,
-                                        num_diss_params::numerical_dissipation_parameters)
-    diffusion_prefactor = num_diss_params.vpa_boundary_buffer_diffusion_coefficient
+function vpa_boundary_buffer_diffusion!(f_out, fvec_in, vpa, vpa_spectral, dt, diffusion_prefactor)
 
     if diffusion_prefactor <= 0.0
         return nothing
@@ -238,14 +293,13 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.1
 ```
 """
 function vpa_dissipation!(f_out, f_in, vpa, spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters) where T_spectral
+                          diffusion_coefficient) where T_spectral
 
-    diffusion_coefficient = num_diss_params.vpa_dissipation_coefficient
     if diffusion_coefficient <= 0.0 || vpa.n == 1
         return nothing
     end
@@ -300,16 +354,15 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vperp_dissipation_coefficient = 0.1
 ```
 """
 function vperp_dissipation!(f_out, f_in, vperp, spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters) where T_spectral
+                            diffusion_coefficient) where T_spectral
     
     begin_s_r_z_vpa_region()
 
-    diffusion_coefficient = num_diss_params.vperp_dissipation_coefficient
     if diffusion_coefficient <= 0.0 || vperp.n == 1
         return nothing
     end
@@ -329,7 +382,7 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 z_dissipation_coefficient = 0.1
 ```
 
@@ -339,9 +392,8 @@ on internal or external element boundaries
 
 """
 function z_dissipation!(f_out, f_in, z, z_spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+                        diffusion_coefficient, scratch_dummy) where T_spectral
 
-    diffusion_coefficient = num_diss_params.z_dissipation_coefficient
     if diffusion_coefficient <= 0.0 || z.n == 1
         return nothing
     end
@@ -374,7 +426,7 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 r_dissipation_coefficient = 0.1
 
 ```
@@ -385,9 +437,8 @@ on internal or external element boundaries
 
 """
 function r_dissipation!(f_out, f_in, r, r_spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+                        diffusion_coefficient, scratch_dummy) where T_spectral
 
-    diffusion_coefficient = num_diss_params.r_dissipation_coefficient
     if diffusion_coefficient <= 0.0 || r.n == 1
         return nothing
     end
@@ -420,14 +471,13 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[neutral_numerical_dissipation]
 vz_dissipation_coefficient = 0.1
 ```
 """
 function vz_dissipation_neutral!(f_out, f_in, vz, spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters) where T_spectral
+                                 diffusion_coefficient) where T_spectral
 
-    diffusion_coefficient = num_diss_params.vz_dissipation_coefficient
     if diffusion_coefficient <= 0.0
         return nothing
     end
@@ -449,7 +499,7 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[neutral_numerical_dissipation]
 z_dissipation_coefficient = 0.1
 ```
 
@@ -459,9 +509,8 @@ on internal or external element boundaries
 
 """
 function z_dissipation_neutral!(f_out, f_in, z, z_spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+                                diffusion_coefficient, scratch_dummy) where T_spectral
 
-    diffusion_coefficient = num_diss_params.z_dissipation_coefficient
     if diffusion_coefficient <= 0.0
         return nothing
     end
@@ -494,7 +543,7 @@ Disabled by default.
 
 The diffusion coefficient is set in the input TOML file by the parameter
 ```
-[numerical_dissipation]
+[neutral_numerical_dissipation]
 r_dissipation_coefficient = 0.1
 
 ```
@@ -505,9 +554,8 @@ on internal or external element boundaries
 
 """
 function r_dissipation_neutral!(f_out, f_in, r, r_spectral::T_spectral, dt,
-        num_diss_params::numerical_dissipation_parameters, scratch_dummy) where T_spectral
+                                diffusion_coefficient, scratch_dummy) where T_spectral
 
-    diffusion_coefficient = num_diss_params.r_dissipation_coefficient
     if diffusion_coefficient <= 0.0 || r.n == 1
         return nothing
     end
@@ -534,17 +582,16 @@ function r_dissipation_neutral!(f_out, f_in, r, r_spectral::T_spectral, dt,
 end
 
 """
-    force_minimum_pdf_value!(f, num_diss_paras::numerical_dissipation_parameters)
+    force_minimum_pdf_value!(f, minval)
 
 Set a minimum value for the pdf-sized array `f`. Any points less than the minimum are
 set to the minimum. By default, no minimum is applied. The minimum value can be set by
 ```
-[numerical_dissipation]
+[ion_numerical_dissipation]
 force_minimum_pdf_value = 0.0
 ```
 """
-function force_minimum_pdf_value!(f, num_diss_params::numerical_dissipation_parameters)
-    minval = num_diss_params.force_minimum_pdf_value
+function force_minimum_pdf_value!(f, minval)
 
     if minval === nothing
         return nothing
@@ -560,17 +607,16 @@ function force_minimum_pdf_value!(f, num_diss_params::numerical_dissipation_para
 end
 
 """
-    force_minimum_pdf_value_neutral!(f, num_diss_paras::numerical_dissipation_parameters)
+    force_minimum_pdf_value_neutral!(f, minval)
 
 Set a minimum value for the neutral-pdf-sized array `f`. Any points less than the minimum
 are set to the minimum. By default, no minimum is applied. The minimum value can be set by
 ```
-[numerical_dissipation]
+[neutral_numerical_dissipation]
 force_minimum_pdf_value = 0.0
 ```
 """
-function force_minimum_pdf_value_neutral!(f, num_diss_params::numerical_dissipation_parameters)
-    minval = num_diss_params.force_minimum_pdf_value
+function force_minimum_pdf_value_neutral!(f, minval)
 
     if minval === nothing
         return nothing
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 003cf0ae9..f57ebb6b6 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -23,7 +23,7 @@ using ..velocity_moments: update_density!, update_upar!, update_ppar!, update_pp
 using ..velocity_moments: update_neutral_density!, update_neutral_qz!
 using ..velocity_moments: update_neutral_uzeta!, update_neutral_uz!, update_neutral_ur!
 using ..velocity_moments: update_neutral_pzeta!, update_neutral_pz!, update_neutral_pr!
-using ..velocity_moments: calculate_moment_derivatives!, calculate_moment_derivatives_neutral!
+using ..velocity_moments: calculate_ion_moment_derivatives!, calculate_neutral_moment_derivatives!
 using ..velocity_moments: calculate_electron_moment_derivatives!
 using ..velocity_moments: update_chodura!
 using ..velocity_grid_transforms: vzvrvzeta_to_vpavperp!, vpavperp_to_vzvrvzeta!
@@ -276,6 +276,9 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
+    ion_mom_diss_coeff = num_diss_params.ion.moment_dissipation_coefficient
+    electron_mom_diss_coeff = num_diss_params.electron.moment_dissipation_coefficient
+    neutral_mom_diss_coeff = num_diss_params.neutral.moment_dissipation_coefficient
     # create array containing coefficients needed for the Runge Kutta time advance
     rk_coefs = setup_runge_kutta_coefficients(t_input.n_rk_stages)
     # create the 'advance' struct to be used in later Euler advance to
@@ -299,8 +302,8 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
     end
     # update the derivatives of the electron moments as these may be needed when
     # computing the electrostatic potential (and components of the electric field)
-    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 
-                                           z_spectral, num_diss_params, composition.electron_physics)
+    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                           electron_mom_diss_coeff, composition.electron_physics)
     # initialize the electrostatic potential
     begin_serial_region()
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
@@ -313,9 +316,10 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
     # 'speed' in advect objects, which are needed for boundary conditions on the
     # distribution function which is then used to (possibly) re-calculate the moments
     # after which the initial values of moment derivatives are re-calculated.
-    calculate_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, num_diss_params)
-    calculate_moment_derivatives_neutral!(moments, scratch[1], scratch_dummy, z,
-                                          z_spectral, num_diss_params)
+    calculate_ion_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                      ion_mom_diss_coeff)
+    calculate_neutral_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                          neutral_mom_diss_coeff)
 
     r_advect = advection_structs.r_advect
     z_advect = advection_structs.z_advect
@@ -427,8 +431,8 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
     if !restarting
         begin_serial_region()
         # ensure initial pdf has no negative values
-        force_minimum_pdf_value!(pdf.ion.norm, num_diss_params)
-        force_minimum_pdf_value_neutral!(pdf.neutral.norm, num_diss_params)
+        force_minimum_pdf_value!(pdf.ion.norm, num_diss_params.ion.force_minimum_pdf_value)
+        force_minimum_pdf_value_neutral!(pdf.neutral.norm, num_diss_params.neutral.force_minimum_pdf_value)
         # enforce boundary conditions and moment constraints to ensure a consistent initial
         # condition
         enforce_boundary_conditions!(
@@ -546,11 +550,12 @@ function setup_time_advance!(pdf, fields, scratch, vz, vr, vzeta, vpa, vperp, z,
         end
     end
 
-    calculate_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, num_diss_params)
-    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, 
-                                           z_spectral, num_diss_params, composition.electron_physics)
-    calculate_moment_derivatives_neutral!(moments, scratch[1], scratch_dummy, z,
-                                          z_spectral, num_diss_params)
+    calculate_ion_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                      ion_mom_diss_coeff)
+    calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                      electron_mom_diss_coeff, composition.electron_physics)
+    calculate_neutral_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
+                                      neutral_mom_diss_coeff)
     # update the electrostatic potential and components of the electric field, as pdfs and moments
     # may have changed due to enforcing boundary/moment constraints                                      
     update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral,
@@ -708,11 +713,11 @@ function setup_advance_flags(moments, composition, t_input, collisions,
         end
 
         # flag to determine if a d^2/dr^2 operator is present
-        r_diffusion = (advance_numerical_dissipation && num_diss_params.r_dissipation_coefficient > 0.0)
+        r_diffusion = (advance_numerical_dissipation && num_diss_params.ion.r_dissipation_coefficient > 0.0)
         # flag to determine if a d^2/dvpa^2 operator is present
-        vpa_diffusion = ((advance_numerical_dissipation && num_diss_params.vpa_dissipation_coefficient > 0.0) || explicit_weakform_fp_collisions)
-        vperp_diffusion = ((advance_numerical_dissipation && num_diss_params.vperp_dissipation_coefficient > 0.0) || explicit_weakform_fp_collisions)
-        vz_diffusion = (advance_numerical_dissipation && num_diss_params.vz_dissipation_coefficient > 0.0)
+        vpa_diffusion = ((advance_numerical_dissipation && num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || explicit_weakform_fp_collisions)
+        vperp_diffusion = ((advance_numerical_dissipation && num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || explicit_weakform_fp_collisions)
+        vz_diffusion = (advance_numerical_dissipation && num_diss_params.neutral.vz_dissipation_coefficient > 0.0)
     end
 
     manufactured_solns_test = manufactured_solns_input.use_for_advance
@@ -1355,7 +1360,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     # Ensure there are no negative values in the pdf before applying boundary
     # conditions, so that negative deviations do not mess up the integral-constraint
     # corrections in the sheath boundary conditions.
-    force_minimum_pdf_value!(new_scratch.pdf, num_diss_params)
+    force_minimum_pdf_value!(new_scratch.pdf, num_diss_params.ion.force_minimum_pdf_value)
 
     # Enforce boundary conditions in z and vpa on the distribution function.
     # Must be done after Runge Kutta update so that the boundary condition applied to
@@ -1406,8 +1411,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                  composition, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
 
-    calculate_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
-                                  num_diss_params)
+    calculate_ion_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
+                                      num_diss_params.ion.moment_dissipation_coefficient)
 
     # update the lowest three electron moments (density, upar and ppar)
     calculate_electron_density!(new_scratch.electron_density, moments.electron.dens_updated, new_scratch.density)
@@ -1433,7 +1438,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                                      new_scratch.electron_density)
     # calculate the corresponding zed derivatives of the moments
     calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
-                                           num_diss_params, composition.electron_physics)
+                                           num_diss_params.electron.moment_dissipation_coefficient, 
+                                           composition.electron_physics)
     # update the electron parallel heat flux
     calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.pdf_electron,
         new_scratch.electron_ppar, new_scratch.electron_upar, moments.electron.vth, moments.electron.dT_dz, 
@@ -1459,7 +1465,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         # Ensure there are no negative values in the pdf before applying boundary
         # conditions, so that negative deviations do not mess up the integral-constraint
         # corrections in the sheath boundary conditions.
-        force_minimum_pdf_value_neutral!(new_scratch.pdf_neutral, num_diss_params)
+        force_minimum_pdf_value_neutral!(new_scratch.pdf_neutral, num_diss_params.neutral.force_minimum_pdf_value)
 
         # Enforce boundary conditions in z and vpa on the distribution function.
         # Must be done after Runge Kutta update so that the boundary condition applied to
@@ -1506,8 +1512,8 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                            r, composition, moments.evolve_density, moments.evolve_upar,
                            moments.evolve_ppar)
 
-        calculate_moment_derivatives_neutral!(moments, new_scratch, scratch_dummy, z,
-                                              z_spectral, num_diss_params)
+        calculate_neutral_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral, 
+                                              num_diss_params.neutral.moment_dissipation_coefficient)
     end
 
     # update the electrostatic potential phi
@@ -1901,19 +1907,19 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     # add numerical dissipation
     if advance.numerical_dissipation
         vpa_dissipation!(fvec_out.pdf, fvec_in.pdf, vpa, vpa_spectral, dt,
-                         num_diss_params)
+                         num_diss_params.ion.vpa_dissipation_coefficient)
         vperp_dissipation!(fvec_out.pdf, fvec_in.pdf, vperp, vperp_spectral, dt,
-                         num_diss_params)
+                         num_diss_params.ion.vperp_dissipation_coefficient)
         z_dissipation!(fvec_out.pdf, fvec_in.pdf, z, z_spectral, dt,
-                       num_diss_params, scratch_dummy)
+                       num_diss_params.ion.z_dissipation_coefficient, scratch_dummy)
         r_dissipation!(fvec_out.pdf, fvec_in.pdf, r, r_spectral, dt,
-                       num_diss_params, scratch_dummy)
+                       num_diss_params.ion.r_dissipation_coefficient, scratch_dummy)
         vz_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, vz,
-                                vz_spectral, dt, num_diss_params)
+                                vz_spectral, dt, num_diss_params.neutral.vz_dissipation_coefficient)
         z_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, z, z_spectral,
-                               dt, num_diss_params, scratch_dummy)
+                               dt, num_diss_params.neutral.z_dissipation_coefficient, scratch_dummy)
         r_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, r, r_spectral,
-                               dt, num_diss_params, scratch_dummy)
+                               dt, num_diss_params.neutral.r_dissipation_coefficient, scratch_dummy)
     end
     # advance with the Fokker-Planck self-collision operator
     if advance.explicit_weakform_fp_collisions
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 35dd16802..52fb0625b 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -52,7 +52,7 @@ using ..moment_kinetics_structs: moments_ion_substruct, moments_electron_substru
 """
 """
 function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
-                            evolve_ppar, ion_source_settings, numerical_dissipation)
+                            evolve_ppar, ion_source_settings, num_diss_params)
     # allocate array used for the particle density
     density = allocate_shared_float(nz, nr, n_species)
     # allocate array of Bools that indicate if the density is updated for each species
@@ -96,8 +96,7 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
         ddens_dz = nothing
         ddens_dz_upwind = nothing
     end
-    if evolve_density &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if evolve_density && num_diss_params.ion.moment_dissipation_coefficient > 0.0
 
         d2dens_dz2 = allocate_shared_float(nz, nr, n_species)
     else
@@ -113,8 +112,7 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
     else
         dupar_dz_upwind = nothing
     end
-    if evolve_upar &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if evolve_upar && num_diss_params.ion.moment_dissipation_coefficient > 0.0
 
         d2upar_dz2 = allocate_shared_float(nz, nr, n_species)
     else
@@ -200,7 +198,7 @@ end
 """
 create a moment struct containing information about the electron moments
 """
-function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
+function create_moments_electron(nz, nr, electron_model, num_diss_params)
     # allocate array used for the particle density
     density = allocate_shared_float(nz, nr)
     # initialise Bool variable that indicates if the density is updated for each species
@@ -241,7 +239,7 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
         dppar_dz_upwind = nothing
         dT_dz_upwind = nothing
     end
-    if numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if num_diss_params.electron.moment_dissipation_coefficient > 0.0
         d2ppar_dz2 = allocate_shared_float(nz, nr)
     else
         d2ppar_dz2 = nothing
@@ -277,7 +275,7 @@ end
     
 function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
                                 evolve_ppar, neutral_source_settings,
-                                numerical_dissipation)
+                                num_diss_params)
     density = allocate_shared_float(nz, nr, n_species)
     density_updated = allocate_bool(n_species)
     density_updated .= false
@@ -320,8 +318,7 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
         ddens_dz = nothing
         ddens_dz_upwind = nothing
     end
-    if evolve_density &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if evolve_density && num_diss_params.neutral.moment_dissipation_coefficient > 0.0
 
         d2dens_dz2 = allocate_shared_float(nz, nr, n_species)
     else
@@ -337,8 +334,7 @@ function create_moments_neutral(nz, nr, n_species, evolve_density, evolve_upar,
     else
         duz_dz_upwind = nothing
     end
-    if evolve_upar &&
-            numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if evolve_upar && num_diss_params.neutral.moment_dissipation_coefficient > 0.0
 
         d2uz_dz2 = allocate_shared_float(nz, nr, n_species)
     else
@@ -864,8 +860,8 @@ end
 """
 Pre-calculate spatial derivatives of the moments that will be needed for the time advance
 """
-function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spectral,
-                                       numerical_dissipation)
+function calculate_ion_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spectral,
+                                           ion_mom_diss_coeff)
     begin_s_r_region()
 
     density = scratch.density
@@ -891,8 +887,7 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
                              dummy_zrs, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
     end
-    if moments.evolve_density &&
-        numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if moments.evolve_density && ion_mom_diss_coeff > 0.0
 
         # centred second derivative for dissipation
         @views derivative_z!(dummy_zrs, density, buffer_r_1, buffer_r_2, buffer_r_3,
@@ -914,8 +909,7 @@ function calculate_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spe
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z)
     end
-    if moments.evolve_upar &&
-        numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if moments.evolve_upar && ion_mom_diss_coeff > 0.0
 
         # centred second derivative for dissipation
         @views derivative_z!(dummy_zrs, upar, buffer_r_1, buffer_r_2, buffer_r_3,
@@ -953,7 +947,7 @@ end
 Pre-calculate spatial derivatives of the electron moments that will be needed for the time advance
 """
 function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy, z, z_spectral,
-                                                numerical_dissipation, electron_model)
+                                                electron_mom_diss_coeff, electron_model)
     begin_r_region()
 
     dens = scratch.electron_density
@@ -982,7 +976,7 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
     end
 
     # centred second derivative for dissipation
-    if numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if electron_mom_diss_coeff > 0.0
         @views derivative_z!(dummy_zr, ppar, buffer_r_1, buffer_r_2, buffer_r_3,
                              buffer_r_4, z_spectral, z)
         @views derivative_z!(moments.electron.d2ppar_dz2, dummy_zr, buffer_r_1,
@@ -1478,8 +1472,8 @@ end
 Pre-calculate spatial derivatives of the neutral moments that will be needed for the time
 advance
 """
-function calculate_moment_derivatives_neutral!(moments, scratch, scratch_dummy, z,
-                                               z_spectral, numerical_dissipation)
+function calculate_neutral_moment_derivatives!(moments, scratch, scratch_dummy, z,
+                                               z_spectral, neutral_mom_diss_coeff)
     begin_sn_r_region()
 
     density = scratch.density_neutral
@@ -1506,8 +1500,7 @@ function calculate_moment_derivatives_neutral!(moments, scratch, scratch_dummy,
                              dummy_zrsn, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z; neutrals=true)
     end
-    if moments.evolve_density &&
-        numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if moments.evolve_density && neutral_mom_diss_coeff > 0.0
 
         # centred second derivative for dissipation
         @views derivative_z!(dummy_zrsn, density, buffer_r_1, buffer_r_2, buffer_r_3,
@@ -1531,8 +1524,7 @@ function calculate_moment_derivatives_neutral!(moments, scratch, scratch_dummy,
                              buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
                              buffer_r_5, buffer_r_6, z_spectral, z; neutrals=true)
     end
-    if moments.evolve_upar &&
-        numerical_dissipation.moment_dissipation_coefficient > 0.0
+    if moments.evolve_upar && neutral_mom_diss_coeff > 0.0
 
         # centred second derivative for dissipation
         @views derivative_z!(dummy_zrsn, uz, buffer_r_1, buffer_r_2, buffer_r_3,
diff --git a/moment_kinetics/test/harrisonthompson.jl b/moment_kinetics/test/harrisonthompson.jl
index 1f196e32b..2a0c4f25f 100644
--- a/moment_kinetics/test/harrisonthompson.jl
+++ b/moment_kinetics/test/harrisonthompson.jl
@@ -141,7 +141,7 @@ test_input_chebyshev_split1 = merge(test_input_chebyshev,
 test_input_chebyshev_split2 = merge(test_input_chebyshev_split1,
                                     Dict("run_name" => "chebyshev_pseudospectral_split2",
                                          "evolve_moments_parallel_flow" => true,
-                                         "numerical_dissipation" => Dict("force_minimum_pdf_value" => 0.0)))
+                                         "ion_numerical_dissipation" => Dict("force_minimum_pdf_value" => 0.0)))
 
 test_input_chebyshev_split3 = merge(test_input_chebyshev_split2,
                                     Dict("run_name" => "chebyshev_pseudospectral_split3",
diff --git a/moment_kinetics/test/recycling_fraction_tests.jl b/moment_kinetics/test/recycling_fraction_tests.jl
index 2e4e0eb2a..671830f9d 100644
--- a/moment_kinetics/test/recycling_fraction_tests.jl
+++ b/moment_kinetics/test/recycling_fraction_tests.jl
@@ -108,7 +108,8 @@ test_input_split3 = merge(test_input_split2,
                                "vpa_nelement" => 31,
                                "vz_nelement" => 31,
                                "evolve_moments_parallel_pressure" => true,
-                               "numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0, "vpa_dissipation_coefficient" => 1e-2)))
+                               "ion_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0, "vpa_dissipation_coefficient" => 1e-2),
+                               "neutral_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0, "vz_dissipation_coefficient" => 1e-2)))
 
 """
 Run a test for a single set of parameters
diff --git a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5.toml b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5.toml
index 0a363e4a9..a10392d9f 100644
--- a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5.toml
+++ b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5.toml
@@ -81,8 +81,14 @@ z_width = 0.125
 source_strength = 2.0
 source_T = 2.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 #force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+#force_minimum_pdf_value = 0.0
diff --git a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split1.toml b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split1.toml
index 8cae6950d..1e44b5482 100644
--- a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split1.toml
+++ b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split1.toml
@@ -81,8 +81,14 @@ z_width = 0.125
 source_strength = 2.0
 source_T = 2.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 #force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+#force_minimum_pdf_value = 0.0
diff --git a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split2.toml b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split2.toml
index 986e8a616..fef873fe8 100644
--- a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split2.toml
+++ b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split2.toml
@@ -81,8 +81,14 @@ z_width = 0.125
 source_strength = 2.0
 source_T = 2.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
diff --git a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split3.toml b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split3.toml
index 28bcc6831..8256f398a 100644
--- a/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split3.toml
+++ b/publication_inputs/2023_EFTC_jto-poster/wall-bc/wall-bc_recyclefraction0.5_split3.toml
@@ -81,8 +81,14 @@ z_width = 0.125
 source_strength = 2.0
 source_T = 2.0
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 1.0e-1
 #vpa_dissipation_coefficient = 1.0e-2
 #vpa_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[ion_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
diff --git a/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml b/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
index 45d596904..9a475346b 100644
--- a/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
+++ b/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
@@ -91,7 +91,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml b/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml
index a465d248c..347dd2b04 100644
--- a/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml
+++ b/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml
@@ -92,7 +92,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
index 6a89faab1..cb8ce6cf2 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
index a90350f22..6d9c31403 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
index 92a8aa331..c0ec6f8f5 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
index 5f63e0aec..dff2bf452 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.0
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
index c8e8df9ee..a5be20d0d 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
index 60f09b084..218ecaf2a 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
@@ -86,7 +86,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.0
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
index 98efd587f..3c716fa21 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = -1.0
 z_dissipation_coefficient = -1.0
 r_dissipation_coefficient = -1.0
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
index 135f5817e..52acf4e29 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = -1.0
 z_dissipation_coefficient = -1.0
 r_dissipation_coefficient = -1.0
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
index a300c9a56..cf3eb5d19 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
index c98077d3a..5bd48c14b 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
index 3ec12c9fe..b13205d70 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
index e936ef06e..d011761cb 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
@@ -87,7 +87,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 vpa_dissipation_coefficient = 0.001
 vperp_dissipation_coefficient = 0.001
 #z_dissipation_coefficient = 0.1
diff --git a/runs/2D-wall_MMS_nel_r_16_z_16_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_16_z_16_vpa_16_vperp_1_diss.toml
index 52f761612..0222c0ac8 100644
--- a/runs/2D-wall_MMS_nel_r_16_z_16_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_16_z_16_vpa_16_vperp_1_diss.toml
@@ -77,7 +77,7 @@ vzeta_L = 12.0
 vzeta_bc = "periodic"
 vzeta_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.01
diff --git a/runs/2D-wall_MMS_nel_r_2_z_2_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_2_z_2_vpa_16_vperp_1_diss.toml
index e67fc2df9..18e004b41 100644
--- a/runs/2D-wall_MMS_nel_r_2_z_2_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_2_z_2_vpa_16_vperp_1_diss.toml
@@ -77,7 +77,7 @@ vzeta_L = 12.0
 vzeta_bc = "periodic"
 vzeta_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.01
diff --git a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
index 5bbb30479..10745c71c 100644
--- a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
@@ -86,7 +86,7 @@ vzeta_discretization = "chebyshev_pseudospectral"
  use_vpabar_in_mms_dfni=true
  alpha_switch=1.0
  type="default"
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.1
diff --git a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss5.toml b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss5.toml
index 0ebf89144..b4b1206a2 100644
--- a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss5.toml
+++ b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss5.toml
@@ -77,7 +77,7 @@ vzeta_L = 12.0
 vzeta_bc = "periodic"
 vzeta_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.01
diff --git a/runs/2D-wall_MMS_nel_r_4_z_4_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_4_z_4_vpa_16_vperp_1_diss.toml
index 7ace633dd..fc0e45f6c 100644
--- a/runs/2D-wall_MMS_nel_r_4_z_4_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_4_z_4_vpa_16_vperp_1_diss.toml
@@ -77,7 +77,7 @@ vzeta_L = 12.0
 vzeta_bc = "periodic"
 vzeta_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.01
diff --git a/runs/2D-wall_MMS_nel_r_8_z_8_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_8_z_8_vpa_16_vperp_1_diss.toml
index 672f7d289..af309dc54 100644
--- a/runs/2D-wall_MMS_nel_r_8_z_8_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_8_z_8_vpa_16_vperp_1_diss.toml
@@ -77,7 +77,7 @@ vzeta_L = 12.0
 vzeta_bc = "periodic"
 vzeta_discretization = "chebyshev_pseudospectral"
 
-[numerical_dissipation]
+[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.1
 #z_dissipation_coefficient = 0.1
 r_dissipation_coefficient = 0.01
diff --git a/runs/2V-evolve_ngrid_3_nel_r_1_z_1_vpa_6_vperp_3_fkpl_test.toml b/runs/2V-evolve_ngrid_3_nel_r_1_z_1_vpa_6_vperp_3_fkpl_test.toml
index 2f32d9ca3..031fd5580 100644
--- a/runs/2V-evolve_ngrid_3_nel_r_1_z_1_vpa_6_vperp_3_fkpl_test.toml
+++ b/runs/2V-evolve_ngrid_3_nel_r_1_z_1_vpa_6_vperp_3_fkpl_test.toml
@@ -89,7 +89,7 @@ vperp_discretization = "gausslegendre_pseudospectral"
 #vzeta_bc = "periodic"
 #vzeta_discretization = "chebyshev_pseudospectral"
 
-#[numerical_dissipation]
+#[ion_numerical_dissipation]
 #vpa_dissipation_coefficient = 0.0
 #vperp_dissipation_coefficient = 0.0
 #z_dissipation_coefficient = 0.1
diff --git a/runs/wall+sheath-bc_boltzmann.toml b/runs/wall+sheath-bc_boltzmann.toml
index 1844deebf..b49b5a328 100644
--- a/runs/wall+sheath-bc_boltzmann.toml
+++ b/runs/wall+sheath-bc_boltzmann.toml
@@ -73,6 +73,14 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 moment_dissipation_coefficient = 0.0001
-vpa_dissipation_coefficient = 0.01
\ No newline at end of file
+vpa_dissipation_coefficient = 0.01
+
+[ion_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01
+
+[neutral_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vz_dissipation_coefficient = 0.01
\ No newline at end of file
diff --git a/runs/wall+sheath-bc_braginskii.toml b/runs/wall+sheath-bc_braginskii.toml
index 60c8cf4d0..616873e38 100644
--- a/runs/wall+sheath-bc_braginskii.toml
+++ b/runs/wall+sheath-bc_braginskii.toml
@@ -75,6 +75,15 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 moment_dissipation_coefficient = 0.0001
-vpa_dissipation_coefficient = 0.002
\ No newline at end of file
+vpa_dissipation_coefficient = 0.002
+
+[ion_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.002
+
+[neutral_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vz_dissipation_coefficient = 0.002
+
diff --git a/runs/wall+sheath-bc_braginskii_boltzmann_test.toml b/runs/wall+sheath-bc_braginskii_boltzmann_test.toml
index 13b94d2ff..ce8b7af1f 100644
--- a/runs/wall+sheath-bc_braginskii_boltzmann_test.toml
+++ b/runs/wall+sheath-bc_braginskii_boltzmann_test.toml
@@ -74,6 +74,14 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 moment_dissipation_coefficient = 0.0001
 vpa_dissipation_coefficient = 0.01
+
+[ion_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01
+
+[neutral_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vz_dissipation_coefficient = 0.01
\ No newline at end of file
diff --git a/runs/wall+sheath-bc_braginskii_colls.toml b/runs/wall+sheath-bc_braginskii_colls.toml
index eb0476bc5..3267b2402 100644
--- a/runs/wall+sheath-bc_braginskii_colls.toml
+++ b/runs/wall+sheath-bc_braginskii_colls.toml
@@ -75,6 +75,14 @@ vz_discretization = "chebyshev_pseudospectral"
 [output]
 ascii_output = true
 
-[numerical_dissipation]
+[electron_numerical_dissipation]
 moment_dissipation_coefficient = 0.0001
 vpa_dissipation_coefficient = 0.01
+
+[ion_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vpa_dissipation_coefficient = 0.01
+
+[neutral_numerical_dissipation]
+moment_dissipation_coefficient = 0.0001
+vz_dissipation_coefficient = 0.01
\ No newline at end of file
diff --git a/util/compare_collision_frequencies.jl b/util/compare_collision_frequencies.jl
index 8396e5cb7..a1fac065e 100644
--- a/util/compare_collision_frequencies.jl
+++ b/util/compare_collision_frequencies.jl
@@ -27,10 +27,10 @@ function compare_collision_frequencies(input_file::String,
     # v_∥ dissipation term is D d^2f/dv_∥^2. Inserting factors of c_ref, this is a bit like
     # pitch angle scattering D cref^2 d^2f/dv_∥^2 ~ D d^2f/dξ^2, so D is similar to a
     # (normalised) collision frequency.
-    if num_diss_params.vpa_dissipation_coefficient < 0.0
+    if num_diss_params.ion.vpa_dissipation_coefficient < 0.0
         nu_vpa_diss = 0.0
     else
-        nu_vpa_diss = num_diss_params.vpa_dissipation_coefficient /
+        nu_vpa_diss = num_diss_params.ion.vpa_dissipation_coefficient /
                       dimensional_parameters["timenorm"]
     end
 
@@ -89,17 +89,17 @@ function compare_collision_frequencies(input_file::String,
     println("classical heat chi_i0 with effective rho_i ", effective_classical_heat_chi_i0)
 
     # Get numerical diffusion parameters
-    if num_diss_params.r_dissipation_coefficient < 0.0
+    if num_diss_params.ion.r_dissipation_coefficient < 0.0
         D_r = 0.0
     else
-        D_r = Unitful.upreferred(num_diss_params.r_dissipation_coefficient *
+        D_r = Unitful.upreferred(num_diss_params.ion.r_dissipation_coefficient *
                                  dimensional_parameters["Lnorm"]^2 /
                                  dimensional_parameters["timenorm"])
     end
-    if num_diss_params.z_dissipation_coefficient < 0.0
+    if num_diss_params.ion.z_dissipation_coefficient < 0.0
         D_z = 0.0
     else
-        D_z = Unitful.upreferred(num_diss_params.z_dissipation_coefficient *
+        D_z = Unitful.upreferred(num_diss_params.ion.z_dissipation_coefficient *
                                  dimensional_parameters["Lnorm"]^2 /
                                  dimensional_parameters["timenorm"])
     end

From 54dd837c4f1366e8ab47a6d7653ebd1a96aadbe2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 17:48:11 +0000
Subject: [PATCH 186/428] Fix source moments to be compatible with
 source_type="energy_source"

Previously, the moments of the source term were introduced in vpa/vz
'advection speed' and in source_terms in terms of
`external_source_amplitude`. The moments were taken assuming that the
source is a Maxwellian with temperature `T_source`. This is not true for
"energy_source" (where the source is defined as roughly 'f_M - f' in
such a way that no particles are introduced). However, the moments of
the source term are already calculated in external_density_source and
external_pressure_source, so these can be used instead to support
more generic source terms.
---
 moment_kinetics/src/load_data.jl            | 22 ++++++++++++++--
 moment_kinetics/src/neutral_vz_advection.jl | 19 +++++++-------
 moment_kinetics/src/source_terms.jl         | 28 ++++++++++++---------
 moment_kinetics/src/vpa_advection.jl        | 19 +++++++-------
 4 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 9e3c18eca..e764cadfe 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3998,8 +3998,14 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         dqpar_dz = get_z_derivative(run_info, "parallel_heat_flux")
         if run_info.external_source_settings.ion.active
             external_source_amplitude = get_variable(run_info, "external_source_amplitude")
+            external_source_density_amplitude = get_variable(run_info, "external_source_density_amplitude")
+            external_source_momentum_amplitude = get_variable(run_info, "external_source_momentum_amplitude")
+            external_source_pressure_amplitude = get_variable(run_info, "external_source_pressure_amplitude")
         else
             external_source_amplitude = zeros(0,0,run_info.nt)
+            external_source_density_amplitude = zeros(0,0,run_info.nt)
+            external_source_momentum_amplitude = zeros(0,0,run_info.nt)
+            external_source_pressure_amplitude = zeros(0,0,run_info.nt)
         end
 
         nz, nr, nspecies, nt = size(vth)
@@ -4022,7 +4028,10 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
                                    dvth_dz=dvth_dz[:,:,:,it],
                                    dqpar_dz=dqpar_dz[:,:,:,it],
                                    vth=vth[:,:,:,it],
-                                   external_source_amplitude=external_source_amplitude[:,:,it]),
+                                   external_source_amplitude=external_source_amplitude[:,:,it],
+                                   external_source_density_amplitude=external_source_density_amplitude[:,:,it],
+                                   external_source_momentum_amplitude=external_source_momentum_amplitude[:,:,it],
+                                   external_source_pressure_amplitude=external_source_pressure_amplitude[:,:,it]),
                              evolve_density=run_info.evolve_density,
                              evolve_upar=run_info.evolve_upar,
                              evolve_ppar=run_info.evolve_ppar)
@@ -4195,8 +4204,14 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         dqz_dz = get_z_derivative(run_info, "qz_neutral")
         if run_info.external_source_settings.neutral.active
             external_source_amplitude = get_variable(run_info, "external_source_neutral_amplitude")
+            external_source_density_amplitude = get_variable(run_info, "external_source_neutral_density_amplitude")
+            external_source_momentum_amplitude = get_variable(run_info, "external_source_neutral_momentum_amplitude")
+            external_source_pressure_amplitude = get_variable(run_info, "external_source_neutral_pressure_amplitude")
         else
             external_source_amplitude = zeros(0,0,run_info.nt)
+            external_source_density_amplitude = zeros(0,0,run_info.nt)
+            external_source_momentum_amplitude = zeros(0,0,run_info.nt)
+            external_source_pressure_amplitude = zeros(0,0,run_info.nt)
         end
 
         nz, nr, nspecies, nt = size(vth)
@@ -4225,7 +4240,10 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
                                        dvth_dz=dvth_dz[:,:,:,it],
                                        dqz_dz=dqz_dz[:,:,:,it],
                                        vth=vth[:,:,:,it],
-                                       external_source_amplitude=external_source_amplitude[:,:,it]),
+                                       external_source_amplitude=external_source_amplitude[:,:,it],
+                                       external_source_density_amplitude=external_source_density_amplitude[:,:,it],
+                                       external_source_momentum_amplitude=external_source_momentum_amplitude[:,:,it],
+                                       external_source_pressure_amplitude=external_source_pressure_amplitude[:,:,it]),
                              evolve_density=run_info.evolve_density,
                              evolve_upar=run_info.evolve_upar,
                              evolve_ppar=run_info.evolve_ppar)
diff --git a/moment_kinetics/src/neutral_vz_advection.jl b/moment_kinetics/src/neutral_vz_advection.jl
index 36ef2c2c9..96b61cd39 100644
--- a/moment_kinetics/src/neutral_vz_advection.jl
+++ b/moment_kinetics/src/neutral_vz_advection.jl
@@ -127,19 +127,21 @@ function update_speed_n_u_p_evolution_neutral!(advect, fvec, moments, vz, z, r,
         end
     end
     if neutral_source_settings.active
-        source_amplitude = moments.neutral.external_source_amplitude
-        source_T = neutral_source_settings.source_T
+        source_density_amplitude = moments.neutral.external_source_density_amplitude
+        source_momentum_amplitude = moments.neutral.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.neutral.external_source_pressure_amplitude
         density = fvec.density_neutral
         uz = fvec.uz_neutral
         pz = fvec.pz_neutral
         vth = moments.neutral.vth
         vz_grid = vz.grid
         @loop_s_r_z is ir iz begin
-            prefactor = source_amplitude[iz,ir]
-            term1 = prefactor * uz[iz,ir,is]/(density[iz,ir,is]*vth[iz,ir,is])
+            term1 = source_density_amplitude[iz,ir] * uz[iz,ir,is]/(density[iz,ir,is]*vth[iz,ir,is])
             term2_over_vpa =
-                0.5 * prefactor * (-(0.5*source_T + uz[iz,ir,is]^2) / pz[iz,ir,is]
-                                   + 1.0/density[iz,ir,is])
+                -0.5 * (source_pressure_amplitude[iz,ir] +
+                        2.0 * uz[iz,ir,is] * source_momentum_amplitude[iz,ir]) /
+                       pz[iz,ir,is] +
+                0.5 * source_density_amplitude[iz,ir] / density[iz,ir,is]
             @loop_vzeta_vr_vz ivzeta ivr ivz begin
                 advect[is].speed[ivz,ivr,ivzeta,iz,ir] += term1 +
                                                           vz_grid[ivz] * term2_over_vpa
@@ -218,13 +220,12 @@ function update_speed_n_u_evolution_neutral!(advect, fvec, moments, vz, z, r, co
         end
     end
     if neutral_source_settings.active
-        source_amplitude = moments.neutral.external_source_amplitude
-        source_T = neutral_source_settings.source_T
+        source_density_amplitude = moments.neutral.external_source_density_amplitude
         density = fvec.density_neutral
         uz = fvec.uz_neutral
         vth = moments.neutral.vth
         @loop_sn_r_z isn ir iz begin
-            term = source_amplitude[iz,ir] * uz[iz,ir,isn] / density[iz,ir,isn]
+            term = source_density_amplitude[iz,ir] * uz[iz,ir,isn] / density[iz,ir,isn]
             @loop_vzeta_vr_vz ivzeta ivr ivz begin
                 advect[isn].speed[ivz,ivr,ivzeta,iz,ir] += term
             end
diff --git a/moment_kinetics/src/source_terms.jl b/moment_kinetics/src/source_terms.jl
index 790603f34..9850d6f23 100644
--- a/moment_kinetics/src/source_terms.jl
+++ b/moment_kinetics/src/source_terms.jl
@@ -65,9 +65,9 @@ function source_terms_evolve_density!(pdf_out, pdf_in, dens, upar, ddens_dz, dup
     end
 
     if ion_source_settings.active
-        source_amplitude = moments.ion.external_source_amplitude
+        source_density_amplitude = moments.ion.external_source_density_amplitude
         @loop_r_z ir iz begin
-            term = dt * source_amplitude[iz,ir] / dens[iz,ir]
+            term = dt * source_density_amplitude[iz,ir] / dens[iz,ir]
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= term * pdf_in[ivpa,ivperp,iz,ir]
             end
@@ -97,11 +97,13 @@ function source_terms_evolve_ppar_no_collisions!(pdf_out, pdf_in, dens, upar, pp
     end
 
     if ion_source_settings.active
-        source_amplitude = moments.ion.external_source_amplitude
-        source_T = ion_source_settings.source_T
+        source_density_amplitude = moments.ion.external_source_density_amplitude
+        source_momentum_amplitude = moments.ion.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.ion.external_source_pressure_amplitude
         @loop_r_z ir iz begin
-            term = dt * source_amplitude[iz,ir] *
-                   (1.5/dens[iz,ir] - (0.25 * source_T + 0.5 * upar[iz,ir]^2) / ppar[iz,ir])
+            term = dt * (1.5 * source_density_amplitude[iz,ir] / dens[iz,ir] -
+                         (0.5 * source_pressure_amplitude[iz,ir] +
+                          source_momentum_amplitude[iz,ir]) / ppar[iz,ir])
             @loop_vperp_vpa ivperp ivpa begin
                 pdf_out[ivpa,ivperp,iz,ir] -= term * pdf_in[ivpa,ivperp,iz,ir]
             end
@@ -191,9 +193,9 @@ function source_terms_evolve_density_neutral!(pdf_out, pdf_in, dens, upar, ddens
     end
 
     if neutral_source_settings.active
-        source_amplitude = moments.neutral.external_source_amplitude
+        source_density_amplitude = moments.neutral.external_source_density_amplitude
         @loop_r_z ir iz begin
-            term = dt * source_amplitude[iz,ir] / dens[iz,ir]
+            term = dt * source_density_amplitude[iz,ir] / dens[iz,ir]
             @loop_vzeta_vr_vz ivzeta ivr ivz begin
                 pdf_out[ivz,ivr,ivzeta,iz,ir] -= term * pdf_in[ivz,ivr,ivzeta,iz,ir]
             end
@@ -222,11 +224,13 @@ function source_terms_evolve_ppar_no_collisions_neutral!(pdf_out, pdf_in, dens,
     end
 
     if neutral_source_settings.active
-        source_amplitude = moments.neutral.external_source_amplitude
-        source_T = neutral_source_settings.source_T
+        source_density_amplitude = moments.neutral.external_source_density_amplitude
+        source_momentum_amplitude = moments.neutral.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.neutral.external_source_pressure_amplitude
         @loop_r_z ir iz begin
-            term = dt * source_amplitude[iz,ir] *
-                   (1.5/dens[iz,ir] - (0.25 * source_T + 0.5 * upar[iz,ir]^2) / ppar[iz,ir])
+            term = dt * (1.5 * source_density_amplitude[iz,ir] / dens[iz,ir] -
+                         (0.5 * source_pressure_amplitude[iz,ir] +
+                          source_momentum_amplitude[iz,ir]) / ppar[iz,ir])
             @loop_vzeta_vr_vz ivzeta ivr ivz begin
                 pdf_out[ivz,ivr,ivzeta,iz,ir] -= term * pdf_in[ivz,ivr,ivzeta,iz,ir]
             end
diff --git a/moment_kinetics/src/vpa_advection.jl b/moment_kinetics/src/vpa_advection.jl
index 169c0ea64..b9e5b4be4 100644
--- a/moment_kinetics/src/vpa_advection.jl
+++ b/moment_kinetics/src/vpa_advection.jl
@@ -150,19 +150,21 @@ function update_speed_n_u_p_evolution!(advect, fvec, moments, vpa, z, r, composi
         end
     end
     if ion_source_settings.active
-        source_amplitude = moments.ion.external_source_amplitude
-        source_T = ion_source_settings.source_T
+        source_density_amplitude = moments.ion.external_source_density_amplitude
+        source_momentum_amplitude = moments.ion.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.ion.external_source_pressure_amplitude
         density = fvec.density
         upar = fvec.upar
         ppar = fvec.ppar
         vth = moments.ion.vth
         vpa_grid = vpa.grid
         @loop_s_r_z is ir iz begin
-            prefactor = source_amplitude[iz,ir]
-            term1 = prefactor * upar[iz,ir,is]/(density[iz,ir,is]*vth[iz,ir,is])
+            term1 = source_density_amplitude[iz,ir] * upar[iz,ir,is]/(density[iz,ir,is]*vth[iz,ir,is])
             term2_over_vpa =
-                0.5 * prefactor * (-(0.5*source_T + upar[iz,ir,is]^2) / ppar[iz,ir,is]
-                                   + 1.0/density[iz,ir,is])
+                -0.5 * (source_pressure_amplitude[iz,ir] +
+                        2.0 * upar[iz,ir,is] * source_momentum_amplitude[iz,ir]) /
+                       ppar[iz,ir,is] +
+                0.5 * source_density_amplitude[iz,ir] / density[iz,ir,is]
             @loop_vperp_vpa ivperp ivpa begin
                 advect[is].speed[ivpa,ivperp,iz,ir] += term1 + vpa_grid[ivpa] * term2_over_vpa
             end
@@ -253,16 +255,15 @@ function update_speed_n_u_evolution!(advect, fvec, moments, vpa, z, r, compositi
         end
     end
     if ion_source_settings.active
-        source_amplitude = moments.ion.external_source_amplitude
+        source_density_amplitude = moments.ion.external_source_density_amplitude
         source_strength = ion_source_settings.source_strength
-        source_T = ion_source_settings.source_T
         r_amplitude = ion_source_settings.r_amplitude
         z_amplitude = ion_source_settings.z_amplitude
         density = fvec.density
         upar = fvec.upar
         vth = moments.ion.vth
         @loop_s_r_z is ir iz begin
-            term = source_amplitude[iz,ir] * upar[iz,ir,is] / density[iz,ir,is]
+            term = source_density_amplitude[iz,ir] * upar[iz,ir,is] / density[iz,ir,is]
             @loop_vperp_vpa ivperp ivpa begin
                 advect[is].speed[ivpa,ivperp,iz,ir] += term
             end

From d68cd4d82032560bc77416a3e39cd755f56d90f2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 27 Mar 2024 22:13:10 +0000
Subject: [PATCH 187/428] Fix typo in docs for source term in neutral
 moment-kinetic equation

The corrected form is similar to the ion one, and has correct
dimensions.
---
 docs/src/moment_kinetic_equations.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/moment_kinetic_equations.md b/docs/src/moment_kinetic_equations.md
index c21735b04..5f5f81a40 100644
--- a/docs/src/moment_kinetic_equations.md
+++ b/docs/src/moment_kinetic_equations.md
@@ -1018,7 +1018,7 @@ and for neutrals where several of the ionization terms cancel
   & \qquad-\frac{w_{\|,n}}{2}\frac{1}{p_{\|,n}}\left(-\frac{\partial q_{\|,n}}{\partial z}
   - R_{in}\left(n_{i}p_{\|,n} - n_{n}p_{\|,i}
   - n_{n}n_{i}\left(u_{n} - u_{i}\right)^{2}\right)
-  + \int dv_\parallel S_{n} + u_{n}^2\int dv_\parallel v_\parallel^2 S_{n}
+  + \int dv_\parallel v_\parallel^2 S_{n} + u_{n}^2\int dv_\parallel S_{n}
   + v_{\mathrm{th},n}w_{\|,n}\frac{\partial p_{\|,n}}{\partial z}\right) \\
   & \qquad\left. + \frac{w_{\parallel,n}}{2}\frac{1}{n_{n}}\int dv_\parallel S_{n}
   + \frac{w_{\|,n}^{2}}{2}\frac{v_{\mathrm{th},n}}{n_{n}}\frac{\partial n_{n}}{\partial z}\right]\frac{\partial g_{n}}{\partial w_{\|,n}} \\

From 14c0f2d9ee1f83ca288214ef329fe0308e706d42 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 18:15:22 +0000
Subject: [PATCH 188/428] Implement external sources for electrons in the same
 way as for ions

Replaces the hard-coded `heat_source` term.
---
 .../src/electron_fluid_equations.jl           |  18 +-
 .../src/electron_kinetic_equation.jl          | 105 ++++++---
 moment_kinetics/src/electron_vpa_advection.jl |  31 ++-
 moment_kinetics/src/external_sources.jl       | 214 +++++++++++++++++-
 moment_kinetics/src/file_io.jl                |  90 +++++++-
 moment_kinetics/src/initial_conditions.jl     |   4 +-
 moment_kinetics/src/load_data.jl              |  27 ++-
 .../src/moment_kinetics_structs.jl            |  12 +-
 moment_kinetics/src/time_advance.jl           |   2 +-
 moment_kinetics/src/velocity_moments.jl       |  16 +-
 10 files changed, 442 insertions(+), 77 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 7be4441d6..9c55a540e 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -121,7 +121,7 @@ an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron_upar,
                                    ion_upar, ion_ppar, density_neutral, uz_neutral,
                                    pz_neutral, moments, collisions, dt, composition,
-                                   num_diss_params, z)
+                                   electron_source_settings, num_diss_params, z)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
     me_over_mi = composition.me_over_mi
@@ -182,16 +182,14 @@ function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron
             end
         end
     end
-    # calculate the external electron heat source, if any
-    calculate_electron_heat_source!(moments.heat_source, ppar_in, moments.dupar_dz,
-                                    density_neutral, collisions.ionization, collisions.ionization_energy,
-                                    electron_density, ion_ppar, collisions.nu_ei, composition.me_over_mi,
-                                    composition.T_wall, z)
-    # add the contribution from the electron heat source                                
-    begin_r_z_region()
-    @loop_r_z ir iz begin
-        ppar_out[iz,ir] += dt * moments.heat_source[iz,ir]
+
+    if electron_source_settings.active
+        source_amplitude = moments.external_source_pressure_amplitude
+        @loop_r_z ir iz begin
+            ppar_out[iz,ir] += dt * source_amplitude[iz,ir]
+        end
     end
+
     return nothing
 end
 
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index fd36f0234..fd589fcdb 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -20,6 +20,7 @@ using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
 using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
+using ..external_sources: external_electron_source!
 using ..file_io: write_initial_electron_state, finish_initial_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
@@ -58,8 +59,9 @@ OUTPUT:
 function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
         phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
-        collisions, composition, num_diss_params, max_electron_pdf_iterations;
-        io_initial_electron=nothing, initial_time=0.0, evolve_ppar=false)
+        collisions, composition, external_source_settings, num_diss_params,
+        max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
+        evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -69,9 +71,10 @@ function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpa
     if solution_method == "artificial_time_derivative"
         return update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi,
             collisions, composition, r, z, vperp, vpa, z_spectral, vperp_spectral,
-            vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params, num_diss_params,
-            max_electron_pdf_iterations; io_initial_electron=io_initial_electron,
-            initial_time=initial_time, evolve_ppar=evolve_ppar)
+            vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
+            external_source_settings, num_diss_params, max_electron_pdf_iterations;
+            io_initial_electron=io_initial_electron, initial_time=initial_time,
+            evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
         return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
             qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
@@ -115,8 +118,9 @@ OUTPUT:
 """
 function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, collisions,
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
-        vpa_advect, scratch_dummy, t_params, num_diss_params, max_electron_pdf_iterations;
-        io_initial_electron=nothing, initial_time=0.0, evolve_ppar=false)
+        vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
+        max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
+        evolve_ppar=false)
 
     begin_r_z_region()
 
@@ -250,8 +254,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                                     moments, z, vperp, vpa, z_spectral,
                                                     vpa_spectral, z_advect, vpa_advect,
                                                     scratch_dummy, collisions,
-                                                    composition, num_diss_params,
-                                                    t_params.dt[];
+                                                    composition, external_source_settings,
+                                                    num_diss_params, t_params.dt[];
                                                     evolve_ppar=evolve_ppar)
             speedup_hack!(scratch[istage+1], scratch[istage], z_speedup_fac, z, vpa;
                           evolve_ppar=evolve_ppar)
@@ -356,7 +360,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
             if t_params.adaptive && istage == t_params.n_rk_stages
                 electron_adaptive_timestep_update!(scratch, time, t_params, moments,
-                                                   z_advect, vpa_advect, r, z, vperp, vpa;
+                                                   z_advect, vpa_advect, r, z, vperp, vpa,
+                                                   external_source_settings;
                                                    evolve_ppar=evolve_ppar)
                 # Re-do this in case electron_adaptive_timestep_update!() re-arranged the
                 # `scratch` vector
@@ -1143,8 +1148,8 @@ Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
 function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_advect,
-                                            vpa_advect, r, z, vperp, vpa;
-                                            evolve_ppar=false)
+                                            vpa_advect, r, z, vperp, vpa,
+                                            external_source_settings; evolve_ppar=false)
     #error_norm_method = "Linf"
     error_norm_method = "L2"
 
@@ -1181,9 +1186,9 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     end
 
     # vpa-advection
-    update_electron_speed_vpa!(vpa_advect[1], moments.electron.ppar, moments.electron.vth,
-                               moments.electron.dppar_dz, moments.electron.dqpar_dz,
-                               moments.electron.dvth_dz, vpa.grid)
+    update_electron_speed_vpa!(vpa_advect[1], moments.electron.dens,
+                               moments.electron.upar, moments.electron.ppar,
+                               moments, vpa.grid, external_source_settings.electron)
     vpa_CFL = get_minimum_CFL_vpa(vpa_advect[1].speed, vpa)
     if block_rank[] == 0
         push!(CFL_limits, t_params.CFL_prefactor * vpa_CFL)
@@ -1262,7 +1267,7 @@ function update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
     # get critical velocities beyond which electrons are lost to the wall
     crit_speed_zmin, crit_speed_zmax = get_electron_critical_velocities(phi, vthe, composition.me_over_mi, z)
     # add the contribution to rhs from the term proportional to the pdf (rather than its derivatives)
-    add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+    add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, upar, dvth_dz, dqpar_dz, vpa.grid, z, external_source_settings.electron)
     # add the contribution to rhs from the wpa advection term
     add_contribution_from_wpa_advection!(rhs, pdf, vthe, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
     # shoot in z from incoming boundary (using sign of zdot to determine direction)
@@ -1381,7 +1386,7 @@ function update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens
         # initialise the RHS to zero
         rhs .= 0.0
         # add the contribution to rhs from the term proportional to the pdf (rather than its derivatives)
-        add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, dvth_dz, dqpar_dz, vpa.grid, z)
+        add_contribution_from_pdf_term!(rhs, pdf, ppar, vthe, dens, ddens_dz, upar, dvth_dz, dqpar_dz, vpa.grid, z, external_source_settings.electron)
         # add the contribution to rhs from the wpa advection term
         #add_contribution_from_wpa_advection!(rhs, pdf, vthe, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
         # loop over wpa locations, solving the linear system at each location;
@@ -1438,8 +1443,8 @@ When `evolve_ppar=true` is passed, also updates the electron parallel pressure.
 function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z, vperp,
                                                  vpa, z_spectral, vpa_spectral, z_advect,
                                                  vpa_advect, scratch_dummy, collisions,
-                                                 composition, num_diss_params, dt;
-                                                 evolve_ppar=false)
+                                                 composition, external_source_settings,
+                                                 num_diss_params, dt; evolve_ppar=false)
     if evolve_ppar
         ppar = fvec_in.electron_ppar
     else
@@ -1451,16 +1456,15 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
                           vpa.grid, z_spectral, scratch_dummy, dt)
 
     # add the contribution from the wpa advection term
-    electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron, ppar,
-                            moments.electron.vth, moments.electron.dppar_dz,
-                            moments.electron.dqpar_dz, moments.electron.dvth_dz,
-                            vpa_advect, vpa, vpa_spectral, scratch_dummy, dt)
+    electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                            moments.electron.dens, moments.electron.upar, ppar,
+                            moments, vpa_advect, vpa, vpa_spectral, scratch_dummy, dt,
+                            external_source_settings.electron)
 
     # add in the contribution to the residual from the term proportional to the pdf
     add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, ppar,
-                                    moments.electron.vth, moments.electron.dens,
-                                    moments.electron.ddens_dz, moments.electron.dvth_dz,
-                                    moments.electron.dqpar_dz, vpa.grid, z, dt)
+                                    moments.electron.dens, moments.electron.upar, moments,
+                                    vpa.grid, z, dt, external_source_settings.electron)
 
     # add in numerical dissipation terms
     add_dissipation_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, scratch_dummy,
@@ -1476,13 +1480,20 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
                                    vperp, vpa, dt)
     end
 
+    if external_source_settings.electron.active
+        external_electron_source!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                                  moments.electron.dens, moments.electron.upar, moments,
+                                  external_source_settings.electron, vperp, vpa, dt)
+    end
+
     if evolve_ppar
         electron_energy_equation!(fvec_out.electron_ppar, fvec_in.electron_ppar,
                                   moments.electron.dens, moments.electron.upar,
                                   moments.ion.upar, moments.ion.ppar,
                                   moments.neutral.dens, moments.neutral.uz,
                                   moments.neutral.pz, moments.electron, collisions, dt,
-                                  composition, num_diss_params, z)
+                                  composition, external_source_settings.electron,
+                                  num_diss_params, z)
     end
 
     return nothing
@@ -1498,7 +1509,8 @@ OUTPUT:
 function electron_kinetic_equation_residual!(residual, max_term, single_term, pdf, dens, upar, vth, ppar, upar_ion,
                                              ddens_dz, dppar_dz, dqpar_dz, dvth_dz, 
                                              z, vperp, vpa, z_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                             collisions, num_diss_params, dt_electron)
+                                             collisions, external_source_settings,
+                                             num_diss_params, dt_electron)
 
     # initialise the residual to zero                                             
     begin_r_vperp_vpa_region()
@@ -1515,7 +1527,8 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #calculate_contribution_from_z_advection!(residual, pdf, vth, z, vpa.grid, z_spectral, scratch_dummy)
     # add in the contribution to the residual from the wpa advection term
     electron_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, 
-                            vpa_advect, vpa, vpa_spectral, scratch_dummy, -1.0)
+                            vpa_advect, vpa, vpa_spectral, scratch_dummy, -1.0,
+                            external_source_settings.electron)
     #dt_max_vadv = simple_vpa_advection!(residual, pdf, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa, dt_electron)
     #@. single_term = residual - single_term
     #max_term .= max.(max_term, abs.(single_term))
@@ -1523,8 +1536,8 @@ function electron_kinetic_equation_residual!(residual, max_term, single_term, pd
     #println("v_adv residual = ", maximum(abs.(single_term)))
     #add_contribution_from_wpa_advection!(residual, pdf, vth, ppar, dppar_dz, dqpar_dz, dvth_dz, vpa, vpa_spectral)
     # add in the contribution to the residual from the term proportional to the pdf
-    add_contribution_from_pdf_term!(residual, pdf, ppar, vth, dens, ddens_dz, dvth_dz,
-                                    dqpar_dz, vpa.grid, z, -1.0)
+    add_contribution_from_pdf_term!(residual, pdf, ppar, dens, moments, vpa.grid, z, -1.0,
+                                    external_source_settings.electron)
     #@. single_term = residual - single_term
     #max_term .= max.(max_term, abs.(single_term))
     #@. single_term = residual
@@ -1654,7 +1667,7 @@ function add_source_term!(source_term, vpa, z, dvth_dz)
     return nothing
 end
 
-function add_dissipation_term!(pdf_out, pdf, scratch_dummy, z_spectral, z, vpa,
+function add_dissipation_term!(pdf_out, pdf_in, scratch_dummy, z_spectral, z, vpa,
                                vpa_spectral, num_diss_params, dt)
     dummy_zr1 = @view scratch_dummy.dummy_zrs[:,:,1]
     dummy_zr2 = @view scratch_dummy.buffer_vpavperpzr_1[1,1,:,:]
@@ -1664,7 +1677,7 @@ function add_dissipation_term!(pdf_out, pdf, scratch_dummy, z_spectral, z, vpa,
     buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
     # add in numerical dissipation terms
     #@loop_vperp_vpa ivperp ivpa begin
-    #    @views derivative_z!(dummy_zr1, pdf[ivpa,ivperp,:,:], buffer_r_1, buffer_r_2, buffer_r_3,
+    #    @views derivative_z!(dummy_zr1, pdf_in[ivpa,ivperp,:,:], buffer_r_1, buffer_r_2, buffer_r_3,
     #                         buffer_r_4, z_spectral, z)
     #    @views derivative_z!(dummy_zr2, dummy_zr1, buffer_r_1, buffer_r_2, buffer_r_3,
     #                         buffer_r_4, z_spectral, z)
@@ -1672,10 +1685,10 @@ function add_dissipation_term!(pdf_out, pdf, scratch_dummy, z_spectral, z, vpa,
     #end
     begin_r_z_vperp_region()
     @loop_r_z_vperp ir iz ivperp begin
-        #@views derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, false)
+        #@views derivative!(vpa.scratch, pdf_in[:,ivperp,iz,ir], vpa, false)
         #@views derivative!(vpa.scratch2, vpa.scratch, vpa, false)
         #@. residual[:,ivperp,iz,ir] -= num_diss_params.vpa_dissipation_coefficient * vpa.scratch2
-        @views second_derivative!(vpa.scratch, pdf[:,ivperp,iz,ir], vpa, vpa_spectral)
+        @views second_derivative!(vpa.scratch, pdf_in[:,ivperp,iz,ir], vpa, vpa_spectral)
         @. pdf_out[:,ivperp,iz,ir] += dt * num_diss_params.vpa_dissipation_coefficient * vpa.scratch
     end
     #stop()
@@ -1880,7 +1893,12 @@ function calculate_pdf_dot_prefactor!(pdf_dot_prefactor, ppar, vth, dens, ddens_
 end
 
 # add contribution to the residual coming from the term proporational to the pdf
-function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, vth, dens, ddens_dz, dvth_dz, dqpar_dz, vpa, z, dt)
+function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, dens, upar, moments, vpa,
+                                         z, dt, electron_source_settings)
+    vth = moments.electron.vth
+    ddens_dz = moments.electron.ddens_dz
+    dvth_dz = moments.electron.dvth_dz
+    dqpar_dz = moments.electron.dqpar_dz
     begin_r_z_vperp_vpa_region()
     @loop_r_z ir iz begin
         this_dqpar_dz = dqpar_dz[iz,ir]
@@ -1898,6 +1916,21 @@ function add_contribution_from_pdf_term!(pdf_out, pdf_in, ppar, vth, dens, ddens
             #pdf_out[ivpa, ivperp, :, :] -= (-0.5 * dqpar_dz[:, :] / ppar[:, :]) * pdf_in[ivpa, ivperp, :, :]
         end
     end
+
+    if electron_source_settings.active
+        source_density_amplitude = moments.electron.external_source_density_amplitude
+        source_momentum_amplitude = moments.electron.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.electron.external_source_pressure_amplitude
+        @loop_r_z ir iz begin
+            term = dt * (1.5 * source_density_amplitude[iz,ir] / dens[iz,ir] -
+                         (0.5 * source_pressure_amplitude[iz,ir] +
+                          source_momentum_amplitude[iz,ir]) / ppar[iz,ir])
+            @loop_vperp_vpa ivperp ivpa begin
+                pdf_out[ivpa,ivperp,iz,ir] -= term * pdf_in[ivpa,ivperp,iz,ir]
+            end
+        end
+    end
+
     return nothing
 end
 
diff --git a/moment_kinetics/src/electron_vpa_advection.jl b/moment_kinetics/src/electron_vpa_advection.jl
index c41e884e4..2d7968286 100644
--- a/moment_kinetics/src/electron_vpa_advection.jl
+++ b/moment_kinetics/src/electron_vpa_advection.jl
@@ -12,8 +12,9 @@ using ..calculus: derivative!, second_derivative!
 calculate the wpa-advection term for the electron kinetic equation 
 = (vthe / 2 ppare * dppare/dz + wpa / 2 ppare * dqpare/dz - wpa^2 * dvthe/dz) * df/dwpa
 """
-function electron_vpa_advection!(pdf_out, pdf_in, ppar, vth, dppar_dz, dqpar_dz, dvth_dz,
-                                 advect, vpa, spectral, scratch_dummy, dt)
+function electron_vpa_advection!(pdf_out, pdf_in, density, upar, ppar, moments, advect,
+                                 vpa, spectral, scratch_dummy, dt,
+                                 electron_source_settings)
     begin_r_z_vperp_region()
 
     # create a reference to a scratch_dummy array to store the wpa-derivative of the electron pdf
@@ -21,7 +22,8 @@ function electron_vpa_advection!(pdf_out, pdf_in, ppar, vth, dppar_dz, dqpar_dz,
     #d2pdf_dvpa2 = scratch_dummy.buffer_vpavperpzr_2
     begin_r_z_vperp_region()
     # get the updated speed along the wpa direction using the current pdf
-    @views update_electron_speed_vpa!(advect[1], ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa.grid)
+    @views update_electron_speed_vpa!(advect[1], density, upar, ppar, moments, vpa.grid,
+                                      electron_source_settings)
     # update adv_fac -- note that there is no factor of dt here because
     # in some cases the electron kinetic equation is solved as a steady-state equation iteratively
     @loop_r_z_vperp ir iz ivperp begin
@@ -51,7 +53,12 @@ end
 """
 calculate the electron advection speed in the wpa-direction at each grid point
 """
-function update_electron_speed_vpa!(advect, ppar, vth, dppar_dz, dqpar_dz, dvth_dz, vpa)
+function update_electron_speed_vpa!(advect, density, upar, ppar, moments, vpa,
+                                    electron_source_settings)
+    vth = moments.electron.vth
+    dppar_dz = moments.electron.dppar_dz
+    dqpar_dz = moments.electron.dqpar_dz
+    dvth_dz = moments.electron.dvth_dz
     # calculate the advection speed in wpa
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
         # TMP FOR TESTING
@@ -59,6 +66,22 @@ function update_electron_speed_vpa!(advect, ppar, vth, dppar_dz, dqpar_dz, dvth_
         advect.speed[ivpa,ivperp,iz,ir] = ((vth[iz,ir] * dppar_dz[iz,ir] + vpa[ivpa] * dqpar_dz[iz,ir]) 
                                            / (2 * ppar[iz,ir]) - vpa[ivpa]^2 * dvth_dz[iz,ir])
     end
+    if electron_source_settings.active
+        source_density_amplitude = moments.electron.external_source_density_amplitude
+        source_momentum_amplitude = moments.electron.external_source_momentum_amplitude
+        source_pressure_amplitude = moments.electron.external_source_pressure_amplitude
+        @loop_r_z ir iz begin
+            term1 = source_density_amplitude[iz,ir] * upar[iz,ir]/(density[iz,ir]*vth[iz,ir])
+            term2_over_vpa =
+                -0.5 * (source_pressure_amplitude[iz,ir] +
+                        2.0 * upar[iz,ir] * source_momentum_amplitude[iz,ir]) /
+                       ppar[iz,ir] +
+                0.5 * source_density_amplitude[iz,ir] / density[iz,ir]
+            @loop_vperp_vpa ivperp ivpa begin
+                advect.speed[ivpa,ivperp,iz,ir] += term1 + vpa[ivpa] * term2_over_vpa
+            end
+        end
+    end
     return nothing
 end
 
diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 8ee1ae980..3059b5570 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -192,7 +192,38 @@ function setup_external_sources!(input_dict, r, z)
                 PI_density_target_rank=PI_density_target_rank)
     end
 
-    return (ion=get_settings(false), neutral=get_settings(true))
+    function get_electron_settings(ion_settings)
+        # Note most settings for the electron source are copied from the ion source,
+        # because we require that the particle sources are the same for ions and
+        # electrons. `source_T` can be set independently, and when using
+        # `source_type="energy"`, the `source_strength` could also be set.
+        input = set_defaults_and_check_section!(
+                     input_dict, "electron_source";
+                     source_strength=ion_settings.source_strength,
+                     source_T=ion_settings.source_T,
+                    )
+        if ion_settings.source_type != "energy"
+            # Need to keep same amplitude for ions and electrons so there is no charge
+            # source.
+            if input["source_strength"] != ion_settings.source_strength
+                println("When not using source_type=\"energy\", source_strength for "
+                        * "electrons must be equal to source_strength for ions to ensure "
+                        * "no charge is injected by the source. Overriding electron "
+                        * "source_strength...")
+            end
+            input["source_strength"] = ion_settings.source_strength
+        end
+        return (; (Symbol(k)=>v for (k,v) ∈ input)..., active=ion_settings.active,
+                r_amplitude=ion_settings.r_amplitude,
+                z_amplitude=ion_settings.z_amplitude,
+                source_type=ion_settings.source_type)
+    end
+
+    ion_settings = get_settings(false)
+    electron_settings = get_electron_settings(ion_settings)
+    neutral_settings = get_settings(true)
+
+    return (ion=ion_settings, electron=electron_settings, neutral=neutral_settings)
 end
 
 """
@@ -234,6 +265,10 @@ Initialize the arrays `moments.ion.external_source_amplitude`,
 `moments.ion.external_source_density_amplitude`,
 `moments.ion.external_source_momentum_amplitude`,
 `moments.ion.external_source_pressure_amplitude`,
+`moments.electron.external_source_amplitude`,
+`moments.electron.external_source_density_amplitude`,
+`moments.electron.external_source_momentum_amplitude`,
+`moments.electron.external_source_pressure_amplitude`,
 `moments.neutral.external_source_amplitude`,
 `moments.neutral.external_source_density_amplitude`,
 `moments.neutral.external_source_momentum_amplitude`, and
@@ -308,6 +343,67 @@ function initialize_external_source_amplitude!(moments, external_source_settings
         end
     end
 
+    electron_source_settings = external_source_settings.electron
+    if electron_source_settings.active
+        if electron_source_settings.source_type == "energy"
+            @loop_r_z ir iz begin
+                moments.electron.external_source_amplitude[iz,ir] =
+                    electron_source_settings.source_strength *
+                    electron_source_settings.r_amplitude[ir] *
+                    electron_source_settings.z_amplitude[iz]
+            end
+            @loop_r_z ir iz begin
+                moments.electron.external_source_density_amplitude[iz,ir] = 0.0
+            end
+            @loop_r_z ir iz begin
+                moments.electron.external_source_momentum_amplitude[iz,ir] =
+                    - moments.electron.dens[iz,ir] * moments.electron.upar[iz,ir] *
+                      electron_source_settings.source_strength *
+                      electron_source_settings.r_amplitude[ir] *
+                      electron_source_settings.z_amplitude[iz]
+            end
+            @loop_r_z ir iz begin
+                moments.electron.external_source_pressure_amplitude[iz,ir] =
+                    (0.5 * electron_source_settings.source_T +
+                     moments.electron.upar[iz,ir]^2 - moments.electron.ppar[iz,ir]) *
+                    electron_source_settings.source_strength *
+                    electron_source_settings.r_amplitude[ir] *
+                    electron_source_settings.z_amplitude[iz]
+            end
+        else
+            @loop_r_z ir iz begin
+                moments.electron.external_source_amplitude[iz,ir] =
+                    moments.ion.external_source_amplitude[iz,ir]
+            end
+            if moments.evolve_density
+                @loop_r_z ir iz begin
+                    moments.electron.external_source_density_amplitude[iz,ir] =
+                        moments.ion.external_source_density_amplitude[iz,ir]
+                end
+            else
+                @loop_r_z ir iz begin
+                    # Note set this using *ion* settings to force electron density source
+                    # to always be equal to ion density source (even when
+                    # evolve_density=false) to ensure the source does not inject charge
+                    # into the simulation.
+                    moments.electron.external_source_density_amplitude[iz,ir] =
+                        ion_source_settings.source_strength *
+                        ion_source_settings.r_amplitude[ir] *
+                        ion_source_settings.z_amplitude[iz]
+                end
+            end
+            @loop_r_z ir iz begin
+                moments.electron.external_source_momentum_amplitude[iz,ir] = 0.0
+            end
+            @loop_r_z ir iz begin
+                moments.electron.external_source_pressure_amplitude[iz,ir] =
+                    (0.5 * electron_source_settings.source_T +
+                     moments.electron.upar[iz,ir]^2) *
+                    moments.electron.external_source_amplitude[iz,ir]
+            end
+        end
+    end
+
     if n_neutral_species > 0
         neutral_source_settings = external_source_settings.neutral
         if neutral_source_settings.active
@@ -504,6 +600,54 @@ function external_ion_source!(pdf, fvec, moments, ion_source_settings, vperp, vp
     return nothing
 end
 
+"""
+    external_electron_source!(pdf, fvec, moments, electron_source_settings, vperp,
+                              vpa, dt)
+
+Add external source term to the electron kinetic equation.
+"""
+function external_electron_source!(pdf_out, pdf_in, electron_density, electron_upar,
+                                   moments, electron_source_settings, vperp, vpa, dt)
+    begin_r_z_vperp_region()
+
+    source_amplitude = moments.electron.external_source_amplitude
+    source_T = electron_source_settings.source_T
+    if vperp.n == 1
+        vth_factor = 1.0 / sqrt(source_T)
+    else
+        vth_factor = 1.0 / source_T^1.5
+    end
+    vpa_grid = vpa.grid
+    vperp_grid = vperp.grid
+
+    vth = moments.electron.vth
+    @loop_r_z ir iz begin
+        this_vth = vth[iz,ir]
+        this_upar = electron_upar[iz,ir]
+        this_prefactor = dt * this_vth / electron_density[iz,ir] * vth_factor *
+                         source_amplitude[iz,ir]
+        @loop_vperp_vpa ivperp ivpa begin
+            # Factor of 1/sqrt(π) (for 1V) or 1/π^(3/2) (for 2V/3V) is absorbed by the
+            # normalisation of F
+            vperp_unnorm = vperp_grid[ivperp] * this_vth
+            vpa_unnorm = vpa_grid[ivpa] * this_vth + this_upar
+            pdf_out[ivpa,ivperp,iz,ir] +=
+                this_prefactor *
+                exp(-(vperp_unnorm^2 + vpa_unnorm^2) / source_T)
+        end
+    end
+
+    if electron_source_settings.source_type == "energy"
+        # Take particles out of pdf so source does not change density
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            pdf_out[ivpa,ivperp,iz,ir] -= dt * source_amplitude[iz,ir] *
+                                             pdf_in[ivpa,ivperp,iz,ir]
+        end
+    end
+
+    return nothing
+end
+
 """
     external_neutral_source!(pdf, fvec, moments, neutral_source_settings, vzeta, vr,
                             vz, dt)
@@ -723,6 +867,74 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
     return nothing
 end
 
+"""
+    external_electron_source_controller!(fvec_in, moments, electron_source_settings, dt)
+
+Calculate the amplitude, e.g. when using a PI controller for the density to set the
+external source amplitude.
+
+As the electron density source must be equal to the ion density source in order not to
+inject charge into the simulation, the electron source (at least in some modes of
+operation) depends on the ion source, so [`external_ion_source_controller`](@ref) must be
+called before this function is called so that `moments.ion.external_source_amplitude` is
+up to date.
+"""
+function external_electron_source_controller!(fvec_in, moments, electron_source_settings,
+                                              dt)
+
+    is = 1
+    electron_moments = moments.electron
+    ion_source_amplitude = moments.ion.external_source_amplitude
+
+    if electron_source_settings.source_type == "Maxwellian"
+        @loop_r_z ir iz begin
+            electron_moments.external_source_pressure_amplitude[iz,ir] =
+                (0.5 * electron_source_settings.source_T +
+                 fvec_in.electron_upar[iz,ir,is]^2) *
+                electron_moments.external_source_amplitude[iz,ir]
+        end
+    elseif electron_source_settings.source_type == "energy"
+        @loop_r_z ir iz begin
+            electron_moments.external_source_momentum_amplitude[iz,ir] =
+                - electron_moments.density[iz,ir] * electron_moments.upar[iz,ir] *
+                  electron_source_settings.source_strength *
+                  electron_source_settings.r_amplitude[ir] *
+                  electron_source_settings.z_amplitude[iz]
+        end
+        @loop_r_z ir iz begin
+            electron_moments.external_source_pressure_amplitude[iz,ir] =
+                (0.5 * electron_source_settings.source_T + electron_moments.upar[iz,ir]^2 -
+                 electron_moments.ppar[iz,ir]) *
+                electron_source_settings.source_strength *
+                electron_source_settings.r_amplitude[ir] *
+                electron_source_settings.z_amplitude[iz]
+        end
+    else
+        @loop_r_z ir iz begin
+            electron_moments.external_source_amplitude[iz,ir] = ion_source_amplitude[iz,ir]
+        end
+        @loop_r_z ir iz begin
+            electron_moments.external_source_momentum_amplitude[iz,ir] =
+                - electron_moments.density[iz,ir] * electron_moments.upar[iz,ir] *
+                  electron_moments.external_source_amplitude[iz,ir]
+        end
+        @loop_r_z ir iz begin
+            electron_moments.external_source_pressure_amplitude[iz,ir] =
+                (0.5 * electron_source_settings.source_T + electron_moments.upar[iz,ir]^2 -
+                 electron_moments.ppar[iz,ir]) *
+                electron_moments.external_source_amplitude[iz,ir]
+        end
+    end
+
+    # Density source is always the same as the ion one
+    @loop_r_z ir iz begin
+        electron_moments.external_source_density_amplitude[iz,ir] =
+            moments.ion.external_source_density_amplitude[iz,ir]
+    end
+
+    return nothing
+end
+
 """
     external_neutral_source_controller!(fvec_in, moments, neutral_source_settings, r,
                                         z, dt)
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 08877aa76..a70ca630e 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -55,9 +55,9 @@ moments & fields only
 """
 struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
                        Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
-                       Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Tconstri, Tconstrn,
-                       Tconstre, Tint, Tfailcause, Telectrontime, Telectronint,
-                       Telectronfailcause}
+                       Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Texte1, Texte2,
+                       Texte3, Texte4, Tconstri, Tconstrn, Tconstre, Tint, Tfailcause,
+                       Telectrontime, Telectronint, Telectronfailcause}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the time variable
@@ -115,6 +115,10 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     external_source_neutral_momentum_amplitude::Textn3
     external_source_neutral_pressure_amplitude::Textn4
     external_source_neutral_controller_integral::Textn5
+    external_source_electron_amplitude::Texte1
+    external_source_electron_density_amplitude::Texte2
+    external_source_electron_momentum_amplitude::Texte3
+    external_source_electron_pressure_amplitude::Texte4
 
     # handles for constraint coefficients
     ion_constraints_A_coefficient::Tconstri
@@ -185,8 +189,8 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr, Telectrontime,
-                                Telectronint, Telectronfailcause}
+struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4,
+                                Tconstr, Telectrontime, Telectronint, Telectronfailcause}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the pseudotime variable
@@ -203,6 +207,11 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Tconstr, Telectrontime,
     electron_parallel_heat_flux::Tmom
     # handle for the electron thermal speed variable
     electron_thermal_speed::Tmom
+    # handles for external source terms
+    external_source_electron_amplitude::Texte1
+    external_source_electron_density_amplitude::Texte2
+    external_source_electron_momentum_amplitude::Texte3
+    external_source_electron_pressure_amplitude::Texte4
     # handles for constraint coefficients
     electron_constraints_A_coefficient::Tconstr
     electron_constraints_B_coefficient::Tconstr
@@ -367,7 +376,12 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
                                                  description="electron distribution function")
 
         io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-        io_electron_vth, io_electron_step_counter, io_electron_dt,
+        io_electron_vth, external_source_electron_amplitude,
+        external_source_electron_density_amplitude,
+        external_source_electron_momentum_amplitude,
+        external_source_electron_pressure_amplitude,
+        electron_constraints_A_coefficient, electron_constraints_B_coefficient,
+        electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
         io_electron_failure_counter, io_electron_failure_caused_by,
         io_electron_limit_caused_by, io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
@@ -406,6 +420,10 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_parallel_pressure"),
                                         getvar("electron_parallel_heat_flux"),
                                         getvar("electron_thermal_speed"),
+                                        getvar("external_source_electron_amplitude"),
+                                        getvar("external_source_electron_density_amplitude"),
+                                        getvar("external_source_electron_momentum_amplitude"),
+                                        getvar("external_source_electron_pressure_amplitude"),
                                         getvar("electron_constraints_A_coefficient"),
                                         getvar("electron_constraints_B_coefficient"),
                                         getvar("electron_constraints_C_coefficient"),
@@ -846,11 +864,14 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                                  evolve_upar, evolve_ppar)
 
         io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-        io_electron_vth, electron_constraints_A_coefficient,
-        electron_constraints_B_coefficient, electron_constraints_C_coefficient,
-        io_electron_step_counter, io_electron_dt, io_electron_failure_counter,
-        io_electron_failure_caused_by, io_electron_limit_caused_by,
-        io_electron_dt_before_last_fail =
+        io_electron_vth, external_source_electron_amplitude,
+        external_source_electron_density_amplitude,
+        external_source_electron_momentum_amplitude,
+        external_source_electron_pressure_amplitude,
+        electron_constraints_A_coefficient, electron_constraints_B_coefficient,
+        electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
+        io_electron_failure_counter, io_electron_failure_caused_by,
+        io_electron_limit_caused_by, io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -925,6 +946,10 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                external_source_neutral_momentum_amplitude,
                                external_source_neutral_pressure_amplitude,
                                external_source_neutral_controller_integral,
+                               external_source_electron_amplitude,
+                               external_source_electron_density_amplitude,
+                               external_source_electron_momentum_amplitude,
+                               external_source_electron_pressure_amplitude,
                                ion_constraints_A_coefficient,
                                ion_constraints_B_coefficient,
                                ion_constraints_C_coefficient,
@@ -1177,6 +1202,23 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                       description="electron species thermal speed",
                                       units="c_ref")
 
+    external_source_electron_amplitude = create_dynamic_variable!(
+        dynamic, "external_source_electron_amplitude", mk_float, z, r;
+        parallel_io=parallel_io, description="Amplitude of the external source for electrons",
+        units="n_ref/c_ref^3*c_ref/L_ref")
+    external_source_electron_density_amplitude = create_dynamic_variable!(
+        dynamic, "external_source_electron_density_amplitude", mk_float, z, r;
+        parallel_io=parallel_io, description="Amplitude of the external density source for electrons",
+        units="n_ref*c_ref/L_ref")
+    external_source_electron_momentum_amplitude = create_dynamic_variable!(
+        dynamic, "external_source_electron_momentum_amplitude", mk_float, z, r;
+        parallel_io=parallel_io, description="Amplitude of the external momentum source for electrons",
+        units="m_ref*n_ref*c_ref*c_ref/L_ref")
+    external_source_electron_pressure_amplitude = create_dynamic_variable!(
+        dynamic, "external_source_electron_pressure_amplitude", mk_float, z, r;
+        parallel_io=parallel_io, description="Amplitude of the external pressure source for electrons",
+        units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
+
     electron_constraints_A_coefficient =
         create_dynamic_variable!(dynamic, "electron_constraints_A_coefficient", mk_float, z, r;
                                parallel_io=parallel_io,
@@ -1232,7 +1274,11 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
     end
 
     return io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-           io_electron_vth, electron_constraints_A_coefficient, electron_constraints_B_coefficient,
+           io_electron_vth, external_source_electron_amplitude,
+           external_source_electron_density_amplitude,
+           external_source_electron_momentum_amplitude,
+           external_source_electron_pressure_amplitude,
+           electron_constraints_A_coefficient, electron_constraints_B_coefficient,
            electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
            io_electron_failure_counter, io_electron_failure_caused_by,
            io_electron_limit_caused_by, io_electron_dt_before_last_fail
@@ -1546,6 +1592,10 @@ function reopen_moments_io(file_info)
                                getvar("external_source_neutral_momentum_amplitude"),
                                getvar("external_source_neutral_pressure_amplitude"),
                                getvar("external_source_neutral_controller_integral"),
+                               getvar("external_source_electron_amplitude"),
+                               getvar("external_source_electron_density_amplitude"),
+                               getvar("external_source_electron_momentum_amplitude"),
+                               getvar("external_source_electron_pressure_amplitude"),
                                getvar("ion_constraints_A_coefficient"),
                                getvar("ion_constraints_B_coefficient"),
                                getvar("ion_constraints_C_coefficient"),
@@ -1665,6 +1715,10 @@ function reopen_dfns_io(file_info)
                                      getvar("external_source_neutral_momentum_amplitude"),
                                      getvar("external_source_neutral_pressure_amplitude"),
                                      getvar("external_source_neutral_controller_integral"),
+                                     getvar("external_source_electron_amplitude"),
+                                     getvar("external_source_electron_density_amplitude"),
+                                     getvar("external_source_electron_momentum_amplitude"),
+                                     getvar("external_source_electron_pressure_amplitude"),
                                      getvar("ion_constraints_A_coefficient"),
                                      getvar("ion_constraints_B_coefficient"),
                                      getvar("ion_constraints_C_coefficient"),
@@ -1902,6 +1956,18 @@ function write_electron_moments_data_to_binary(moments, t_params,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
                               t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.external_source_electron_amplitude,
+                              moments.electron.external_source_amplitude, t_idx,
+                              parallel_io, z, r)
+        append_to_dynamic_var(io_moments.external_source_electron_density_amplitude,
+                              moments.electron.external_source_density_amplitude,
+                              t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.external_source_electron_momentum_amplitude,
+                              moments.electron.external_source_momentum_amplitude,
+                              t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.external_source_electron_pressure_amplitude,
+                              moments.electron.external_source_pressure_amplitude,
+                              t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_constraints_A_coefficient,
                               moments.electron.constraints_A_coefficient, t_idx,
                               parallel_io, z, r)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6d3880156..824742701 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -511,8 +511,8 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         moments.electron.dvth_dz, r, z, vperp, vpa,
                                         z_spectral, vperp_spectral, vpa_spectral,
                                         z_advect, vpa_advect, scratch_dummy, t_params,
-                                        collisions, composition, num_diss_params,
-                                        max_electron_pdf_iterations;
+                                        collisions, composition, external_source_settings,
+                                        num_diss_params, max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
                                         initial_time=code_time, evolve_ppar=true)
 
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index e764cadfe..86148dc32 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4103,11 +4103,24 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
     elseif variable_name == "electron_vpa_advect_speed"
         # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
         # to get_variable() in this case. Instead select a slice of the result.
+        density = get_variable(run_info, "electron_density")
+        upar = get_variable(run_info, "electron_parallel_flow")
         ppar = get_variable(run_info, "electron_parallel_pressure")
         vth = get_variable(run_info, "electron_thermal_speed")
         dppar_dz = get_z_derivative(run_info, "electron_parallel_pressure")
         dvth_dz = get_z_derivative(run_info, "electron_thermal_speed")
         dqpar_dz = get_z_derivative(run_info, "electron_parallel_heat_flux")
+        if run_info.external_source_settings.electron.active
+            external_source_amplitude = get_variable(run_info, "external_source_electron_amplitude")
+            external_source_density_amplitude = get_variable(run_info, "external_source_electron_density_amplitude")
+            external_source_momentum_amplitude = get_variable(run_info, "external_source_electron_momentum_amplitude")
+            external_source_pressure_amplitude = get_variable(run_info, "external_source_electron_pressure_amplitude")
+        else
+            external_source_amplitude = zeros(0,0,run_info.nt)
+            external_source_density_amplitude = zeros(0,0,run_info.nt)
+            external_source_momentum_amplitude = zeros(0,0,run_info.nt)
+            external_source_pressure_amplitude = zeros(0,0,run_info.nt)
+        end
 
         nz, nr, nt = size(vth)
         nvperp = run_info.vperp.n
@@ -4122,9 +4135,17 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             begin_serial_region()
             # Only need some struct with a 'speed' variable
             advect = (speed=@view(speed[:,:,:,:,it]),)
-            @views update_electron_speed_vpa!(advect, ppar[:,:,it], vth[:,:,it],
-                                              dppar_dz[:,:,it], dqpar_dz[:,:,it],
-                                              dvth_dz[:,:,it], run_info.vpa.grid)
+            moments = (electron=(vth=vth[:,:,it],
+                                 dppar_dz=dppar_dz[:,:,it],
+                                 dqpar_dz=dqpar_dz[:,:,it],
+                                 dvth_dz=dvth_dz[:,:,it],
+                                 external_source_amplitude=external_source_amplitude[:,:,it],
+                                 external_source_density_amplitude=external_source_density_amplitude[:,:,it],
+                                 external_source_momentum_amplitude=external_source_momentum_amplitude[:,:,it],
+                                 external_source_pressure_amplitude=external_source_pressure_amplitude[:,:,it]),)
+            @views update_electron_speed_vpa!(advect, density[:,:,it], upar[:,:,it],
+                                              ppar[:,:,it], moments, run_info.vpa.grid,
+                                              run_info.external_source_settings.electron)
         end
 
         variable = speed
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index 7fe440e40..ed3bf13b3 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -171,8 +171,16 @@ struct moments_electron_substruct
     vth::MPISharedArray{mk_float,2}
     # this is the parallel friction force between ions and electrons
     parallel_friction::MPISharedArray{mk_float,2}
-    # this is the electron heat source
-    heat_source::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the external source term
+    external_source_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the density moment of the external source term
+    external_source_density_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel momentum moment of the external source
+    # term
+    external_source_momentum_amplitude::MPISharedArray{mk_float,2}
+    # Spatially varying amplitude of the parallel pressure moment of the external source
+    # term
+    external_source_pressure_amplitude::MPISharedArray{mk_float,2}
     # if evolve_ppar = true, then the velocity variable is (vpa - upa)/vth, which introduces
     # a factor of vth for each power of wpa in velocity space integrals.
     # v_norm_fac accounts for this: it is vth if using the above definition for the parallel velocity,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 3ff3888ce..97bc064d2 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2478,7 +2478,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                   fvec_in.ppar, fvec_in.density_neutral,
                                   fvec_in.uz_neutral, fvec_in.pz_neutral,
                                   moments.electron, collisions, dt, composition,
-                                  num_diss_params, z)
+                                  external_source_settings.electron, num_diss_params, z)
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 35dd16802..f92c0cf0a 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -223,8 +223,11 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
     parallel_heat_flux_updated = Ref(false)
     # allocate array used for the election-ion parallel friction force
     parallel_friction_force = allocate_shared_float(nz, nr)
-    # allocate array used for electron heat source
-    heat_source = allocate_shared_float(nz, nr)
+    # allocate arrays used for external sources
+    external_source_amplitude = allocate_shared_float(nz, nr)
+    external_source_density_amplitude = allocate_shared_float(nz, nr)
+    external_source_momentum_amplitude = allocate_shared_float(nz, nr)
+    external_source_pressure_amplitude = allocate_shared_float(nz, nr)
     # allocate array used for the thermal speed
     thermal_speed = allocate_shared_float(nz, nr)
     # if evolving the electron pdf, it will be a function of the vth-normalised peculiar velocity
@@ -264,10 +267,11 @@ function create_moments_electron(nz, nr, electron_model, numerical_dissipation)
         parallel_flow_updated, parallel_pressure, parallel_pressure_updated,
         temperature, temperature_updated, 
         parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
-        parallel_friction_force, heat_source, v_norm_fac,
-        ddens_dz, dupar_dz, dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, 
-        dT_dz, dT_dz_upwind, dvth_dz, constraints_A_coefficient,
-        constraints_B_coefficient, constraints_C_coefficient)
+        parallel_friction_force, external_source_amplitude,
+        external_source_density_amplitude, external_source_momentum_amplitude,
+        external_source_pressure_amplitude, v_norm_fac, ddens_dz, dupar_dz, dppar_dz,
+        dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dT_dz, dT_dz_upwind, dvth_dz,
+        constraints_A_coefficient, constraints_B_coefficient, constraints_C_coefficient)
 end
 
 # neutral particles have natural mean velocities 

From fcc6ea870fe2cd1544df4234dd891e2abe931940 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 21:39:30 +0000
Subject: [PATCH 189/428] Simplify post-processing of 'per-step' from
 'cumulative' variable

Introduce utility function that can be re-used several times.
---
 moment_kinetics/src/load_data.jl | 59 ++++++++++----------------------
 1 file changed, 19 insertions(+), 40 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 86148dc32..5e360fac7 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3891,6 +3891,17 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         return variable
     end
 
+    # Get a 'per step' value from a saved 'cumulative' value. E.g. 'iterations per step'
+    # from a saved 'cumulative total iterations'
+    function get_per_step_from_cumulative_variable(run_info, varname::String; kwargs...)
+        variable = get_variable(run_info, varname; kwargs...)
+        tdim = ndims(variable)
+        for i ∈ size(variable, tdim):-1:2
+            selectdim(variable, tdim, i) .-= selectdim(variable, tdim, i-1)
+        end
+        return variable
+    end
+
     if variable_name == "temperature"
         vth = postproc_load_variable(run_info, "thermal_speed"; kwargs...)
         variable = vth.^2
@@ -4278,29 +4289,13 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         variable = speed
         variable = select_slice_of_variable(variable; kwargs...)
     elseif variable_name == "steps_per_output"
-        steps_per_output = get_variable(run_info, "step_counter"; kwargs...)
-        for i ∈ length(steps_per_output):-1:2
-            steps_per_output[i] -= steps_per_output[i-1]
-        end
-        variable = steps_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "step_counter"; kwargs...)
     elseif variable_name == "failures_per_output"
-        failures_per_output = get_variable(run_info, "failure_counter"; kwargs...)
-        for i ∈ length(failures_per_output):-1:2
-            failures_per_output[i] -= failures_per_output[i-1]
-        end
-        variable = failures_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "failure_counter"; kwargs...)
     elseif variable_name == "failure_caused_by_per_output"
-        failure_caused_by_per_output = get_variable(run_info, "failure_caused_by"; kwargs...)
-        for i ∈ size(failure_caused_by_per_output,2):-1:2
-            failure_caused_by_per_output[:,i] .-= failure_caused_by_per_output[:,i-1]
-        end
-        variable = failure_caused_by_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "failure_caused_by"; kwargs...)
     elseif variable_name == "limit_caused_by_per_output"
-        limit_caused_by_per_output = get_variable(run_info, "limit_caused_by"; kwargs...)
-        for i ∈ size(limit_caused_by_per_output,2):-1:2
-            limit_caused_by_per_output[:,i] .-= limit_caused_by_per_output[:,i-1]
-        end
-        variable = limit_caused_by_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "limit_caused_by"; kwargs...)
     elseif variable_name == "average_successful_dt"
         steps_per_output = get_variable(run_info, "steps_per_output"; kwargs...)
         failures_per_output = get_variable(run_info, "failures_per_output"; kwargs...)
@@ -4317,29 +4312,13 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             variable[1] = 0.0
         end
     elseif variable_name == "electron_steps_per_output"
-        electron_steps_per_output = get_variable(run_info, "electron_step_counter"; kwargs...)
-        for i ∈ length(electron_steps_per_output):-1:2
-            electron_steps_per_output[i] -= electron_steps_per_output[i-1]
-        end
-        variable = electron_steps_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "electron_step_counter"; kwargs...)
     elseif variable_name == "electron_failures_per_output"
-        electron_failures_per_output = get_variable(run_info, "electron_failure_counter"; kwargs...)
-        for i ∈ length(electron_failures_per_output):-1:2
-            electron_failures_per_output[i] -= electron_failures_per_output[i-1]
-        end
-        variable = electron_failures_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "electron_failure_counter"; kwargs...)
     elseif variable_name == "electron_failure_caused_by_per_output"
-        electron_failure_caused_by_per_output = get_variable(run_info, "electron_failure_caused_by"; kwargs...)
-        for i ∈ size(electron_failure_caused_by_per_output,2):-1:2
-            electron_failure_caused_by_per_output[:,i] .-= electron_failure_caused_by_per_output[:,i-1]
-        end
-        variable = electron_failure_caused_by_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "electron_failure_caused_by"; kwargs...)
     elseif variable_name == "electron_limit_caused_by_per_output"
-        electron_limit_caused_by_per_output = get_variable(run_info, "electron_limit_caused_by"; kwargs...)
-        for i ∈ size(electron_limit_caused_by_per_output,2):-1:2
-            electron_limit_caused_by_per_output[:,i] .-= electron_limit_caused_by_per_output[:,i-1]
-        end
-        variable = electron_limit_caused_by_per_output
+        variable = get_per_step_from_cumulative_variable(run_info, "electron_limit_caused_by"; kwargs...)
     elseif variable_name == "electron_average_successful_dt"
         electron_steps_per_output = get_variable(run_info, "electron_steps_per_output"; kwargs...)
         electron_failures_per_output = get_variable(run_info, "electron_failures_per_output"; kwargs...)

From 60889343abdaf8fcc746f36b1f905e83a6ef0b5e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 25 Mar 2024 21:44:13 +0000
Subject: [PATCH 190/428] Fix per-step values when post-processing restarted
 simulation

Fix is a bit hacky, as values at the restart points are just set to
zero.
---
 moment_kinetics/src/load_data.jl | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 5e360fac7..44b875ffc 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3899,6 +3899,16 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         for i ∈ size(variable, tdim):-1:2
             selectdim(variable, tdim, i) .-= selectdim(variable, tdim, i-1)
         end
+
+        # Per-step count does not make sense for the first step, so make sure element-1 is
+        # zero.
+        selectdim(variable, tdim, 1) .= zero(first(variable))
+
+        # Assume cumulative variables always increase, so if any value in the 'per-step'
+        # variable is negative, it is because there was a restart where the cumulative
+        # variable started over
+        variable .= max.(variable, zero(first(variable)))
+
         return variable
     end
 
@@ -4307,6 +4317,11 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         end
 
         variable = delta_t ./ successful_steps_per_output
+        for i ∈ eachindex(successful_steps_per_output)
+            if successful_steps_per_output[i] == 0
+                variable[i] = 0.0
+            end
+        end
         if successful_steps_per_output[1] == 0
             # Don't want a meaningless Inf...
             variable[1] = 0.0
@@ -4330,6 +4345,11 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         end
 
         variable = delta_t ./ electron_successful_steps_per_output
+        for i ∈ eachindex(electron_successful_steps_per_output)
+            if electron_successful_steps_per_output[i] == 0
+                variable[i] = 0.0
+            end
+        end
         if electron_successful_steps_per_output[1] == 0
             # Don't want a meaningless Inf...
             variable[1] = 0.0

From d04c6d85abc008070083cbea68022396ec72e258 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 26 Mar 2024 09:02:30 +0000
Subject: [PATCH 191/428] Better error messages when `iv0` cannot be found

---
 moment_kinetics/src/electron_kinetic_equation.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index fd589fcdb..46179e203 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -721,6 +721,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
         phi[1,ir] = me_over_mi * vmax^2
         iv0 = findfirst(x -> x>0.0, vpa_unnorm)
+        if iv0 === nothing
+            error("All unnormalised vpa values at lower-z sheath entrance are negative. "
+                  * "Cannot apply electron boundary condition.")
+        end
         pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
         #println("check reversed change ", reversed_pdf[iv0:end])
         #println("reversed_pdf ", reversed_pdf)
@@ -972,6 +976,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
         #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
         phi[end,ir] = me_over_mi * vmin^2
         iv0 = findlast(x -> x<0.0, vpa_unnorm)
+        if iv0 === nothing
+            error("All unnormalised vpa values at upper-z sheath entrance are positive. "
+                  * "Cannot apply electron boundary condition.")
+        end
         pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
         #println("after pdf ", pdf[:,1,end,ir])
         # obtain the normalisation constants needed to ensure the zeroth, first and second moments

From 8d691d2df45e786aee0da86e1ee854792c8eeaa8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 27 Mar 2024 16:36:39 +0000
Subject: [PATCH 192/428] 2-step initialization of electron pdf

In the first step, evolve electron_ppar along with pdf_electron until it
is fairly close to steady state. Fully reaching steady state while
evolving electron_ppar takes a very large number of iterations, and
anyway it is time advanced along with the ions, so do not require
electron_ppar to 'exactly' reach steady state. Instead, do a second step
where pdf_electron is advanced with electron_ppar held fixed, until
pdf_electron reaches steady state (within tolerance).
---
 .../src/electron_kinetic_equation.jl          | 32 +++++++------
 moment_kinetics/src/initial_conditions.jl     | 46 ++++++++++++++++---
 moment_kinetics/src/moment_kinetics_input.jl  |  1 +
 moment_kinetics/src/time_advance.jl           |  4 +-
 4 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 46179e203..f902ab3fb 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -61,7 +61,7 @@ function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpa
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
-        evolve_ppar=false)
+        initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -74,7 +74,8 @@ function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpa
             vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
             external_source_settings, num_diss_params, max_electron_pdf_iterations;
             io_initial_electron=io_initial_electron, initial_time=initial_time,
-            evolve_ppar=evolve_ppar)
+            initial_output_counter=initial_output_counter,
+            residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
         return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
             qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
@@ -120,7 +121,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
-        evolve_ppar=false)
+        initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
 
     begin_r_z_region()
 
@@ -160,14 +161,12 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                          buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
     time = initial_time
-    if initial_time > 0.0
-        # Make sure that output times are set relative to this initial_time (before this they
-        # will be set relative to 0.0).
-        t_params.moments_output_times .+= initial_time
-        t_params.dfns_output_times .+= initial_time
-    end
+    # Make sure that output times are set relative to this initial_time (the values in
+    # t_params are set relative to 0.0).
+    moments_output_times = t_params.moments_output_times .+ initial_time
+    dfns_output_times = t_params.dfns_output_times .+ initial_time
     if io_initial_electron !== nothing
-        t_params.next_output_time[] = t_params.dfns_output_times[1]
+        t_params.next_output_time[] = dfns_output_times[1]
     end
 
     #z_speedup_fac = 20.0
@@ -219,7 +218,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
     end
 
-    output_counter = 0
+    output_counter = initial_output_counter
     begin_serial_region()
     output_counter += 1
     @serial_region begin
@@ -405,7 +404,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
             if global_rank[] == 0
-                electron_pdf_converged = abs(residual) < t_params.converged_residual_value
+                if residual_tolerance === nothing
+                    residual_tolerance = t_params.converged_residual_value
+                end
+                electron_pdf_converged = abs(residual) < residual_tolerance
             end
             electron_pdf_converged = MPI.Bcast(electron_pdf_converged, 0, comm_world)
         end
@@ -445,7 +447,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 println("iteration: ", t_params.step_counter[], " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
             end
         end
-        if (time ≥ t_params.dfns_output_times[output_counter] - epsilon)
+        if (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)
             begin_serial_region()
             @serial_region begin
                 if text_output
@@ -471,10 +473,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
             output_counter += 1
-            if output_counter ≤ length(t_params.dfns_output_times)
+            if output_counter - initial_output_counter ≤ length(dfns_output_times)
                 @serial_region begin
                     t_params.next_output_time[] =
-                        t_params.dfns_output_times[output_counter]
+                        dfns_output_times[output_counter - initial_output_counter]
                 end
             end
             @serial_region begin
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 824742701..021c377a5 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -226,8 +226,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                                vperp_spectral, vpa_spectral, collisions,
                                external_source_settings, scratch_dummy, scratch, t_params,
-                               num_diss_params, advection_structs, io_input, input_dict;
-                               restart_from_Boltzmann_electrons=false)
+                               t_input, num_diss_params, advection_structs, io_input,
+                               input_dict; restart_from_Boltzmann_electrons=false)
     
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
@@ -334,7 +334,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
-                             num_diss_params, t_params.electron, io_input, input_dict)
+                             num_diss_params, t_params.electron, t_input.electron_t_input,
+                             io_input, input_dict)
 
     return nothing
 end
@@ -429,7 +430,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                   vz, z_spectral, vperp_spectral, vpa_spectral, z_advect,
                                   vpa_advect, scratch_dummy, collisions, composition,
                                   geometry, external_source_settings, num_diss_params,
-                                  t_params, io_input, input_dict)
+                                  t_params, t_input, io_input, input_dict)
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
@@ -494,12 +495,42 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
             scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
             scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+
         # now that we have our initial guess for the electron pdf, we iterate
         # using the time-independent electron kinetic equation to find a self-consistent
-        # solution for the electron pdf
+        # solution for the electron pdf.
+        # First run with evolve_ppar=true to get electron_ppar close to steady state.
+        # electron_ppar does not have to be exactly steady state as it will be
+        # time-evolved along with the ions.
         max_electron_pdf_iterations = 2000000
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
+        if global_rank[] == 0
+            println("Initializing electrons - evolving both pdf_electron and electron_ppar")
+        end
+        electron_pseudotime, n_debug_outputs =
+            @views update_electron_pdf!(scratch, pdf.electron.norm, moments,
+                                        moments.electron.dens, moments.electron.vth,
+                                        moments.electron.ppar, moments.electron.qpar,
+                                        moments.electron.qpar_updated, phi,
+                                        moments.electron.ddens_dz,
+                                        moments.electron.dppar_dz,
+                                        moments.electron.dqpar_dz,
+                                        moments.electron.dvth_dz, r, z, vperp, vpa,
+                                        z_spectral, vperp_spectral, vpa_spectral,
+                                        z_advect, vpa_advect, scratch_dummy, t_params,
+                                        collisions, composition, external_source_settings,
+                                        num_diss_params, max_electron_pdf_iterations;
+                                        io_initial_electron=io_initial_electron,
+                                        initial_time=code_time,
+                                        residual_tolerance=t_input.initialization_residual_value,
+                                        evolve_ppar=true)
+
+        # Now run without evolve_ppar=true to get pdf_electron fully to steady state,
+        # ready for the start of the ion time advance.
+        if global_rank[] == 0
+            println("Initializing electrons - evolving pdf_electron only to steady state")
+        end
         electron_pseudotime, n_debug_outputs =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments,
                                         moments.electron.dens, moments.electron.vth,
@@ -514,11 +545,12 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         collisions, composition, external_source_settings,
                                         num_diss_params, max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
-                                        initial_time=code_time, evolve_ppar=true)
+                                        initial_time=electron_pseudotime,
+                                        initial_output_counter=n_debug_outputs)
 
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
-        t_idx = n_debug_outputs+1
+        t_idx = n_debug_outputs + 1
         write_initial_electron_state(pdf.electron.norm, moments, t_params,
                                      electron_pseudotime, io_initial_electron, t_idx, r,
                                      z, vperp, vpa)
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 246d3cced..1d886f18a 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -244,6 +244,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         minimum_dt=timestepping_section["minimum_dt"] * sqrt(composition.me_over_mi),
         maximum_dt=timestepping_section["maximum_dt"] * sqrt(composition.me_over_mi),
         high_precision_error_sum=timestepping_section["high_precision_error_sum"],
+        initialization_residual_value=1.0,
        )
     if electron_timestepping_section["nwrite"] === nothing
         electron_timestepping_section["nwrite"] = electron_timestepping_section["nstep"]
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 97bc064d2..fe58cb129 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -449,8 +449,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions,
                               external_source_settings, scratch_dummy, scratch, t_params,
-                              num_diss_params, advection_structs, io_input, input_dict;
-                              restart_from_Boltzmann_electrons=restarting)
+                              t_input, num_diss_params, advection_structs, io_input,
+                              input_dict; restart_from_Boltzmann_electrons=restarting)
     end
 
     # update the derivatives of the electron moments as these may be needed when

From f5bf19f55a3bc0240cb30c4658b9a6245b84f5cb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 08:15:31 +0000
Subject: [PATCH 193/428] Finish removing upwinded derivative in electron
 energy equation

Don't need to calculate `dppar_dz_upwind` in
`calculate_electron_moment_derivatives!()`.
---
 moment_kinetics/src/velocity_moments.jl | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index f92c0cf0a..79830b26e 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -969,22 +969,10 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
     buffer_r_3 = @view scratch_dummy.buffer_rs_3[:,1]
     buffer_r_4 = @view scratch_dummy.buffer_rs_4[:,1]
-    buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
-    buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
        
     @views derivative_z!(moments.electron.dupar_dz, upar, buffer_r_1,
                          buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
-    # Upwinded using upar as advection velocity, to be used in energy equation
-    @loop_r_z ir iz begin
-        dummy_zr[iz,ir] = -upar[iz,ir]
-    end
-    if electron_model ∈ (braginskii_fluid, kinetic_electrons)
-        @views derivative_z!(moments.electron.dppar_dz_upwind, ppar, dummy_zr,
-                             buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
-                             buffer_r_5, buffer_r_6, z_spectral, z)
-    end
-
     # centred second derivative for dissipation
     if numerical_dissipation.moment_dissipation_coefficient > 0.0
         @views derivative_z!(dummy_zr, ppar, buffer_r_1, buffer_r_2, buffer_r_3,

From e7766757514589db2fb49d11e7177c9d4848e957 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 08:22:59 +0000
Subject: [PATCH 194/428] Simplify arguments to calculate_electron_qpar!()

---
 moment_kinetics/src/electron_fluid_equations.jl |  8 ++++++--
 moment_kinetics/src/initial_conditions.jl       | 14 +++++++-------
 moment_kinetics/src/time_advance.jl             |  6 +++---
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 9c55a540e..6f36c4570 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -276,10 +276,14 @@ output:
     qpar_e = updated parallel electron heat flux
     qpar_updated = flag indicating that the parallel electron heat flux is updated
 """
-function calculate_electron_qpar!(qpar_e, qpar_updated, pdf, ppar_e, upar_e, vth_e, dTe_dz, upar_i, 
-                                  nu_ei, me_over_mi, electron_model, vpa)
+function calculate_electron_qpar!(electron_moments, pdf, ppar_e, upar_e, upar_i, nu_ei,
+                                  me_over_mi, electron_model, vpa)
     # only calculate qpar_e if needs updating
+    qpar_updated = electron_moments.qpar_updated
     if !qpar_updated[]
+        qpar_e = electron_moments.qpar
+        vth_e = electron_moments.vth
+        dTe_dz = electron_moments.dT_dz
         if electron_model == braginskii_fluid
             begin_r_z_region()
             # use the classical Braginskii expression for the electron heat flux
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 021c377a5..3f5b3d9d5 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -299,9 +299,9 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
     moments.electron.qpar_updated[] = false
-    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
-        moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
-        collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+    calculate_electron_qpar!(moments.electron, pdf.electron, moments.electron.ppar,
+        moments.electron.upar, moments.ion.upar, collisions.nu_ei, composition.me_over_mi,
+        composition.electron_physics, vpa)
     # calculate the zed derivative of the initial electron parallel heat flux
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
         scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
@@ -486,10 +486,10 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         end
 
         moments.electron.qpar_updated[] = false
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
-            moments.electron.ppar, moments.electron.upar, moments.electron.vth, 
-            moments.electron.dT_dz, moments.ion.upar, 
-            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+        calculate_electron_qpar!(moments.electron, pdf.electron, moments.electron.ppar,
+                                 moments.electron.upar, moments.ion.upar,
+                                 collisions.nu_ei, composition.me_over_mi,
+                                 composition.electron_physics, vpa)
         # update dqpar/dz for electrons
         # calculate the zed derivative of the initial electron parallel heat flux
         @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index fe58cb129..e08a71ae9 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1697,9 +1697,9 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                            num_diss_params, composition.electron_physics)
     # update the electron parallel heat flux
-    calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, new_scratch.pdf_electron,
-        new_scratch.electron_ppar, new_scratch.electron_upar, moments.electron.vth, moments.electron.dT_dz, 
-        new_scratch.upar, collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+    calculate_electron_qpar!(moments.electron, new_scratch.pdf_electron,
+        new_scratch.electron_ppar, new_scratch.electron_upar, new_scratch.upar,
+        collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
     # update the electron parallel friction force
     calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
         new_scratch.electron_upar, new_scratch.upar, moments.electron.dT_dz, composition.me_over_mi,

From f991c7e094541fa0192e37d1e076f5503eddd622 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 08:43:37 +0000
Subject: [PATCH 195/428] Simplify arguments to update_electron_pdf!()

---
 .../src/electron_kinetic_equation.jl          | 21 +++++++++--
 moment_kinetics/src/initial_conditions.jl     | 36 +++++++------------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index f902ab3fb..1192402fa 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -56,8 +56,7 @@ The electron kinetic equation is:
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
-function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpar_updated,
-        phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, r, z, vperp, vpa, z_spectral,
+function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_spectral,
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
@@ -77,10 +76,28 @@ function update_electron_pdf!(scratch, pdf, moments, dens, vthe, ppar, qpar, qpa
             initial_output_counter=initial_output_counter,
             residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
+        dens = moments.electron.dens
+        vthe = moments.electron.vth
+        ppar = moments.electron.ppar
+        qpar = moments.electron.qpar
+        qpar_updated = moments.electron.qpar_updated
+        ddens_dz = moments.electron.ddens_dz
+        dppar_dz = moments.electron.dppar_dz
+        dqpar_dz = moments.electron.dqpar_dz
+        dvth_dz = moments.electron.dvth_dz
         return update_electron_pdf_with_shooting_method!(pdf, dens, vthe, ppar, qpar,
             qpar_updated, phi, ddens_dz, dppar_dz, dqpar_dz, dvth_dz, z, vpa,
             vpa_spectral, scratch_dummy, composition)
     elseif solution_method == "picard_iteration"
+        dens = moments.electron.dens
+        vthe = moments.electron.vth
+        ppar = moments.electron.ppar
+        qpar = moments.electron.qpar
+        qpar_updated = moments.electron.qpar_updated
+        ddens_dz = moments.electron.ddens_dz
+        dppar_dz = moments.electron.dppar_dz
+        dqpar_dz = moments.electron.dqpar_dz
+        dvth_dz = moments.electron.dvth_dz
         return update_electron_pdf_with_picard_iteration!(pdf, dens, vthe, ppar, ddens_dz,
             dppar_dz, dqpar_dz, dvth_dz, z, vpa, vpa_spectral, scratch_dummy,
             max_electron_pdf_iterations)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 3f5b3d9d5..da3484e6c 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -509,18 +509,12 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             println("Initializing electrons - evolving both pdf_electron and electron_ppar")
         end
         electron_pseudotime, n_debug_outputs =
-            @views update_electron_pdf!(scratch, pdf.electron.norm, moments,
-                                        moments.electron.dens, moments.electron.vth,
-                                        moments.electron.ppar, moments.electron.qpar,
-                                        moments.electron.qpar_updated, phi,
-                                        moments.electron.ddens_dz,
-                                        moments.electron.dppar_dz,
-                                        moments.electron.dqpar_dz,
-                                        moments.electron.dvth_dz, r, z, vperp, vpa,
-                                        z_spectral, vperp_spectral, vpa_spectral,
-                                        z_advect, vpa_advect, scratch_dummy, t_params,
-                                        collisions, composition, external_source_settings,
-                                        num_diss_params, max_electron_pdf_iterations;
+            @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
+                                        vperp, vpa, z_spectral, vperp_spectral,
+                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                        t_params, collisions, composition,
+                                        external_source_settings, num_diss_params,
+                                        max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
                                         initial_time=code_time,
                                         residual_tolerance=t_input.initialization_residual_value,
@@ -532,18 +526,12 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             println("Initializing electrons - evolving pdf_electron only to steady state")
         end
         electron_pseudotime, n_debug_outputs =
-            @views update_electron_pdf!(scratch, pdf.electron.norm, moments,
-                                        moments.electron.dens, moments.electron.vth,
-                                        moments.electron.ppar, moments.electron.qpar,
-                                        moments.electron.qpar_updated, phi,
-                                        moments.electron.ddens_dz,
-                                        moments.electron.dppar_dz,
-                                        moments.electron.dqpar_dz,
-                                        moments.electron.dvth_dz, r, z, vperp, vpa,
-                                        z_spectral, vperp_spectral, vpa_spectral,
-                                        z_advect, vpa_advect, scratch_dummy, t_params,
-                                        collisions, composition, external_source_settings,
-                                        num_diss_params, max_electron_pdf_iterations;
+            @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
+                                        vperp, vpa, z_spectral, vperp_spectral,
+                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                        t_params, collisions, composition,
+                                        external_source_settings, num_diss_params,
+                                        max_electron_pdf_iterations;
                                         io_initial_electron=io_initial_electron,
                                         initial_time=electron_pseudotime,
                                         initial_output_counter=n_debug_outputs)

From 59ee104119ad3ef1898485180cd2668286b6d81e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 08:47:39 +0000
Subject: [PATCH 196/428] Include update of electron distribution function in
 ion/neutral time advance

---
 moment_kinetics/src/time_advance.jl | 35 +++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index e08a71ae9..c48c39c95 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -40,6 +40,7 @@ using ..neutral_vz_advection: update_speed_neutral_vz!, neutral_advection_vz!
 using ..vperp_advection: update_speed_vperp!, vperp_advection!
 using ..vpa_advection: update_speed_vpa!, vpa_advection!
 using ..charge_exchange: charge_exchange_collisions_1V!, charge_exchange_collisions_3V!
+using ..electron_kinetic_equation: update_electron_pdf!
 using ..ionization: ionization_collisions_1V!, ionization_collisions_3V!, constant_ionization_source!
 using ..krook_collisions: krook_collisions!
 using ..external_sources
@@ -856,9 +857,9 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                 advance_neutral_energy = true
             end
         end
-        # if treating the electrons as a fluid with Braginskii closure, 
-        # then advance the electron energy equation
-        if composition.electron_physics == braginskii_fluid
+        # if treating the electrons as a fluid with Braginskii closure, or
+        # moment-kinetically then advance the electron energy equation
+        if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
             advance_electron_energy = true
         end
 
@@ -1678,7 +1679,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     # if electron model is braginskii_fluid, then ppar is evolved via the energy equation
     # and is already updated;
     # otherwise update assuming electron temperature is fixed in time
-    if composition.electron_physics == braginskii_fluid
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
         @loop_r_z ir iz begin
             new_scratch.electron_ppar[iz,ir] = (rk_coefs[1]*moments.electron.ppar[iz,ir] 
                 + rk_coefs[2]*old_scratch.electron_ppar[iz,ir] + rk_coefs[3]*new_scratch.electron_ppar[iz,ir])
@@ -1700,6 +1701,32 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     calculate_electron_qpar!(moments.electron, new_scratch.pdf_electron,
         new_scratch.electron_ppar, new_scratch.electron_upar, new_scratch.upar,
         collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+    if composition.electron_physics == kinetic_electrons
+        max_electron_pdf_iterations = 100000
+
+        # Copy ion and electron moments from `scratch` into `moments` to be used in
+        # electron kinetic equation update
+        begin_r_z_region()
+        @loop_s_r_z is ir iz begin
+            moments.ion.dens[iz,ir,is] = new_scratch.density[iz,ir,is]
+            moments.ion.upar[iz,ir,is] = new_scratch.upar[iz,ir,is]
+            moments.ion.ppar[iz,ir,is] = new_scratch.ppar[iz,ir,is]
+        end
+        @loop_sn_r_z isn ir iz begin
+            moments.neutral.dens[iz,ir,isn] = new_scratch.density_neutral[iz,ir,isn]
+            moments.neutral.uz[iz,ir,isn] = new_scratch.uz_neutral[iz,ir,isn]
+            moments.neutral.pz[iz,ir,isn] = new_scratch.pz_neutral[iz,ir,isn]
+        end
+        @loop_r_z ir iz begin
+            moments.electron.ppar[iz,ir] = new_scratch.electron_ppar[iz,ir]
+        end
+
+        update_electron_pdf!(scratch, pdf.electron.norm, moments, fields.phi, r, z, vperp,
+                             vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
+                             vpa_advect, scratch_dummy, t_params.electron, collisions,
+                             composition, external_source_settings, num_diss_params,
+                             max_electron_pdf_iterations)
+    end
     # update the electron parallel friction force
     calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
         new_scratch.electron_upar, new_scratch.upar, moments.electron.dT_dz, composition.me_over_mi,

From f52d5456a2242c95c893b4164e25e1bc8496065b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 11:17:15 +0000
Subject: [PATCH 197/428] Add electron quantities to standard post-processing

---
 .../src/makie_post_processing.jl              | 22 +++++++++++++++++--
 moment_kinetics/src/load_data.jl              |  3 +++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 204b46ffe..b54363d76 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -36,8 +36,9 @@ using moment_kinetics.manufactured_solns: manufactured_solutions,
                                           manufactured_electric_fields
 using moment_kinetics.load_data: close_run_info, get_run_info_no_setup, get_variable,
                                  timestep_diagnostic_variables, em_variables,
-                                 ion_moment_variables, neutral_moment_variables,
-                                 all_moment_variables, ion_dfn_variables,
+                                 ion_moment_variables, electron_moment_variables,
+                                 neutral_moment_variables, all_moment_variables,
+                                 ion_dfn_variables, electron_dfn_variables,
                                  neutral_dfn_variables, all_dfn_variables, ion_variables,
                                  neutral_variables, all_variables
 using moment_kinetics.initial_conditions: vpagrid_to_dzdt
@@ -206,6 +207,17 @@ function makie_post_process(run_dir::Union{String,Tuple},
     has_rdim = any(ri !== nothing && ri.r.n > 1 for ri ∈ run_info_moments)
     has_zdim = any(ri !== nothing && ri.z.n > 1 for ri ∈ run_info_moments)
 
+    # Only plot electron stuff if some runs have electrons
+    if any(ri !== nothing for ri ∈ run_info_moments)
+        has_electrons = any(r.composition.electron_physics
+                            ∈ (braginskii_fluid, kinetic_electrons)
+                            for r in run_info_moments)
+    else
+        has_electrons = any(r.composition.electron_physics
+                            ∈ (braginskii_fluid, kinetic_electrons)
+                            for r in run_info_dfns)
+    end
+
     # Only plot neutral stuff if all runs have neutrals
     if any(ri !== nothing for ri ∈ run_info_moments)
         has_neutrals = all(r.n_neutral_species > 0 for r in run_info_moments)
@@ -220,6 +232,9 @@ function makie_post_process(run_dir::Union{String,Tuple},
     #############################
 
     moment_variable_list = tuple(em_variables..., ion_moment_variables...)
+    if has_electrons
+        moment_variable_list = tuple(moment_variable_list..., electron_moment_variables...)
+    end
     if has_neutrals
         moment_variable_list = tuple(moment_variable_list..., neutral_moment_variables...)
     end
@@ -295,6 +310,9 @@ function makie_post_process(run_dir::Union{String,Tuple},
     ############################################
     if any(ri !== nothing for ri in run_info_dfns)
         dfn_variable_list = ion_dfn_variables
+        if has_electrons
+            dfn_variable_list = tuple(dfn_variable_list..., electron_dfn_variables...)
+        end
         if has_neutrals
             dfn_variable_list = tuple(dfn_variable_list..., neutral_dfn_variables...)
         end
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 44b875ffc..ab33e908c 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3927,6 +3927,9 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         n = postproc_load_variable(run_info, "electron_density"; kwargs...)
         vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
         variable = get_collision_frequency_ei(run_info.collisions, n, vth)
+    elseif variable_name == "electron_temperature"
+        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        variable = vth.^2
     elseif variable_name == "temperature_neutral"
         vth = postproc_load_variable(run_info, "thermal_speed_neutral"; kwargs...)
         variable = vth.^2

From 79f613cb4b19b7ad7407d49a8fc4072a7aea0765 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 13:08:30 +0000
Subject: [PATCH 198/428] Add electron_ppar to error check used for adaptive
 timestep control

---
 moment_kinetics/src/file_io.jl      | 19 +++++++++++--------
 moment_kinetics/src/time_advance.jl | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index a70ca630e..5c63f5d3d 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -387,7 +387,7 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar, true)
+                                                      evolve_ppar, kinetic_electrons)
 
         close(fid)
 
@@ -843,7 +843,7 @@ define dynamic (time-evolving) moment variables for writing to the hdf5 file
 function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                           r::coordinate, z::coordinate, parallel_io,
                                           external_source_settings, evolve_density,
-                                          evolve_upar, evolve_ppar, kinetic_electrons)
+                                          evolve_upar, evolve_ppar, electron_physics)
     @serial_region begin
         dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
 
@@ -875,7 +875,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar, kinetic_electrons)
+                                                      evolve_ppar, electron_physics)
 
         io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
         io_thermal_speed_neutral, external_source_neutral_amplitude,
@@ -911,6 +911,9 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         if n_neutral_species > 0
             n_failure_vars *= 2
         end
+        if electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+            n_failure_vars += 1
+        end
         io_failure_caused_by = create_dynamic_variable!(
             dynamic, "failure_caused_by", mk_int; diagnostic_var_size=n_failure_vars,
             parallel_io=parallel_io,
@@ -1168,7 +1171,7 @@ define dynamic (time-evolving) electron moment variables for writing to the hdf5
 """
 function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordinate,
         parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar,
-        kinetic_electrons)
+        electron_physics)
 
     dynamic = get_group(fid, "dynamic_data")
 
@@ -1232,7 +1235,7 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                parallel_io=parallel_io,
                                description="'C' coefficient enforcing pressure constraint for electrons")
 
-    if kinetic_electrons
+    if electron_physics == kinetic_electrons
         io_electron_step_counter = create_dynamic_variable!(
             dynamic, "electron_step_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of electron pseudo-timesteps for the run")
@@ -1428,7 +1431,7 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
                                                       evolve_ppar,
-                                                      composition.electron_physics == kinetic_electrons)
+                                                      composition.electron_physics)
 
         dynamic = get_group(fid, "dynamic_data")
 
@@ -1539,7 +1542,7 @@ function setup_moments_io(prefix, binary_format, vz, vr, vzeta, vpa, vperp, r, z
         io_moments = define_dynamic_moment_variables!(
             fid, composition.n_ion_species, composition.n_neutral_species, r, z,
             parallel_io, external_source_settings, evolve_density, evolve_upar,
-            evolve_ppar, composition.electron_physics == kinetic_electrons)
+            evolve_ppar, composition.electron_physics)
 
         close(fid)
 
@@ -2495,7 +2498,7 @@ function debug_dump(vz::coordinate, vr::coordinate, vzeta::coordinate, vpa::coor
                                                           external_source_settings,
                                                           evolve_density, evolve_upar,
                                                           evolve_ppar,
-                                                          composition.electron_physics == kinetic_electrons)
+                                                          composition.electron_physics)
             io_dfns = define_dynamic_dfn_variables!(
                 fid, r, z, vperp, vpa, vzeta, vr, vz, composition.n_ion_species,
                 composition.n_neutral_species, false, external_source_settings,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 08ab3c705..94b74c785 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -400,6 +400,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
         # ion pressure
         push!(t_params.failure_caused_by, 0)
     end
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        # electron pressure
+        push!(t_params.failure_caused_by, 0)
+    end
     if composition.n_neutral_species > 0
         # neutral pdf
         push!(t_params.limit_caused_by, 0, 0)
@@ -2016,6 +2020,20 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
         push!(total_points, z.n_global * r.n_global * n_ion_species)
     end
 
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        begin_r_z_region()
+        rk_error_variable!(scratch, :electron_ppar, t_params)
+        electron_p_err = local_error_norm(scratch[2].electron_ppar,
+                                          scratch[t_params.n_rk_stages+1].electron_ppar,
+                                          t_params.rtol, t_params.atol;
+                                          method=error_norm_method,
+                                          skip_r_inner=skip_r_inner,
+                                          skip_z_lower=skip_z_lower,
+                                          error_sum_zero=t_params.error_sum_zero)
+        push!(error_norms, electron_p_err)
+        push!(total_points, z.n_global * r.n_global)
+    end
+
     if n_neutral_species > 0
         # neutral z-advection
         # Don't parallelise over species here, because get_minimum_CFL_*() does an MPI

From 75c911b10548b10ff5f5fde2000cfa3ba2b47103 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 13:09:25 +0000
Subject: [PATCH 199/428] Better info print for electron timestep diagnostics
 post-processing

---
 .../makie_post_processing/src/makie_post_processing.jl      | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index b54363d76..dfcf3933c 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -6946,7 +6946,11 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
             run_info = (run_info,)
         end
 
-        println("Making timestep diagnostics plots")
+        if electron
+            println("Making electron timestep diagnostics plots")
+        else
+            println("Making timestep diagnostics plots")
+        end
 
         input = Dict_to_NamedTuple(input_dict["timestep_diagnostics"])
 

From 882c4140137c8b7a670f1f940ba82712aeaa3c42 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 14:07:41 +0000
Subject: [PATCH 200/428] step_counter[] for electrons tracks global number of
 pseudo-timesteps

---
 .../src/electron_kinetic_equation.jl          | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 23a61ee66..0eb2fdafa 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -193,8 +193,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     text_output = false
 
     epsilon = 1.e-11
-    # initialise the number of iterations in the solution of the electron kinetic equation to be 1
-    t_params.step_counter[] = 1
+    # Store the initial number of iterations in the solution of the electron kinetic
+    # equation
+    initial_step_counter = t_params.step_counter[]
+    t_params.step_counter[] += 1
     # initialise the electron pdf convergence flag to false
     electron_pdf_converged = false
 
@@ -221,10 +223,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     println(io_pdf,"")
                 end
                 @loop_z iz begin
-                    println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                    println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[], " dens: ", dens[iz,1])
+                    println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                 end
                 println(io_upar,"")
                 println(io_qpar,"")
@@ -246,7 +248,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
-    while !electron_pdf_converged && (t_params.step_counter[] < max_electron_pdf_iterations) && t_params.dt[] > 0.0
+    while (!electron_pdf_converged
+           && (t_params.step_counter[] - initial_step_counter < max_electron_pdf_iterations)
+           && t_params.dt[] > 0.0)
         for istage ∈ 1:t_params.n_rk_stages
             # Set the initial values for this stage to the final values from the previous
             # stage
@@ -447,21 +451,21 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
 
         if text_output
-            if (mod(t_params.step_counter[],t_params.nwrite)==1)
+            if (mod(t_params.step_counter[] - initial_step_counter, t_params.nwrite)==1)
                 begin_serial_region()
                 @serial_region begin
                     @loop_vpa ivpa begin
-                        println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa,1,end,1], " iteration: ", t_params.step_counter[], " flag: ", 1)
+                        println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa,1,end,1], " iteration: ", t_params.step_counter[] - initial_step_counter, " flag: ", 1)
                     end
                     println(io_pdf_stages,"")
                 end
             end
         end
 
-        if (mod(t_params.step_counter[],100) == 0)
+        if (mod(t_params.step_counter[] - initial_step_counter,100) == 0)
             begin_serial_region()
             if global_rank[] == 0
-                println("iteration: ", t_params.step_counter[], " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
+                println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
             end
         end
         if (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)
@@ -478,10 +482,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                         println(io_pdf,"")
                     end
                     @loop_z iz begin
-                        println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                        println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[])
-                        println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[], " dens: ", dens[iz,1])
+                        println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                     end
                     println(io_upar,"")
                     println(io_qpar,"")

From 04628d4d745c4ffa10eb16f5401f97e77b62de78 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 14:28:40 +0000
Subject: [PATCH 201/428] Distinguish electrons/ions in adaptive timestep
 diagnostic print

---
 moment_kinetics/src/electron_kinetic_equation.jl | 3 ++-
 moment_kinetics/src/runge_kutta.jl               | 6 ++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 0eb2fdafa..fe0cd9cbd 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1258,7 +1258,8 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     end
 
     adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
-                                       total_points, current_dt, error_norm_method)
+                                       total_points, current_dt, error_norm_method,
+                                       electron=true)
 
     return nothing
 end
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 5f7329da1..207366f04 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -732,7 +732,8 @@ end
 Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_params`.
 """
 function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
-                                            total_points, current_dt, error_norm_method)
+                                            total_points, current_dt, error_norm_method;
+                                            electron=false)
     # Get global minimum of CFL limits
     CFL_limit = nothing
     this_limit_caused_by = nothing
@@ -907,7 +908,8 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
                 t_params.limit_caused_by[this_limit_caused_by] += 1
 
                 if (t_params.step_counter[] % 1000 == 0) && global_rank[] == 0
-                    println("step ", t_params.step_counter[], ": t=",
+                    prefix = electron ? "electron" : "ion"
+                    println("$prefix step ", t_params.step_counter[], ": t=",
                             round(t, sigdigits=6), ", nfail=", t_params.failure_counter[],
                             ", dt=", t_params.dt[])
                 end

From 0ae774b0596704bccf2bd40fbd02d4eb41fa33d3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 15:05:38 +0000
Subject: [PATCH 202/428] Fix post-proc calculation of 'electron_temperature'

...was missing a factor of me_over_mi before.
---
 moment_kinetics/src/load_data.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index ab33e908c..e9447bd35 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3929,7 +3929,7 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         variable = get_collision_frequency_ei(run_info.collisions, n, vth)
     elseif variable_name == "electron_temperature"
         vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
-        variable = vth.^2
+        variable = run_info.composition.me_over_mi .* vth.^2
     elseif variable_name == "temperature_neutral"
         vth = postproc_load_variable(run_info, "thermal_speed_neutral"; kwargs...)
         variable = vth.^2

From 7303db9c9e9ad905933bbe649ff60d748ea79548 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 17:32:32 +0000
Subject: [PATCH 203/428] Increase dissipation, use adaptive timestepping in
 wall+sheath-bc_kinetic_loworder.toml

---
 .../wall+sheath-bc_kinetic_loworder.toml            | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
index e9b8a56c5..c1e8a7637 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_loworder.toml
@@ -72,14 +72,14 @@ vz_bc = "zero"
 vz_discretization = "chebyshev_pseudospectral"
 
 [output]
-ascii_output = true
+#ascii_output = true
 
 [electron_numerical_dissipation]
 #moment_dissipation_coefficient = 0.0001
 #moment_dissipation_coefficient = 1.0
 #vpa_dissipation_coefficient = 0.002
-vpa_dissipation_coefficient = 0.2
-#vpa_dissipation_coefficient = 2.0
+#vpa_dissipation_coefficient = 0.2
+vpa_dissipation_coefficient = 2.0
 #vpa_dissipation_coefficient = 20.0
 
 [timestepping]
@@ -91,10 +91,11 @@ nwrite_dfns = 10000
 type = "SSPRK4"
 
 [electron_timestepping]
-nstep = 50000
+nstep = 5000000
 #nstep = 1
 dt = 2.0e-8
 nwrite = 1000
 nwrite_dfns = 1000
-type = "SSPRK4"
-#type = "Fekete4(3)"
+#type = "SSPRK4"
+type = "Fekete4(3)"
+minimum_dt = 1.0e-8

From c69b35022ddec083950dc25895fb25d475750a8d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 17:33:12 +0000
Subject: [PATCH 204/428] More example input files for kinetic electrons

---
 ...fraction0.5_split3_boltzmann-vpadiss0.toml |  95 +++++++++++++++
 ...c_recyclefraction0.5_split3_boltzmann.toml |  95 +++++++++++++++
 ...lefraction0.5_split3_kinetic-vpadiss0.toml | 112 ++++++++++++++++++
 3 files changed, 302 insertions(+)
 create mode 100644 examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
 create mode 100644 examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
 create mode 100644 examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
new file mode 100644
index 000000000..97e534663
--- /dev/null
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
@@ -0,0 +1,95 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+boltzmann_electron_response = true
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+krook_collisions_option = "reference_parameters"
+T_e = 0.2 # 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = -1.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.5
+constant_ionization_rate = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+z_nelement = 32
+z_nelement_local = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_element_spacing_option = "sqrt"
+vpa_ngrid = 6
+vpa_nelement = 63
+vpa_L = 36.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 63
+vz_L = 36.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "Fekete4(3)"
+#nstep = 50000
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-6
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+
+[ion_source]
+active = true
+z_profile = "gaussian"
+z_width = 0.125
+source_strength = 2.0
+source_T = 2.0
+
+[ion_numerical_dissipation]
+#vpa_dissipation_coefficient = 1.0e-1
+#vpa_dissipation_coefficient = 1.0e-2
+#vpa_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
new file mode 100644
index 000000000..23ba4265d
--- /dev/null
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
@@ -0,0 +1,95 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+boltzmann_electron_response = true
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+krook_collisions_option = "reference_parameters"
+T_e = 0.2 # 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = -1.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.5
+constant_ionization_rate = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+z_nelement = 32
+z_nelement_local = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_element_spacing_option = "sqrt"
+vpa_ngrid = 6
+vpa_nelement = 63
+vpa_L = 36.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 63
+vz_L = 36.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "Fekete4(3)"
+#nstep = 50000
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-6
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+
+[ion_source]
+active = true
+z_profile = "gaussian"
+z_width = 0.125
+source_strength = 2.0
+source_T = 2.0
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e-1
+#vpa_dissipation_coefficient = 1.0e-2
+#vpa_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
new file mode 100644
index 000000000..d4b93cd94
--- /dev/null
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
@@ -0,0 +1,112 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "kinetic_electrons"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+krook_collisions_option = "reference_parameters"
+T_e = 0.2 # 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = -1.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.5
+constant_ionization_rate = false
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+z_nelement = 32
+#z_nelement_local = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_element_spacing_option = "sqrt"
+vpa_ngrid = 6
+vpa_nelement = 63
+vpa_L = 36.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 63
+vz_L = 36.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "Fekete4(3)"
+#nstep = 50000
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-6
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+
+[electron_timestepping]
+nstep = 5000000
+#nstep = 1
+dt = 2.0e-8
+nwrite = 10000
+nwrite_dfns = 100000
+#type = "SSPRK4"
+type = "Fekete4(3)"
+rtol = 1.0e-3
+minimum_dt = 1.0e-8
+initialization_residual_value = 2.5
+converged_residual_value = 0.1 #1.0e-3
+
+[ion_source]
+active = true
+z_profile = "gaussian"
+z_width = 0.125
+source_strength = 2.0
+source_T = 2.0
+
+[ion_numerical_dissipation]
+#vpa_dissipation_coefficient = 1.0e-1
+#vpa_dissipation_coefficient = 1.0e-2
+#vpa_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient = 2.0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0

From ee8f7fd805eecbca1bfb90ca0ea866d9cc7103ad Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 17:53:27 +0000
Subject: [PATCH 205/428] Fix call to calculate_electron_qpar!() when not
 restarting

...had been missed in e7766757 "Simplify arguments to
calculate_electron_qpar!()".
---
 moment_kinetics/src/time_advance.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 94b74c785..52c3d9c7c 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -686,9 +686,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
             scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
             scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         # calculate the electron parallel heat flux
-        calculate_electron_qpar!(moments.electron.qpar, moments.electron.qpar_updated, pdf.electron,
-            moments.electron.ppar, moments.electron.upar, moments.electron.vth, moments.electron.dT_dz, moments.ion.upar, 
-            collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
+        calculate_electron_qpar!(moments.electron, pdf.electron, moments.electron.ppar,
+            moments.electron.upar, moments.ion.upar, collisions.nu_ei,
+            composition.me_over_mi, composition.electron_physics, vpa)
         # calculate the electron-ion parallel friction force
         calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
             moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,

From 864bd0b4b16e8c39bc1e864c89c3c9a4f3b8b797 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 28 Mar 2024 17:50:14 +0000
Subject: [PATCH 206/428] Distributed-MPI support for in kinetic electrons

---
 moment_kinetics/src/calculus.jl               |  62 +-
 .../src/electron_kinetic_equation.jl          | 916 +++++++++---------
 moment_kinetics/src/em_fields.jl              |   4 +-
 moment_kinetics/src/initial_conditions.jl     |  31 +-
 4 files changed, 519 insertions(+), 494 deletions(-)

diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl
index 4809182ed..1a10376db 100644
--- a/moment_kinetics/src/calculus.jl
+++ b/moment_kinetics/src/calculus.jl
@@ -354,35 +354,41 @@ function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims},
 	else
 		println("ERROR: invalid key in assign_endpoint!")
 	end
-    # test against coord name -- make sure to use exact string delimiters e.g. "x" not 'x'
-	# test against Ndims (autodetermined) to choose which array slices to use in assigning endpoints
-	#println("DEBUG MESSAGE: coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
-	if coord.name == "z" && Ndims==2
-		df1d[j,:] .= receive_buffer[:]
-		#println("ASSIGNING DATA")
+        # test against coord name -- make sure to use exact string delimiters e.g. "x" not 'x'
+        # test against Ndims (autodetermined) to choose which array slices to use in assigning endpoints
+        #println("DEBUG MESSAGE: coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
+        if coord.name == "z" && Ndims==2
+            df1d[j,:] .= receive_buffer[:]
+            #println("ASSIGNING DATA")
         elseif coord.name == "z" && Ndims==3
-		df1d[j,:,:] .= receive_buffer[:,:]
-		#println("ASSIGNING DATA")
-	elseif coord.name == "z" && Ndims==5
-		df1d[:,:,j,:,:] .= receive_buffer[:,:,:,:]
-		#println("ASSIGNING DATA")
-    elseif coord.name == "z" && Ndims==6
-		df1d[:,:,:,j,:,:] .= receive_buffer[:,:,:,:,:]
-		#println("ASSIGNING DATA")
-	elseif coord.name == "r" && Ndims==2
-		df1d[:,j] .= receive_buffer[:]
-		#println("ASSIGNING DATA")
-	elseif coord.name == "r" && Ndims==3
-		df1d[:,j,:] .= receive_buffer[:,:]
-		#println("ASSIGNING DATA")
-	elseif coord.name == "r" && Ndims==5
-		df1d[:,:,:,j,:] .= receive_buffer[:,:,:,:]
-		#println("ASSIGNING DATA")
-	elseif coord.name == "r" && Ndims==6
-		df1d[:,:,:,:,j,:] .= receive_buffer[:,:,:,:,:]
-		#println("ASSIGNING DATA")
-	else
-        println("ERROR: failure to assign endpoints in reconcile_element_boundaries_MPI! (centered): coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
+            df1d[j,:,:] .= receive_buffer[:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "z" && Ndims==4
+            df1d[:,:,j,:] .= receive_buffer[:,:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "z" && Ndims==5
+            df1d[:,:,j,:,:] .= receive_buffer[:,:,:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "z" && Ndims==6
+            df1d[:,:,:,j,:,:] .= receive_buffer[:,:,:,:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "r" && Ndims==2
+            df1d[:,j] .= receive_buffer[:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "r" && Ndims==3
+            df1d[:,j,:] .= receive_buffer[:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "r" && Ndims==4
+            df1d[:,:,:,j] .= receive_buffer[:,:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "r" && Ndims==5
+            df1d[:,:,:,j,:] .= receive_buffer[:,:,:,:]
+            #println("ASSIGNING DATA")
+        elseif coord.name == "r" && Ndims==6
+            df1d[:,:,:,:,j,:] .= receive_buffer[:,:,:,:,:]
+            #println("ASSIGNING DATA")
+        else
+        error("ERROR: failure to assign endpoints in reconcile_element_boundaries_MPI! (centered): coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
     end
 end
 
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index fe0cd9cbd..cd0f203b5 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -299,8 +299,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             # enforce the boundary condition(s) on the electron pdf
             enforce_boundary_condition_on_electron_pdf!(scratch[istage+1].pdf_electron, phi,
                                                         moments.electron.vth,
-                                                        moments.electron.upar, vperp, vpa,
-                                                        vperp_spectral, vpa_spectral,
+                                                        moments.electron.upar, z, vperp,
+                                                        vpa, vperp_spectral, vpa_spectral,
                                                         vpa_advect, moments,
                                                         num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
                                                         composition.me_over_mi)
@@ -603,7 +603,7 @@ function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=fal
     return nothing
 end
 
-function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp, vpa,
+function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, moments, vpa_diffusion,
                                                      me_over_mi)
@@ -650,226 +650,228 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
     # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
     #ivpa_zero = (vpa.n+1)÷2
 
-    @loop_r ir begin
-        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
-        #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
-
-        # Want to construct the w-grid corresponding to -vpa.
-        #   wpa(vpa) = (vpa - upar)/vth
-        #   ⇒ vpa = vth*wpa(vpa) + upar
-        #   wpa(-vpa) = (-vpa - upar)/vth
-        #             = (-(vth*wpa(vpa) + upar) - upar)/vth
-        #             = (-vth*wpa - 2*upar)/vth
-        #             = -wpa - 2*upar/vth
-        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
-        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
-        # interpolate the pdf onto this grid
-        #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
-        @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-        reverse!(reversed_pdf)
-        # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
-        #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
-        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
-        # interpolate back onto the original wpa grid
-        #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-        # construct wpa * pdf
-        #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
-        # calculate the first moment of the normalised pdf
-        #first_vspace_moment = 0.0
-        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
-        ivpa = 0
-        ivpa_max = vpa.n + 1
-        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-        # if the first moment is positive, then the cutoff speed needs to be reduced
-        upar_over_vth = upar[1,ir] / vthe[1,ir]
-        #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
-        #println("$first_vspace_moment, u/vth=$upar_over_vth")
-        vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
-        upar_integral = 0.0
-        #while first_vspace_moment > upar_over_vth # > 0.0
-        #    # zero out the pdf at the current cutoff velocity
-        #    pdf[ivpa_max,1,1,ir] = 0.0
-        #    # update wpa * pdf
-        #    vpa.scratch3[ivpa_max] = 0.0
-        #    # calculate the updated first moment of the normalised pdf
-        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
-        #    if first_vspace_moment > upar_over_vth #0.0
-        #        ivpa_max -= 1
-        #    end
-        #end
-        upar0 = upar[1,ir]
-        #println("before pdf left ", pdf[:,1,1,ir])
-        while upar_integral > upar0 && ivpa_max > 1
-            ivpa += 1
-            ivpa_max -= 1
-            # zero out the reversed pdf at the current cutoff velocity
-            #reversed_pdf[ivpa_max] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
-            #println("left ", ivpa, " ", upar_integral, " ", upar0)
-        end
-        integral_excess = upar_integral - upar0
-        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
-        #println("fraction_of_pdf=", fraction_of_pdf)
-
-        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-        # integral from the point at ivpa is required to make upar_integral<upar0) the
-        # cut-off velocity is half way between ivpa and ivpa+1, while when
-        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-        # ivpa is required to make upar_integral<upar0) the cut-off is half way between
-        # ivpa-1 and ivpa.
-        vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
-               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
-
-        #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
-        #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
-        #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
-        #@loop_vpa ivpa begin
-        #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
-        #end
-        reversed_pdf[ivpa_max+1:end] .= 0
-        reversed_pdf[ivpa_max] *= fraction_of_pdf
-        #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
-        #println("done first cutoff loop")
-        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
-        phi[1,ir] = me_over_mi * vmax^2
-        iv0 = findfirst(x -> x>0.0, vpa_unnorm)
-        if iv0 === nothing
-            error("All unnormalised vpa values at lower-z sheath entrance are negative. "
-                  * "Cannot apply electron boundary condition.")
-        end
-        pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
-        #println("check reversed change ", reversed_pdf[iv0:end])
-        #println("reversed_pdf ", reversed_pdf)
-        #println("after pdf left ", pdf[:,1,1,ir])
-        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-        # of the modified pdf are 1, 0 and 1/2 respectively
-        # will need vpa / vthe = wpa + upar/vthe
-        @. vpa.scratch2 = vpa.grid + upar[1,ir] / vthe[1,ir]
-        # first need to calculate int dwpa pdf = zeroth_moment
-        zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
-        # calculate int dwpa wpa^2 * pdf = wpa2_moment
-        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
-        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * pdf = vpa2_moment
-        @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
-        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
-        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
-        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        
-        normalisation_constant_A = 0.0
-        normalisation_constant_B = 0.0
-        normalisation_constant_C = 0.0
-        if pdf_adjustment_option == "absvpa"
-            # calculate int dwpa |vpa| * pdf = absvpa_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
-            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            # calculate the 'B' normalisation constant
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-            # calculate the 'A' normalisation constant
-            normalisation_constant_A = (1 + normalisation_constant_B * 
-                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-            # calculate the 'C' normalisation constant
-            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                + vpa.scratch2^2 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4"
-            # calculate int dwpa vpa^4 * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                + vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4_gaussian"
-            afac = 0.1
-            bfac = 0.2
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+    if z.irank == 0
+        @loop_r ir begin
+            # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+            #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
+            #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+
+            # Want to construct the w-grid corresponding to -vpa.
+            #   wpa(vpa) = (vpa - upar)/vth
+            #   ⇒ vpa = vth*wpa(vpa) + upar
+            #   wpa(-vpa) = (-vpa - upar)/vth
+            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+            #             = (-vth*wpa - 2*upar)/vth
+            #             = -wpa - 2*upar/vth
+            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+            # Need to reverse vpa.grid because the grid passed as the second argument of
+            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
+            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
+            reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+            # interpolate the pdf onto this grid
+            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+            reverse!(reversed_pdf)
+            # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
+            #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
+            # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+            #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
+            # interpolate back onto the original wpa grid
+            #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+            # construct wpa * pdf
+            #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
+            # calculate the first moment of the normalised pdf
+            #first_vspace_moment = 0.0
+            #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
+            ivpa = 0
+            ivpa_max = vpa.n + 1
+            # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+            # if the first moment is positive, then the cutoff speed needs to be reduced
+            upar_over_vth = upar[1,ir] / vthe[1,ir]
+            #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
+            #println("$first_vspace_moment, u/vth=$upar_over_vth")
+            vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
+            upar_integral = 0.0
+            #while first_vspace_moment > upar_over_vth # > 0.0
+            #    # zero out the pdf at the current cutoff velocity
+            #    pdf[ivpa_max,1,1,ir] = 0.0
+            #    # update wpa * pdf
+            #    vpa.scratch3[ivpa_max] = 0.0
+            #    # calculate the updated first moment of the normalised pdf
+            #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
+            #    if first_vspace_moment > upar_over_vth #0.0
+            #        ivpa_max -= 1
+            #    end
+            #end
+            upar0 = upar[1,ir]
+            #println("before pdf left ", pdf[:,1,1,ir])
+            while upar_integral > upar0 && ivpa_max > 1
+                ivpa += 1
+                ivpa_max -= 1
+                # zero out the reversed pdf at the current cutoff velocity
+                #reversed_pdf[ivpa_max] = 0.0
+                # calculate the updated first moment of the normalised pdf
+                upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
+                #println("left ", ivpa, " ", upar_integral, " ", upar0)
+            end
+            integral_excess = upar_integral - upar0
+            fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
+            #println("fraction_of_pdf=", fraction_of_pdf)
+
+            # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+            # integral from the point at ivpa is required to make upar_integral<upar0) the
+            # cut-off velocity is half way between ivpa and ivpa+1, while when
+            # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+            # ivpa is required to make upar_integral<upar0) the cut-off is half way between
+            # ivpa-1 and ivpa.
+            vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
+                   0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
+
+            #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
+            #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
+            #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
+            #@loop_vpa ivpa begin
+            #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
+            #end
+            reversed_pdf[ivpa_max+1:end] .= 0
+            reversed_pdf[ivpa_max] *= fraction_of_pdf
+            #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
+            #println("done first cutoff loop")
+            # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+            #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
+            phi[1,ir] = me_over_mi * vmax^2
+            iv0 = findfirst(x -> x>0.0, vpa_unnorm)
+            if iv0 === nothing
+                error("All unnormalised vpa values at lower-z sheath entrance are negative. "
+                      * "Cannot apply electron boundary condition.")
+            end
+            pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
+            #println("check reversed change ", reversed_pdf[iv0:end])
+            #println("reversed_pdf ", reversed_pdf)
+            #println("after pdf left ", pdf[:,1,1,ir])
+            # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+            # of the modified pdf are 1, 0 and 1/2 respectively
+            # will need vpa / vthe = wpa + upar/vthe
+            @. vpa.scratch2 = vpa.grid + upar[1,ir] / vthe[1,ir]
+            # first need to calculate int dwpa pdf = zeroth_moment
+            zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
+            # calculate int dwpa wpa^2 * pdf = wpa2_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
+            wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
             vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
+            # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
             vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
+            # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
             vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #                           (vpa2_wpa2_moment
-            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-            #normalisation_constant_A = (1 - normalisation_constant_B 
-            #                                * vpa2_moment) / zeroth_moment
-            #normalisation_constant_C = 0.0
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
-            @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
-        else
-            println("pdf_adjustment_option not recognised")
-            stop()
-        end
+            
+            normalisation_constant_A = 0.0
+            normalisation_constant_B = 0.0
+            normalisation_constant_C = 0.0
+            if pdf_adjustment_option == "absvpa"
+                # calculate int dwpa |vpa| * pdf = absvpa_moment
+                @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
+                absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+                # calculate the 'B' normalisation constant
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                    (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                    + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+                # calculate the 'A' normalisation constant
+                normalisation_constant_A = (1 + normalisation_constant_B * 
+                    (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+                # calculate the 'C' normalisation constant
+                normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+                # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+                @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                    + vpa.scratch2^2 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "vpa4"
+                # calculate int dwpa vpa^4 * pdf = vpa4_moment
+                @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
+                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+                normalisation_constant_A = (1 + normalisation_constant_B 
+                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                    + vpa.scratch2^4 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "vpa4_gaussian"
+                afac = 0.1
+                bfac = 0.2
+                # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+                @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+                vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+                @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
+                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+                normalisation_constant_A = (1 + normalisation_constant_B 
+                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                #                           (vpa2_wpa2_moment
+                #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+                #normalisation_constant_A = (1 - normalisation_constant_B 
+                #                                * vpa2_moment) / zeroth_moment
+                #normalisation_constant_C = 0.0
+                @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                    + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "no1st_vpa2"
+                normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                           (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+                normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
+                @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
+            else
+                println("pdf_adjustment_option not recognised")
+                stop()
+            end
 
-        moments.electron.constraints_A_coefficient[1,ir] = normalisation_constant_A
-        moments.electron.constraints_B_coefficient[1,ir] = normalisation_constant_B
-        moments.electron.constraints_C_coefficient[1,ir] = normalisation_constant_C
+            moments.electron.constraints_A_coefficient[1,ir] = normalisation_constant_A
+            moments.electron.constraints_B_coefficient[1,ir] = normalisation_constant_B
+            moments.electron.constraints_C_coefficient[1,ir] = normalisation_constant_C
 
-        # smooth the pdf at the boundary
-        #for ivpa ∈ 2:ivpa_max-1
-        #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
-        #end
+            # smooth the pdf at the boundary
+            #for ivpa ∈ 2:ivpa_max-1
+            #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
+            #end
+        end
     end
 
     # next enforce the boundary condition at z_max.
@@ -892,248 +894,250 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, vperp
     # end
     # println(io_pdf_stages,"")
 
-    @loop_r ir begin
-        # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-        #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-
-        # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-        #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-        # Need to reverse vpa.grid because the grid passed as the second argument of
-        # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
-        #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
-        reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
-        # interpolate the pdf onto this grid
-        #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
-        @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-        reverse!(reversed_pdf)
-        # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
-        #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
-        # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-        #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
-        # interpolate back onto the original wpa grid
-        #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
-        # end
-        # println(io_pdf_stages,"")
-
-        # construct wpa * pdf
-        #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
-        # calculate the first moment of the normalised pdf
-        #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
-        ivpa = vpa.n+1
-        ivpa_min = 0
-        # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-        # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
-        upar_over_vth = upar[end,ir] / vthe[end,ir]
-        #println("$first_vspace_moment, u/vth=$upar_over_vth")
-        vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
-        upar_integral = 0.0
-        #while first_vspace_moment < upar_over_vth # < 0.0
-        #    # zero out the pdf at the current cutoff velocity
-        #    pdf[ivpa_min,1,end,ir] = 0.0
-        #    # update wpa * pdf
-        #    vpa.scratch3[ivpa_min] = 0.0
-        #    # calculate the updated first moment of the normalised pdf
-        #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        #    if first_vspace_moment < upar_over_vth
-        #        ivpa_min += 1
-        #    end
-        #end
-        upar_end = upar[end,ir]
-        #println("before pdf ", pdf[:,1,end,ir])
-        while upar_integral < upar_end && ivpa > 1
-            ivpa -= 1
-            ivpa_min += 1
-            # zero out the reversed pdf at the current cutoff velocity
-            #reversed_pdf[ivpa_min] = 0.0
-            # calculate the updated first moment of the normalised pdf
-            upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
-            #println("right ", ivpa, " ", upar_integral, " ", upar_end)
-        end
-        integral_excess = upar_integral - upar_end
-        fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
-        #println("B fraction_of_pdf=", fraction_of_pdf)
-
-        # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-        # integral from the point at ivpa is required to make upar_integral>upar_end) the
-        # cut-off velocity is half way between ivpa-1 and ivpa, while when
-        # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-        # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
-        # ivpa and ivpa+1.
-        vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
-               0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
-
-        #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
-        #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
-        #@loop_vpa ivpa begin
-        #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
-        #end
-        reversed_pdf[1:ivpa_min-1] .= 0
-        reversed_pdf[ivpa_min] *= fraction_of_pdf
-        #println("done second cutoff loop")
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
-        # end
-        # println(io_pdf_stages,"")
-
-        # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-        #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
-        phi[end,ir] = me_over_mi * vmin^2
-        iv0 = findlast(x -> x<0.0, vpa_unnorm)
-        if iv0 === nothing
-            error("All unnormalised vpa values at upper-z sheath entrance are positive. "
-                  * "Cannot apply electron boundary condition.")
-        end
-        pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
-        #println("after pdf ", pdf[:,1,end,ir])
-        # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-        # of the modified pdf are 1, 0 and 1/2 respectively
-        # will need vpa / vthe = wpa + upar/vthe
-        @. vpa.scratch2 = vpa.grid + upar[end,ir] / vthe[end,ir]
-        # first need to calculate int dwpa pdf = zeroth_moment
-        zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
-        # calculate int dwpa wpa^2 * pdf = wpa2_moment
-        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
-        wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * pdf = vpa2_moment
-        @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
-        vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-        @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
-        vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-        @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
-        vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        
-        normalisation_constant_A = 0.0
-        normalisation_constant_B = 0.0
-        normalisation_constant_C = 0.0
-        if pdf_adjustment_option == "absvpa"
-            # calculate int dwpa |vpa| * pdf = absvpa_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
-            absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            # calculate the 'B' normalisation constant
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-            # calculate the 'A' normalisation constant
-            normalisation_constant_A = (1 + normalisation_constant_B * 
-                (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-            # calculate the 'C' normalisation constant
-            normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-            # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                + vpa.scratch2^2 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4"
-            # calculate int dwpa vpa^4 * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                + vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "vpa4_gaussian"
-            afac = 0.1
-            bfac = 0.2
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+    if z.irank == z.nrank - 1
+        @loop_r ir begin
+            # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
+            #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
+            # Need to reverse vpa.grid because the grid passed as the second argument of
+            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+
+            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+            # Need to reverse vpa.grid because the grid passed as the second argument of
+            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
+            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
+            reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+            # interpolate the pdf onto this grid
+            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+            reverse!(reversed_pdf)
+            # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
+            #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
+            # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
+            #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
+            # interpolate back onto the original wpa grid
+            #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
+
+            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @. vpa.scratch3 *= vpa.grid
+            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @loop_vpa ivpa begin
+            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
+            # end
+            # println(io_pdf_stages,"")
+
+            # construct wpa * pdf
+            #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
+            # calculate the first moment of the normalised pdf
+            #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
+            ivpa = vpa.n+1
+            ivpa_min = 0
+            # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
+            # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
+            upar_over_vth = upar[end,ir] / vthe[end,ir]
+            #println("$first_vspace_moment, u/vth=$upar_over_vth")
+            vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
+            upar_integral = 0.0
+            #while first_vspace_moment < upar_over_vth # < 0.0
+            #    # zero out the pdf at the current cutoff velocity
+            #    pdf[ivpa_min,1,end,ir] = 0.0
+            #    # update wpa * pdf
+            #    vpa.scratch3[ivpa_min] = 0.0
+            #    # calculate the updated first moment of the normalised pdf
+            #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            #    if first_vspace_moment < upar_over_vth
+            #        ivpa_min += 1
+            #    end
+            #end
+            upar_end = upar[end,ir]
+            #println("before pdf ", pdf[:,1,end,ir])
+            while upar_integral < upar_end && ivpa > 1
+                ivpa -= 1
+                ivpa_min += 1
+                # zero out the reversed pdf at the current cutoff velocity
+                #reversed_pdf[ivpa_min] = 0.0
+                # calculate the updated first moment of the normalised pdf
+                upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
+                #println("right ", ivpa, " ", upar_integral, " ", upar_end)
+            end
+            integral_excess = upar_integral - upar_end
+            fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
+            #println("B fraction_of_pdf=", fraction_of_pdf)
+
+            # Define so that when fraction_of_pdf=1 (when all of the contribution to the
+            # integral from the point at ivpa is required to make upar_integral>upar_end) the
+            # cut-off velocity is half way between ivpa-1 and ivpa, while when
+            # fraction_of_pdf=0 (none of the contribution to the integral from the point at
+            # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
+            # ivpa and ivpa+1.
+            vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
+                   0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
+
+            #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
+            #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
+            #@loop_vpa ivpa begin
+            #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
+            #end
+            reversed_pdf[1:ivpa_min-1] .= 0
+            reversed_pdf[ivpa_min] *= fraction_of_pdf
+            #println("done second cutoff loop")
+
+            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @. vpa.scratch3 *= vpa.grid
+            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @loop_vpa ivpa begin
+            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
+            # end
+            # println(io_pdf_stages,"")
+
+            # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
+            #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
+            phi[end,ir] = me_over_mi * vmin^2
+            iv0 = findlast(x -> x<0.0, vpa_unnorm)
+            if iv0 === nothing
+                error("All unnormalised vpa values at upper-z sheath entrance are positive. "
+                      * "Cannot apply electron boundary condition.")
+            end
+            pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
+            #println("after pdf ", pdf[:,1,end,ir])
+            # obtain the normalisation constants needed to ensure the zeroth, first and second moments
+            # of the modified pdf are 1, 0 and 1/2 respectively
+            # will need vpa / vthe = wpa + upar/vthe
+            @. vpa.scratch2 = vpa.grid + upar[end,ir] / vthe[end,ir]
+            # first need to calculate int dwpa pdf = zeroth_moment
+            zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
+            # calculate int dwpa wpa^2 * pdf = wpa2_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
+            wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # calculate int dwpa vpa^2 * pdf = vpa2_moment
+            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
             vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
+            # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
+            @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
             vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
+            # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
+            @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
             vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-            @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
-            vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-            @. vpa.scratch3 *= vpa.grid
-            vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-            normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                 (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                 + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-            normalisation_constant_A = (1 + normalisation_constant_B 
-                * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-            normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-            #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-            #                           (vpa2_wpa2_moment
-            #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-            #normalisation_constant_A = (1 - normalisation_constant_B 
-            #                                * vpa2_moment) / zeroth_moment
-            #normalisation_constant_C = 0.0
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-        elseif pdf_adjustment_option == "no1st_vpa2"
-            normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                       (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-            normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
-            @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
-        else
-            println("pdf_adjustment_option not recognised")
-            stop()
-        end
+            
+            normalisation_constant_A = 0.0
+            normalisation_constant_B = 0.0
+            normalisation_constant_C = 0.0
+            if pdf_adjustment_option == "absvpa"
+                # calculate int dwpa |vpa| * pdf = absvpa_moment
+                @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
+                absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+                # calculate the 'B' normalisation constant
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
+                    (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
+                    + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
+                # calculate the 'A' normalisation constant
+                normalisation_constant_A = (1 + normalisation_constant_B * 
+                    (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
+                # calculate the 'C' normalisation constant
+                normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
+                # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
+                @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
+                                    + vpa.scratch2^2 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "vpa4"
+                # calculate int dwpa vpa^4 * pdf = vpa4_moment
+                @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
+                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+                normalisation_constant_A = (1 + normalisation_constant_B 
+                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
+                                    + vpa.scratch2^4 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "vpa4_gaussian"
+                afac = 0.1
+                bfac = 0.2
+                # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
+                @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
+                vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
+                @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
+                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
+                @. vpa.scratch3 *= vpa.grid
+                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
+                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
+                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
+                normalisation_constant_A = (1 + normalisation_constant_B 
+                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
+                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
+                #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
+                #                           (vpa2_wpa2_moment
+                #                            - wpa2_moment / zeroth_moment * vpa2_moment)
+                #normalisation_constant_A = (1 - normalisation_constant_B 
+                #                                * vpa2_moment) / zeroth_moment
+                #normalisation_constant_C = 0.0
+                @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
+                                    + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
+            elseif pdf_adjustment_option == "no1st_vpa2"
+                normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
+                                           (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
+                normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
+                @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
+            else
+                println("pdf_adjustment_option not recognised")
+                stop()
+            end
 
-        moments.electron.constraints_A_coefficient[end,ir] = normalisation_constant_A
-        moments.electron.constraints_B_coefficient[end,ir] = normalisation_constant_B
-        moments.electron.constraints_C_coefficient[end,ir] = normalisation_constant_C
-
-        # smooth the pdf at the boundary
-        #for ivpa ∈ ivpa_min+1:vpa.n-1
-        #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
-        #end
-
-        # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-        # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-        # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @. vpa.scratch3 *= vpa.grid
-        # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-        # @loop_vpa ivpa begin
-        #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-        #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
-        # end
-        # println(io_pdf_stages,"")
+            moments.electron.constraints_A_coefficient[end,ir] = normalisation_constant_A
+            moments.electron.constraints_B_coefficient[end,ir] = normalisation_constant_B
+            moments.electron.constraints_C_coefficient[end,ir] = normalisation_constant_C
+
+            # smooth the pdf at the boundary
+            #for ivpa ∈ ivpa_min+1:vpa.n-1
+            #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
+            #end
+
+            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
+            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
+            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @. vpa.scratch3 *= vpa.grid
+            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
+            # @loop_vpa ivpa begin
+            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
+            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
+            # end
+            # println(io_pdf_stages,"")
+        end
     end
 
     # # initialise the electron pdf for positive vpa to be the mirror reflection about vpa = 0
diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 91912cc2c..7c2a5f243 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -16,6 +16,8 @@ using ..velocity_moments: update_density!
 using ..derivatives: derivative_r!, derivative_z!
 using ..electron_fluid_equations: calculate_Epar_from_electron_force_balance!
 
+using MPI
+
 """
 """
 function setup_em_fields(nz, nr, force_phi, drive_amplitude, drive_frequency, force_Er_zero)
@@ -166,7 +168,7 @@ function calculate_phi_from_Epar!(phi, Epar, z)
         # this one.
         this_delta_phi = phi[end,:] .- phi[1,:]
         for irank ∈ 1:z.nrank-1
-            delta_phi = MPI.Bcast(this_delta_phi, irank, z.comm)
+            delta_phi = MPI.bcast(this_delta_phi, z.comm; root=irank)
             if z.irank > irank
                 @loop_r_z ir iz begin
                     phi[iz,ir] += delta_phi[ir]
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6877643dc..6cd49acfa 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -294,7 +294,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
             [(speed=speed,)], moments, num_diss_params,
-            composition.me_over_mi)
+            composition.me_over_mi, scratch_dummy)
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
@@ -1209,7 +1209,7 @@ this 'initital' value for the electron will just be the first guess in an iterat
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
         vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, moments, num_diss_params,
-        me_over_mi; restart_from_boltzmann=false)
+        me_over_mi, scratch_dummy; restart_from_boltzmann=false)
 
     if z.bc == "wall"
         begin_r_region()
@@ -1224,12 +1224,28 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
         end
         # Apply the sheath boundary condition to get cut-off boundary distribution
         # functions and boundary values of phi
-        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, vperp, vpa,
+        enforce_boundary_condition_on_electron_pdf!(pdf, phi, vth, upar, z, vperp, vpa,
                                                     vperp_spectral, vpa_spectral,
                                                     vpa_advect, moments,
                                                     num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
                                                     me_over_mi)
-        begin_r_region()
+
+        # Distribute the z-boundary pdf values to every process
+        begin_serial_region()
+        pdf_lower = scratch_dummy.buffer_vpavperprs_1
+        pdf_upper = scratch_dummy.buffer_vpavperprs_2
+        @serial_region begin
+            if z.irank == 0
+                pdf_lower .= pdf[:,:,1,:]
+            end
+            MPI.Bcast!(pdf_lower, z.comm; root=0)
+            if z.irank == z.nrank - 1
+                pdf_upper .= pdf[:,:,end,:]
+            end
+            MPI.Bcast!(pdf_upper, z.comm; root=z.nrank-1)
+        end
+
+        begin_r_z_region()
         @loop_r ir begin
             # get critical velocities beyond which electrons are lost to the wall
             #vpa_crit_zmin, vpa_crit_zmax = get_electron_critical_velocities(phi, vth, me_over_mi, z)
@@ -1237,9 +1253,6 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
             # Blend boundary distribution function into bulk of domain to avoid
             # discontinuities (as much as possible)
             blend_fac = 100
-            if z.nrank > 1
-                error("Distributed MPI not supported in this init yet")
-            end
             @loop_z_vperp iz ivperp begin
                 #@. pdf[:,ivperp,iz] = exp(-30*z.grid[iz]^2)
                 #@. pdf[:,ivperp,iz] = (density[iz] / vth[iz]) *
@@ -1247,8 +1260,8 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
                 blend_fac_lower = exp(-blend_fac*(z.grid[iz] + 0.5*z.L)^2)
                 blend_fac_upper = exp(-blend_fac*(z.grid[iz] - 0.5*z.L)^2)
                 @. pdf[:,ivperp,iz,ir] = (1.0 - blend_fac_lower) * (1.0 - blend_fac_upper) * pdf[:,ivperp,iz,ir] +
-                                        blend_fac_lower * pdf[:,ivperp,1,ir] +
-                                        blend_fac_upper * pdf[:,ivperp,end,ir]
+                                        blend_fac_lower * pdf_lower[:,ivperp,ir] +
+                                        blend_fac_upper * pdf_upper[:,ivperp,ir]
                 #@. pdf[:,ivperp,iz,ir] = exp(-vpa.grid^2) * (
                 #                         (1 - exp(-blend_fac*(z.grid[iz] - z.grid[1])^2) *
                 #                          tanh(sharp_fac*(vpa.grid-vpa_crit_zmin))) *

From 5f0ee516ae8d899be2b11a50a1ae2393bb18ba26 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 19:00:09 +0100
Subject: [PATCH 207/428] Initialize several arrays to avoid errors from debug
 checks

The arrays were not actually used uninitialized, only communicated
before being overwritten, but fixing lets debug checks work properly.
---
 moment_kinetics/src/initial_conditions.jl | 10 ++++++++++
 moment_kinetics/src/load_data.jl          |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6cd49acfa..cca5bdf8b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -147,6 +147,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
             # initialise pressures assuming isotropic distribution
             @. moments.ion.ppar = 0.5 * moments.ion.dens * moments.ion.vth^2
             @. moments.ion.pperp = moments.ion.ppar
+            if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+                @. moments.ion.constraints_A_coefficient = 1.0
+                @. moments.ion.constraints_B_coefficient = 0.0
+                @. moments.ion.constraints_C_coefficient = 0.0
+            end
             if n_neutral_species > 0
                 # initialise the neutral density profile
                 init_density!(moments.neutral.dens, z, r, species.neutral, n_neutral_species)
@@ -162,6 +167,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                 @. moments.neutral.pz = 0.5 * moments.neutral.dens * moments.neutral.vth^2
                 # calculate the total neutral pressure
                 @. moments.neutral.ptot = 1.5 * moments.neutral.dens * moments.neutral.vth^2
+                if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+                    @. moments.neutral.constraints_A_coefficient = 1.0
+                    @. moments.neutral.constraints_B_coefficient = 0.0
+                    @. moments.neutral.constraints_C_coefficient = 0.0
+                end
             end
         end
         # reflect the fact that the ion moments have now been updated
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index e9447bd35..aa7185b16 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -641,6 +641,10 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                               restart_r_spectral, restart_z,
                                               restart_z_spectral, interpolation_needed)
             moments.ion.ppar_updated .= true
+            moments.ion.pperp .= reload_moment("perpendicular_pressure", dynamic,
+                                               time_index, r, z, r_range, z_range,
+                                               restart_r, restart_r_spectral, restart_z,
+                                               restart_z_spectral, interpolation_needed)
             moments.ion.qpar .= reload_moment("parallel_heat_flux", dynamic, time_index,
                                               r, z, r_range, z_range, restart_r,
                                               restart_r_spectral, restart_z,
@@ -759,6 +763,9 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                         restart_vpa_spectral, interpolation_needed,
                                         restart_evolve_density, restart_evolve_upar,
                                         restart_evolve_ppar)
+            elseif pdf.electron !== nothing
+                # The electron distribution function will be initialized later
+                pdf.electron.norm .= 0.0
             end
 
             if composition.n_neutral_species > 0

From 41d5e9896602feba30663bf1379527ddb2cbdd45 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 19:14:59 +0100
Subject: [PATCH 208/428] Fix shared-memory usage in
 electron_adaptive_timestep_update!()

---
 moment_kinetics/src/electron_kinetic_equation.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index cd0f203b5..82adb74a4 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1210,7 +1210,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     #
     # z-advection
     # No need to synchronize here, as we just called _block_synchronize()
-    begin_r_z_vperp_vpa_region(; no_synchronize=true)
+    begin_r_vperp_vpa_region(; no_synchronize=true)
     update_electron_speed_z!(z_advect[1], moments.electron.upar, moments.electron.vth,
                              vpa.grid)
     z_CFL = get_minimum_CFL_z(z_advect[1].speed, z)
@@ -1221,6 +1221,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     end
 
     # vpa-advection
+    begin_r_z_vperp_region()
     update_electron_speed_vpa!(vpa_advect[1], moments.electron.dens,
                                moments.electron.upar, moments.electron.ppar,
                                moments, vpa.grid, external_source_settings.electron)

From a27706ed6a2ac7079b445a4be96de3f53b42f14f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 22:55:33 +0100
Subject: [PATCH 209/428] Initialise electron moments and pdf before they are
 used

Initialized to 0.0 at an early stage, so still need to be set properly
before being used. This might hide some failure to initialize proprely,
but helps ensure reproducibility at least.

Also only write external source terms for electrons to output file if
electron external sources are active.
---
 moment_kinetics/src/file_io.jl            | 66 +++++++++++++----------
 moment_kinetics/src/initial_conditions.jl | 17 ++++++
 2 files changed, 55 insertions(+), 28 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 5c63f5d3d..9d694b989 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1205,22 +1205,30 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                       description="electron species thermal speed",
                                       units="c_ref")
 
-    external_source_electron_amplitude = create_dynamic_variable!(
-        dynamic, "external_source_electron_amplitude", mk_float, z, r;
-        parallel_io=parallel_io, description="Amplitude of the external source for electrons",
-        units="n_ref/c_ref^3*c_ref/L_ref")
-    external_source_electron_density_amplitude = create_dynamic_variable!(
-        dynamic, "external_source_electron_density_amplitude", mk_float, z, r;
-        parallel_io=parallel_io, description="Amplitude of the external density source for electrons",
-        units="n_ref*c_ref/L_ref")
-    external_source_electron_momentum_amplitude = create_dynamic_variable!(
-        dynamic, "external_source_electron_momentum_amplitude", mk_float, z, r;
-        parallel_io=parallel_io, description="Amplitude of the external momentum source for electrons",
-        units="m_ref*n_ref*c_ref*c_ref/L_ref")
-    external_source_electron_pressure_amplitude = create_dynamic_variable!(
-        dynamic, "external_source_electron_pressure_amplitude", mk_float, z, r;
-        parallel_io=parallel_io, description="Amplitude of the external pressure source for electrons",
-        units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
+    electron_source_settings = external_source_settings.electron
+    if electron_source_settings.active
+        external_source_electron_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_electron_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external source for electrons",
+            units="n_ref/c_ref^3*c_ref/L_ref")
+        external_source_electron_density_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_electron_density_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external density source for electrons",
+            units="n_ref*c_ref/L_ref")
+        external_source_electron_momentum_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_electron_momentum_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external momentum source for electrons",
+            units="m_ref*n_ref*c_ref*c_ref/L_ref")
+        external_source_electron_pressure_amplitude = create_dynamic_variable!(
+            dynamic, "external_source_electron_pressure_amplitude", mk_float, z, r;
+            parallel_io=parallel_io, description="Amplitude of the external pressure source for electrons",
+            units="m_ref*n_ref*c_ref^2*c_ref/L_ref")
+    else
+        external_source_electron_amplitude = nothing
+        external_source_electron_density_amplitude = nothing
+        external_source_electron_momentum_amplitude = nothing
+        external_source_electron_pressure_amplitude = nothing
+    end
 
     electron_constraints_A_coefficient =
         create_dynamic_variable!(dynamic, "electron_constraints_A_coefficient", mk_float, z, r;
@@ -1959,18 +1967,20 @@ function write_electron_moments_data_to_binary(moments, t_params,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
                               t_idx, parallel_io, z, r)
-        append_to_dynamic_var(io_moments.external_source_electron_amplitude,
-                              moments.electron.external_source_amplitude, t_idx,
-                              parallel_io, z, r)
-        append_to_dynamic_var(io_moments.external_source_electron_density_amplitude,
-                              moments.electron.external_source_density_amplitude,
-                              t_idx, parallel_io, z, r)
-        append_to_dynamic_var(io_moments.external_source_electron_momentum_amplitude,
-                              moments.electron.external_source_momentum_amplitude,
-                              t_idx, parallel_io, z, r)
-        append_to_dynamic_var(io_moments.external_source_electron_pressure_amplitude,
-                              moments.electron.external_source_pressure_amplitude,
-                              t_idx, parallel_io, z, r)
+        if io_moments.external_source_electron_amplitude !== nothing
+            append_to_dynamic_var(io_moments.external_source_electron_amplitude,
+                                  moments.electron.external_source_amplitude, t_idx,
+                                  parallel_io, z, r)
+            append_to_dynamic_var(io_moments.external_source_electron_density_amplitude,
+                                  moments.electron.external_source_density_amplitude,
+                                  t_idx, parallel_io, z, r)
+            append_to_dynamic_var(io_moments.external_source_electron_momentum_amplitude,
+                                  moments.electron.external_source_momentum_amplitude,
+                                  t_idx, parallel_io, z, r)
+            append_to_dynamic_var(io_moments.external_source_electron_pressure_amplitude,
+                                  moments.electron.external_source_pressure_amplitude,
+                                  t_idx, parallel_io, z, r)
+        end
         append_to_dynamic_var(io_moments.electron_constraints_A_coefficient,
                               moments.electron.constraints_A_coefficient, t_idx,
                               parallel_io, z, r)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index cca5bdf8b..a266010e4 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -197,6 +197,23 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                      pdf.ion.norm, vpa, vperp, z, r, composition,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
 
+        begin_serial_region()
+        @serial_region begin
+            # If electrons are being used, they will be initialized properly later. Here
+            # we only set the values to avoid false positives from the debug checks
+            # (when @debug_track_initialized is active).
+            moments.electron.dens .= 0.0
+            moments.electron.upar .= 0.0
+            moments.electron.ppar .= 0.0
+            moments.electron.qpar .= 0.0
+            moments.electron.constraints_A_coefficient .= 1.0
+            moments.electron.constraints_B_coefficient .= 0.0
+            moments.electron.constraints_C_coefficient .= 0.0
+            if composition.electron_physics == kinetic_electrons
+                pdf.electron.norm .= 0.0
+            end
+        end
+
         initialize_external_source_amplitude!(moments, external_source_settings, vperp,
                                               vzeta, vr, n_neutral_species)
         initialize_external_source_controller_integral!(moments, external_source_settings,

From dc16805ab124e040fefbc4f7c8d3c326fe71fb68 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 31 Mar 2024 21:22:11 +0100
Subject: [PATCH 210/428] Better status prints from kinetic electrons with
 distributed MPI

When using distributed MPI, the lower and upper boundary values of phi
are not on the same process. To deal with this, print status updates
from processes that contain the first point in z and from processes that
contain the last point in z.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 82adb74a4..74199042a 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -464,8 +464,12 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
         if (mod(t_params.step_counter[] - initial_step_counter,100) == 0)
             begin_serial_region()
-            if global_rank[] == 0
-                println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
+            @serial_region begin
+                if z.irank == 0 && z.irank == z.nrank - 1
+                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
+                elseif z.irank == 0
+                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary_lower: ", phi[1,1], " residual: ", residual)
+                end
             end
         end
         if (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)

From fea14997b9fc7c61ef9ee1875d4394455c7150aa Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 1 Apr 2024 10:21:37 +0100
Subject: [PATCH 211/428] Set debug flags in ldiv!() method for
 MPIDebugSharedArray

---
 moment_kinetics/src/communication.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/communication.jl b/moment_kinetics/src/communication.jl
index cd2bc53cd..4280ad4a8 100644
--- a/moment_kinetics/src/communication.jl
+++ b/moment_kinetics/src/communication.jl
@@ -508,12 +508,19 @@ end
     # SparseArray A
     import LinearAlgebra: ldiv!, Factorization
     function ldiv!(Y::DebugMPISharedArray, A::Factorization, B::DebugMPISharedArray)
+        @debug_track_initialized begin
+            Y.is_initialized .= 1
+        end
+        Y.is_written .= true
+        Y.accessed[] = true
         return ldiv!(Y.data, A, B.data)
     end
 
     import MPI: Buffer
     function Buffer(A::DebugMPISharedArray)
-        A.is_initialized .= 1
+        @debug_track_initialized begin
+            A.is_initialized .= 1
+        end
         A.is_read .= true
         A.is_written .= true
         A.accessed[] = true

From fdbeafcb327a97f1e762dd93d3520d02faf8b2fd Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 2 Apr 2024 08:51:09 +0100
Subject: [PATCH 212/428] Allow `nstep = 0`

Lets code be run just to write out the initial profiles.
---
 moment_kinetics/src/time_advance.jl | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 5b5a1c6c0..31a656ca2 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -292,13 +292,21 @@ function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_relo
 
     end_time = code_time + t_input.dt * t_input.nstep
     epsilon = 1.e-11
-    moments_output_times = [code_time + i*t_input.dt
-                            for i ∈ t_input.nwrite:t_input.nwrite:t_input.nstep]
+    if t_input.nwrite == 0
+        moments_output_times = [end_time]
+    else
+        moments_output_times = [code_time + i*t_input.dt
+                                for i ∈ t_input.nwrite:t_input.nwrite:t_input.nstep]
+    end
     if moments_output_times[end] < end_time - epsilon
         push!(moments_output_times, end_time)
     end
-    dfns_output_times = [code_time + i*t_input.dt
-                         for i ∈ t_input.nwrite_dfns:t_input.nwrite_dfns:t_input.nstep]
+    if t_input.nwrite_dfns == 0
+        dfns_output_times = [end_time]
+    else
+        dfns_output_times = [code_time + i*t_input.dt
+                             for i ∈ t_input.nwrite_dfns:t_input.nwrite_dfns:t_input.nstep]
+    end
     if dfns_output_times[end] < end_time - epsilon
         push!(dfns_output_times, end_time)
     end
@@ -1172,6 +1180,11 @@ function time_advance!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyr
     iwrite_dfns = 2
     finish_now = false
     t_params.step_counter[] = 1
+    if t ≥ t_params.end_time - epsilon
+        # User must have requested zero output steps, i.e. to just write out the initial
+        # profiles
+        return nothing
+    end
     while true
         if t_params.split_operators
             # MRH NOT SUPPORTED

From 7a5aa6acbd4910afbbeee6a60752d0b4d00daac0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 2 Apr 2024 09:29:49 +0100
Subject: [PATCH 213/428] Set electrons t_params.next_output_time[] when not
 doing initial output

When not doing initial electron output, the
`t_params.next_output_time[]` was not set correctly, so its value would
be either 0.0, or whatever was left over from the electron
initialisation. This might terminate the kinetic electron update
incorrectly, or otherwise cause errors. Instead set it to `Inf` so that
the output-writing blocks in update_electron_pdf_with_time_advance!()
are never triggered.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 74199042a..d50258b34 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -182,7 +182,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     # t_params are set relative to 0.0).
     moments_output_times = t_params.moments_output_times .+ initial_time
     dfns_output_times = t_params.dfns_output_times .+ initial_time
-    if io_initial_electron !== nothing
+    if io_initial_electron === nothing
+        t_params.next_output_time[] = Inf
+    else
         t_params.next_output_time[] = dfns_output_times[1]
     end
 

From 79303fd30d664daaa7946fa4db9988a2f12b26be Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 2 Apr 2024 12:20:18 +0100
Subject: [PATCH 214/428] Do bounds checks in electron boundary condition

If something goes wrong, danger of out-of-bounds array accesses after
distribution cut-off index is calculated. Add checks for this to avoid
silent memory corruption.
---
 .../src/electron_kinetic_equation.jl          | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index d50258b34..27dce02e7 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -714,6 +714,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             #    end
             #end
             upar0 = upar[1,ir]
+            if upar0 >= 0.0
+                error("In lower-z boundary condition, upar0=$upar0 has the wrong sign.")
+            elseif isnan(upar0)
+                error("In lower-z boundary condition, upar0=$upar0.")
+            end
             #println("before pdf left ", pdf[:,1,1,ir])
             while upar_integral > upar0 && ivpa_max > 1
                 ivpa += 1
@@ -724,6 +729,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                 upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
                 #println("left ", ivpa, " ", upar_integral, " ", upar0)
             end
+            if ivpa ≥ vpa.n
+                error("In lower-z boundary condition, upar_integral failed to reach upar0")
+            elseif ivpa ≤ 1
+                error("In lower-z boundary condition, ivpa=$ivpa ≤ 1")
+            end
             integral_excess = upar_integral - upar0
             fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
             #println("fraction_of_pdf=", fraction_of_pdf)
@@ -961,6 +971,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             #    end
             #end
             upar_end = upar[end,ir]
+            if upar_end <= 0.0
+                error("In upper-z boundary condition, upar_end=$upar_end has the wrong sign.")
+            elseif isnan(upar_end)
+                error("In upper-z boundary condition, upar_end=$upar_end.")
+            end
             #println("before pdf ", pdf[:,1,end,ir])
             while upar_integral < upar_end && ivpa > 1
                 ivpa -= 1
@@ -971,6 +986,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                 upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
                 #println("right ", ivpa, " ", upar_integral, " ", upar_end)
             end
+            if ivpa ≤ 1
+                error("In upper-z boundary condition, upar_integral failed to reach upar_end")
+            elseif ivpa ≥ vpa.n
+                error("In upper-z boundary condition, ivpa=$ivpa ≥ vpa.n=$(vpa.n)")
+            end
             integral_excess = upar_integral - upar_end
             fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
             #println("B fraction_of_pdf=", fraction_of_pdf)

From f22dec2b45f3f746bff9bbf1484464aed72101ef Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 2 Apr 2024 13:22:35 +0100
Subject: [PATCH 215/428] Set `restart_had_kinetic_electrons` on all processes

---
 moment_kinetics/src/load_data.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index aa7185b16..c8e2464ca 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -15,7 +15,7 @@ export read_distributed_zr_data!
 
 using ..array_allocation: allocate_float, allocate_int
 using ..calculus: derivative!
-using ..communication: setup_distributed_memory_MPI
+using ..communication
 using ..coordinates: coordinate, define_coordinate
 using ..electron_vpa_advection: update_electron_speed_vpa!
 using ..electron_z_advection: update_electron_speed_z!
@@ -875,6 +875,9 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
         end
     end
 
+    restart_had_kinetic_electrons = MPI.Bcast(restart_had_kinetic_electrons, 0,
+                                              comm_block[])
+
     return code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
            previous_runs_info, time_index, restart_had_kinetic_electrons
 end

From d5488a0c73884468514e117b952f2d537d736fb3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 2 Apr 2024 15:12:07 +0100
Subject: [PATCH 216/428] Move `using Quadmath` from moment_kinetics_input to
 time_advance

This should have been moved when setup of `time_info` objects was
refactored.
---
 moment_kinetics/src/moment_kinetics_input.jl | 1 -
 moment_kinetics/src/time_advance.jl          | 1 +
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 6dd0488ac..09fe51e59 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -22,7 +22,6 @@ using ..reference_parameters
 using ..geo: init_magnetic_geometry, setup_geometry_input
 
 using MPI
-using Quadmath
 using TOML
 using UUIDs
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 31a656ca2..39173b377 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -8,6 +8,7 @@ export allocate_advection_structs
 export setup_dummy_and_buffer_arrays
 
 using MPI
+using Quadmath
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_shared_float, allocate_shared_bool
 using ..communication

From a5c324216d69d3b6c0f52adc44b6c20932cff114 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 19:56:00 +0100
Subject: [PATCH 217/428] Check that t_params.dt[] does not become NaN

When using adaptive timestepping, if somehow `t_params.dt[]` becomes
`NaN`, the timestep will fail continuously until the maximum number of
iterations is reached, which would be a waste of time, so check
explicitly for NaN values and terminate the time-loop if they are found.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 moment_kinetics/src/time_advance.jl              | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 27dce02e7..b5df3c3c8 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -252,7 +252,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     # evolve (artificially) in time until the residual is less than the tolerance
     while (!electron_pdf_converged
            && (t_params.step_counter[] - initial_step_counter < max_electron_pdf_iterations)
-           && t_params.dt[] > 0.0)
+           && t_params.dt[] > 0.0 && !isnan(t_params.dt[]))
         for istage ∈ 1:t_params.n_rk_stages
             # Set the initial values for this stage to the final values from the previous
             # stage
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 39173b377..262464e15 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1203,7 +1203,7 @@ function time_advance!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyr
         # update the time
         t += t_params.previous_dt[]
 
-        if t ≥ t_params.end_time - epsilon || t_params.dt[] < 0.0
+        if t ≥ t_params.end_time - epsilon || t_params.dt[] < 0.0 || isnan(t_params.dt[])
             # Ensure all output is written at the final step
             # Negative t_params.dt[] indicates the time stepping has failed, so stop and
             # write output.

From 052e8de0bce16cccba4e48c3249251da61b95dc9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 19:57:48 +0100
Subject: [PATCH 218/428] Only activate electron sources when electrons are
 Brakinskii or kinetic

Otherwise uninitialized arrays might be read.
---
 moment_kinetics/src/external_sources.jl      | 10 +++++++---
 moment_kinetics/src/moment_kinetics_input.jl |  3 ++-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index f6ea63dc6..0a620eb2e 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -20,7 +20,7 @@ using ..array_allocation: allocate_float, allocate_shared_float
 using ..calculus
 using ..communication
 using ..coordinates
-using ..input_structs: set_defaults_and_check_section!, Dict_to_NamedTuple
+using ..input_structs
 using ..looping
 
 using MPI
@@ -38,7 +38,7 @@ and z-coordinates.
 Returns a NamedTuple `(ion=ion_source_settings, neutral=neutral_source_settings)`
 containing two NamedTuples of settings.
 """
-function setup_external_sources!(input_dict, r, z)
+function setup_external_sources!(input_dict, r, z, electron_physics)
     function get_settings(neutrals)
         input = set_defaults_and_check_section!(
                      input_dict, neutrals ? "neutral_source" : "ion_source";
@@ -220,7 +220,11 @@ function setup_external_sources!(input_dict, r, z)
     end
 
     ion_settings = get_settings(false)
-    electron_settings = get_electron_settings(ion_settings)
+    if electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        electron_settings = get_electron_settings(ion_settings)
+    else
+        electron_settings = (active=false,)
+    end
     neutral_settings = get_settings(true)
 
     return (ion=ion_settings, electron=electron_settings, neutral=neutral_settings)
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 09fe51e59..ac7eff62b 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -705,7 +705,8 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     vzeta, vzeta_spectral = define_coordinate(vzeta_immutable, io_immutable.parallel_io;
                                               ignore_MPI=ignore_MPI)
 
-    external_source_settings = setup_external_sources!(scan_input, r, z)
+    external_source_settings = setup_external_sources!(scan_input, r, z,
+                                                       composition.electron_physics)
 
     if global_rank[] == 0 && save_inputs_to_txt
         # Make file to log some information about inputs into.

From fb746554e375643ce1a1626d615ebac85c224bb1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 19:58:32 +0100
Subject: [PATCH 219/428] Set region type in
 external_electron_source_controller() function

---
 moment_kinetics/src/external_sources.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 0a620eb2e..998d93d71 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -891,6 +891,7 @@ up to date.
 """
 function external_electron_source_controller!(fvec_in, moments, electron_source_settings,
                                               dt)
+    begin_r_z_region()
 
     is = 1
     electron_moments = moments.electron

From 591c2a3f567cf232263e33734e4cf8d98f7da3c2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 21:55:39 +0100
Subject: [PATCH 220/428] Debug tests for kinetic electrons

---
 .../debug_test/kinetic_electron_inputs.jl     | 96 +++++++++++++++++++
 .../debug_test/kinetic_electron_tests.jl      | 23 +++++
 2 files changed, 119 insertions(+)
 create mode 100644 moment_kinetics/debug_test/kinetic_electron_inputs.jl
 create mode 100644 moment_kinetics/debug_test/kinetic_electron_tests.jl

diff --git a/moment_kinetics/debug_test/kinetic_electron_inputs.jl b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
new file mode 100644
index 000000000..d5af2dd64
--- /dev/null
+++ b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
@@ -0,0 +1,96 @@
+test_type = "Kinetic electron"
+
+test_input = Dict("n_ion_species" => 1,
+                  "n_neutral_species" => 1,
+                  "electron_physics" => "kinetic_electrons",
+                  "run_name" => "kinetic_electron",
+                  "base_directory" => test_output_directory,
+                  "evolve_moments_density" => true,
+                  "evolve_moments_parallel_flow" => true,
+                  "evolve_moments_parallel_pressure" => true,
+                  "evolve_moments_conservation" => true,
+                  "recycling_fraction" => 0.5,
+                  "krook_collisions" => true,
+                  "T_e" => 0.2,
+                  "T_wall" => 0.1,
+                  "initial_density1" => 1.0,
+                  "initial_temperature1" => 1.0,
+                  "z_IC_option1" => "gaussian",
+                  "z_IC_density_amplitude1" => 0.001,
+                  "z_IC_density_phase1" => 0.0,
+                  "z_IC_upar_amplitude1" => 1.0,
+                  "z_IC_upar_phase1" => 0.0,
+                  "z_IC_temperature_amplitude1" => 0.0,
+                  "z_IC_temperature_phase1" => 0.0,
+                  "vpa_IC_option1" => "gaussian",
+                  "vpa_IC_density_amplitude1" => 1.0,
+                  "vpa_IC_density_phase1" => 0.0,
+                  "vpa_IC_upar_amplitude1" => 0.0,
+                  "vpa_IC_upar_phase1" => 0.0,
+                  "vpa_IC_temperature_amplitude1" => 0.0,
+                  "vpa_IC_temperature_phase1" => 0.0,
+                  "initial_density2" => 1.0,
+                  "initial_temperature2" => 1.0,
+                  "z_IC_option2" => "gaussian",
+                  "z_IC_density_amplitude2" => 0.001,
+                  "z_IC_density_phase2" => 0.0,
+                  "z_IC_upar_amplitude2" => -1.0,
+                  "z_IC_upar_phase2" => 0.0,
+                  "z_IC_temperature_amplitude2" => 0.0,
+                  "z_IC_temperature_phase2" => 0.0,
+                  "vpa_IC_option2" => "gaussian",
+                  "vpa_IC_density_amplitude2" => 1.0,
+                  "vpa_IC_density_phase2" => 0.0,
+                  "vpa_IC_upar_amplitude2" => 0.0,
+                  "vpa_IC_upar_phase2" => 0.0,
+                  "vpa_IC_temperature_amplitude2" => 0.0,
+                  "vpa_IC_temperature_phase2" => 0.0,
+                  "charge_exchange_frequency" => 0.75,
+                  "ionization_frequency" => 0.5,
+                  "constant_ionization_rate" => false,
+                  "timestepping" => Dict{String,Any}("type" => "Fekete4(3)",
+                                                     "nstep" => 3,
+                                                     "dt" => 1.0e-8,
+                                                     "minimum_dt" => 1.0e-8,
+                                                     "CFL_prefactor" => 1.0,
+                                                     "step_update_prefactor" => 0.4,
+                                                     "nwrite" => 2,
+                                                     "split_operators" => false),
+                  "electron_timestepping" => Dict{String,Any}("type" => "Fekete4(3)",
+                                                              "nstep" => 10,
+                                                              "dt" => 2.0e-11,
+                                                              "minimum_dt" => 2.0e-11,
+                                                              "initialization_residual_value" => 1.e10,
+                                                              "converged_residual_value" => 1.e10,
+                                                              "nwrite" => 10000,
+                                                              "nwrite_dfns" => 10000,
+                                                              "no_restart" => true),
+                  "r_ngrid" => 1,
+                  "r_nelement" => 1,
+                  "z_ngrid" => 3,
+                  "z_nelement" => 24,
+                  "z_bc" => "wall",
+                  "z_discretization" => "chebyshev_pseudospectral",
+                  "z_element_spacing_option" => "sqrt",
+                  "vpa_ngrid" => 3,
+                  "vpa_nelement" => 4,
+                  "vpa_L" => 6.0,
+                  "vpa_bc" => "zero",
+                  "vpa_discretization" => "chebyshev_pseudospectral",
+                  "vz_ngrid" => 3,
+                  "vz_nelement" => 4,
+                  "vz_L" => 6.0,
+                  "vz_bc" => "zero",
+                  "vz_discretization" => "chebyshev_pseudospectral",
+                  "ion_source" => Dict("active" => true,
+                                       "z_profile" => "gaussian",
+                                       "z_width" => 0.125,
+                                       "source_strength" => 2.0,
+                                       "source_T" => 2.0),
+                  "numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0,
+                                                              "vpa_dissipation_coefficient" => 1e-2))
+
+
+test_input_list = [
+     test_input,
+    ]
diff --git a/moment_kinetics/debug_test/kinetic_electron_tests.jl b/moment_kinetics/debug_test/kinetic_electron_tests.jl
new file mode 100644
index 000000000..f246eed9a
--- /dev/null
+++ b/moment_kinetics/debug_test/kinetic_electron_tests.jl
@@ -0,0 +1,23 @@
+module KineticElectronDebug
+
+# Debug test using wall boundary conditions.
+
+include("setup.jl")
+
+# Create a temporary directory for test output
+test_output_directory = get_MPI_tempdir()
+mkpath(test_output_directory)
+
+
+# Input parameters for the test
+include("kinetic_electron_inputs.jl")
+
+# Defines the test functions, using variables defined in the *_inputs.jl file
+include("runtest_template.jl")
+
+end # KineticElectronDebug
+
+
+using .KineticElectronDebug
+
+KineticElectronDebug.runtests()

From e8f60a2841e028fbe32971914e53feac69292d83 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 23:05:20 +0100
Subject: [PATCH 221/428] Add input option to skip possibly restarting for
 kinetic electron init

Useful for debug tests, which re-run in the same directory, and which
also cannot use glob() because the test directory needs to use an
absolute path.
---
 moment_kinetics/src/initial_conditions.jl    | 8 ++++++--
 moment_kinetics/src/moment_kinetics_input.jl | 1 +
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index a266010e4..e310c58d9 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -463,8 +463,12 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
     # if using kinetic electrons
     if composition.electron_physics == kinetic_electrons
         begin_serial_region()
-        restart_filename = get_default_restart_filename(io_input, "initial_electron";
-                                                        error_if_no_file_found=false)
+        if t_input.no_restart
+            restart_filename = nothing
+        else
+            restart_filename = get_default_restart_filename(io_input, "initial_electron";
+                                                            error_if_no_file_found=false)
+        end
         if restart_filename === nothing
             # No file to restart from
             previous_runs_info = nothing
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index ac7eff62b..ae2691f10 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -244,6 +244,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         maximum_dt=timestepping_section["maximum_dt"] * sqrt(composition.me_over_mi),
         high_precision_error_sum=timestepping_section["high_precision_error_sum"],
         initialization_residual_value=1.0,
+        no_restart=false,
        )
     if electron_timestepping_section["nwrite"] === nothing
         electron_timestepping_section["nwrite"] = electron_timestepping_section["nstep"]

From 7b623a6917b1f6f1fee188b0854aedc38d5d3d12 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 6 Apr 2024 23:06:59 +0100
Subject: [PATCH 222/428] Shared-mem fix for dfns_output_times for electrons

---
 moment_kinetics/src/electron_kinetic_equation.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index b5df3c3c8..d5a33d8d0 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -182,10 +182,12 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     # t_params are set relative to 0.0).
     moments_output_times = t_params.moments_output_times .+ initial_time
     dfns_output_times = t_params.dfns_output_times .+ initial_time
-    if io_initial_electron === nothing
-        t_params.next_output_time[] = Inf
-    else
-        t_params.next_output_time[] = dfns_output_times[1]
+    @serial_region begin
+        if io_initial_electron === nothing
+            t_params.next_output_time[] = Inf
+        else
+            t_params.next_output_time[] = dfns_output_times[1]
+        end
     end
 
     #z_speedup_fac = 20.0

From c01cb4df26dbf5b62964989d5b2a2e68594a8633 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 09:40:31 +0100
Subject: [PATCH 223/428] Let value in setindex!() method for
 DebugMPISharedArray be any Number

Allows `setindex!()` for DebugMPISharedArray to be used with, for
example, `Float128` values, not just `mk_float`. Any type that is OK for
the underlying `setindex!()` of `Array{mk_float,N}` should be supported.
---
 moment_kinetics/src/communication.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/communication.jl b/moment_kinetics/src/communication.jl
index 4280ad4a8..2550d0358 100644
--- a/moment_kinetics/src/communication.jl
+++ b/moment_kinetics/src/communication.jl
@@ -457,7 +457,7 @@ end
         A.accessed[] = true
         return getindex(A.data, I...)
     end
-    function Base.setindex!(A::DebugMPISharedArray{T, N}, v::T, I::Vararg{mk_int,N}) where {T, N}
+    function Base.setindex!(A::DebugMPISharedArray{T, N}, v::Number, I::Vararg{mk_int,N}) where {T, N}
         @debug_track_initialized begin
             A.is_initialized[I...] = 1
         end

From d2d58eb1a2d4bce29b420c131fe0579dc76f2103 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 09:42:46 +0100
Subject: [PATCH 224/428] Reload dSdt when restarting

Prevents writing uninitialized values at the zero'th time point of the
restarted simulation.
---
 moment_kinetics/src/load_data.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index c8e2464ca..571974718 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -654,6 +654,10 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                              r_range, z_range, restart_r,
                                              restart_r_spectral, restart_z,
                                              restart_z_spectral, interpolation_needed)
+            moments.ion.dSdt .= reload_moment("entropy_production", dynamic, time_index,
+                                              r, z, r_range, z_range, restart_r,
+                                              restart_r_spectral, restart_z,
+                                              restart_z_spectral, interpolation_needed)
             if z.irank == 0
                 if "chodura_integral_lower" ∈ keys(dynamic)
                     moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower",

From 8fab804d9a5ea05b68104856504d1686bd858944 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 09:49:58 +0100
Subject: [PATCH 225/428] Reload *_constraints_*_coefficient, to avoid using
 uninitialized values

---
 moment_kinetics/src/load_data.jl | 85 ++++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 571974718..be2e803db 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -658,6 +658,35 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                               r, z, r_range, z_range, restart_r,
                                               restart_r_spectral, restart_z,
                                               restart_z_spectral, interpolation_needed)
+            if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+                if "ion_constraints_A_coefficient" ∈ keys(dynamic)
+                    moments.ion.constraints_A_coefficient .=
+                        reload_moment("ion_constraints_A_coefficient", dynamic,
+                                      time_index, r, z, r_range, z_range, restart_r,
+                                      restart_r_spectral, restart_z, restart_z_spectral,
+                                      interpolation_needed)
+                elseif moments.ion.constraints_A_coefficient !== nothing
+                    moments.ion.constraints_A_coefficient .= 0.0
+                end
+                if "ion_constraints_B_coefficient" ∈ keys(dynamic)
+                    moments.ion.constraints_B_coefficient .=
+                        reload_moment("ion_constraints_B_coefficient", dynamic,
+                                      time_index, r, z, r_range, z_range, restart_r,
+                                      restart_r_spectral, restart_z, restart_z_spectral,
+                                      interpolation_needed)
+                elseif moments.ion.constraints_B_coefficient !== nothing
+                    moments.ion.constraints_B_coefficient .= 0.0
+                end
+                if "ion_constraints_C_coefficient" ∈ keys(dynamic)
+                    moments.ion.constraints_C_coefficient .=
+                        reload_moment("ion_constraints_C_coefficient", dynamic,
+                                      time_index, r, z, r_range, z_range, restart_r,
+                                      restart_r_spectral, restart_z, restart_z_spectral,
+                                      interpolation_needed)
+                elseif moments.ion.constraints_C_coefficient !== nothing
+                    moments.ion.constraints_C_coefficient .= 0.0
+                end
+            end
             if z.irank == 0
                 if "chodura_integral_lower" ∈ keys(dynamic)
                     moments.ion.chodura_integral_lower .= load_slice(dynamic, "chodura_integral_lower",
@@ -750,6 +779,33 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                                            restart_r_spectral, restart_z,
                                                            restart_z_spectral,
                                                            interpolation_needed)
+            if "electron_constraints_A_coefficient" ∈ keys(dynamic)
+                moments.electron.constraints_A_coefficient .=
+                    reload_electron_moment("electron_constraints_A_coefficient", dynamic,
+                                           time_index, r, z, r_range, z_range, restart_r,
+                                           restart_r_spectral, restart_z,
+                                           restart_z_spectral, interpolation_needed)
+            else
+                moments.electron.constraints_A_coefficient .= 0.0
+            end
+            if "electron_constraints_B_coefficient" ∈ keys(dynamic)
+                moments.electron.constraints_B_coefficient .=
+                    reload_electron_moment("electron_constraints_B_coefficient", dynamic,
+                                           time_index, r, z, r_range, z_range, restart_r,
+                                           restart_r_spectral, restart_z,
+                                           restart_z_spectral, interpolation_needed)
+            else
+                moments.electron.constraints_B_coefficient .= 0.0
+            end
+            if "electron_constraints_C_coefficient" ∈ keys(dynamic)
+                moments.electron.constraints_C_coefficient .=
+                    reload_electron_moment("electron_constraints_C_coefficient", dynamic,
+                                           time_index, r, z, r_range, z_range, restart_r,
+                                           restart_r_spectral, restart_z,
+                                           restart_z_spectral, interpolation_needed)
+            else
+                moments.electron.constraints_C_coefficient .= 0.0
+            end
 
             # For now, electrons are always fully moment_kinetic
             restart_electron_evolve_density, restart_electron_evolve_upar,
@@ -802,6 +858,35 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                                      restart_r, restart_r_spectral,
                                                      restart_z, restart_z_spectral,
                                                      interpolation_needed)
+                if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
+                    if "neutral_constraints_A_coefficient" ∈ keys(dynamic)
+                        moments.neutral.constraints_A_coefficient .=
+                            reload_moment("neutral_constraints_A_coefficient", dynamic,
+                                          time_index, r, z, r_range, z_range, restart_r,
+                                          restart_r_spectral, restart_z, restart_z_spectral,
+                                          interpolation_needed)
+                    elseif moments.neutral.constraints_A_coefficient !== nothing
+                        moments.neutral.constraints_A_coefficient .= 0.0
+                    end
+                    if "neutral_constraints_B_coefficient" ∈ keys(dynamic)
+                        moments.neutral.constraints_B_coefficient .=
+                            reload_moment("neutral_constraints_B_coefficient", dynamic,
+                                          time_index, r, z, r_range, z_range, restart_r,
+                                          restart_r_spectral, restart_z, restart_z_spectral,
+                                          interpolation_needed)
+                    elseif moments.neutral.constraints_B_coefficient !== nothing
+                        moments.neutral.constraints_B_coefficient .= 0.0
+                    end
+                    if "neutral_constraints_C_coefficient" ∈ keys(dynamic)
+                        moments.neutral.constraints_C_coefficient .=
+                            reload_moment("neutral_constraints_C_coefficient", dynamic,
+                                          time_index, r, z, r_range, z_range, restart_r,
+                                          restart_r_spectral, restart_z, restart_z_spectral,
+                                          interpolation_needed)
+                    elseif moments.neutral.constraints_C_coefficient !== nothing
+                        moments.neutral.constraints_C_coefficient .= 0.0
+                    end
+                end
 
                 if "external_source_neutral_controller_integral" ∈ get_variable_keys(dynamic) &&
                         length(moments.neutral.external_source_controller_integral) == 1

From a9b0f2516dbbfc546ef826fe84b5b2aac7f22662 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 10:15:32 +0100
Subject: [PATCH 226/428] Zero-initialize moments.electron.temp

Avoid using uninitialized values.
---
 moment_kinetics/src/initial_conditions.jl | 1 +
 moment_kinetics/src/moment_kinetics.jl    | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index e310c58d9..47c8f9416 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -206,6 +206,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
             moments.electron.upar .= 0.0
             moments.electron.ppar .= 0.0
             moments.electron.qpar .= 0.0
+            moments.electron.temp .= 0.0
             moments.electron.constraints_A_coefficient .= 1.0
             moments.electron.constraints_B_coefficient .= 0.0
             moments.electron.constraints_C_coefficient .= 0.0
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 4abe967eb..ac52b2911 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -310,6 +310,11 @@ function setup_moment_kinetics(input_dict::AbstractDict;
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
                                     vz)
 
+        begin_serial_region()
+        @serial_region begin
+            @. moments.electron.temp = moments.electron.vth^2
+        end
+
         # Re-initialize the source amplitude here instead of loading it from the restart
         # file so that we can change the settings between restarts.
         initialize_external_source_amplitude!(moments, external_source_settings, vperp,

From a129274f329d92b94e8233a8a95324338c73b25f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 11:26:26 +0100
Subject: [PATCH 227/428] Fix region type when updating electron_ppar

---
 moment_kinetics/src/time_advance.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 262464e15..4fb05ee7b 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1707,11 +1707,13 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         # and is already updated;
         # otherwise update assuming electron temperature is fixed in time
         if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+            begin_r_z_region()
             @loop_r_z ir iz begin
                 new_scratch.electron_ppar[iz,ir] = (rk_coefs[1]*moments.electron.ppar[iz,ir] 
                     + rk_coefs[2]*old_scratch.electron_ppar[iz,ir] + rk_coefs[3]*new_scratch.electron_ppar[iz,ir])
             end
         else
+            begin_r_z_region()
             @loop_r_z ir iz begin
                 new_scratch.electron_ppar[iz,ir] = 0.5 * new_scratch.electron_density[iz,ir] *
                                                    moments.electron.vth[iz,ir]^2

From e79a1e37e7082d741d28f51057e02ec37f2bfb9b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Apr 2024 12:11:25 +0100
Subject: [PATCH 228/428] When ion timestep fails, reset electron pdf to its
 value before the step

Hopefully will reduce the number of iterations to converge the electron
pdf after a failed step, by providing a better (?) initial guess.
---
 .../src/electron_fluid_equations.jl           |  4 +--
 moment_kinetics/src/initial_conditions.jl     | 15 ++++++++--
 moment_kinetics/src/moment_kinetics.jl        |  7 +++++
 .../src/moment_kinetics_structs.jl            | 10 ++++++-
 moment_kinetics/src/time_advance.jl           | 29 ++++++++++++++++---
 5 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 11e553ac9..33300b6e9 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -12,7 +12,7 @@ using ..communication
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
 using ..input_structs: braginskii_fluid, kinetic_electrons
-using ..moment_kinetics_structs: pdf_substruct
+using ..moment_kinetics_structs: electron_pdf_substruct
 using ..velocity_moments: integrate_over_vspace
 
 using MPI
@@ -300,7 +300,7 @@ function calculate_electron_qpar!(electron_moments, pdf, ppar_e, upar_e, upar_i,
             end
         elseif electron_model == kinetic_electrons
             # use the modified electron pdf to calculate the electron heat flux
-            if isa(pdf, pdf_substruct)
+            if isa(pdf, electron_pdf_substruct)
                 electron_pdf = pdf.norm
             else
                 electron_pdf = pdf
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 47c8f9416..7051aaf32 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -25,8 +25,8 @@ using ..em_fields: update_phi!
 using ..file_io: setup_initial_electron_io, write_initial_electron_state,
                  finish_initial_electron_io
 using ..load_data: reload_electron_data!
-using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, pdf_struct, moments_struct,
-                                 boundary_distributions_struct
+using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, electron_pdf_substruct,
+                                 pdf_struct, moments_struct, boundary_distributions_struct
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
 using ..velocity_moments: create_moments_ion, create_moments_electron, create_moments_neutral
@@ -106,7 +106,10 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
         # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
         # to emulate ion and neutral behaviour
         pdf_electron_buffer = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
-        electron_substruct = pdf_substruct(pdf_electron_norm, pdf_electron_buffer)
+        pdf_before_ion_timestep = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
+        electron_substruct = electron_pdf_substruct(pdf_electron_norm,
+                                                    pdf_electron_buffer,
+                                                    pdf_before_ion_timestep)
     else
         electron_substruct = nothing
     end
@@ -568,6 +571,12 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         initial_time=electron_pseudotime,
                                         initial_output_counter=n_debug_outputs)
 
+        begin_r_z_vperp_vpa_region()
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir] =
+                pdf.electron.norm[ivpa,ivperp,iz,ir]
+        end
+
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
         t_idx = n_debug_outputs + 1
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index ac52b2911..2aabcafc5 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -314,6 +314,13 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         @serial_region begin
             @. moments.electron.temp = moments.electron.vth^2
         end
+        if composition.electron_physics == kinetic_electrons
+            begin_r_z_vperp_vpa_region()
+            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+                pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir] =
+                    pdf.electron.norm[ivpa,ivperp,iz,ir]
+            end
+        end
 
         # Re-initialize the source amplitude here instead of loading it from the restart
         # file so that we can change the settings between restarts.
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index ed3bf13b3..edcd847a5 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -314,6 +314,14 @@ struct pdf_substruct{n_distribution}
     buffer::MPISharedArray{mk_float,n_distribution} # for collision operator terms when pdfs must be interpolated onto different velocity space grids
 end
 
+"""
+"""
+struct electron_pdf_substruct{n_distribution}
+    norm::MPISharedArray{mk_float,n_distribution}
+    buffer::MPISharedArray{mk_float,n_distribution} # for collision operator terms when pdfs must be interpolated onto different velocity space grids
+    pdf_before_ion_timestep::MPISharedArray{mk_float,n_distribution}
+end
+
 # struct of structs neatly contains i+n info?
 """
 """
@@ -321,7 +329,7 @@ struct pdf_struct
     #ion particles: s + r + z + vperp + vpa
     ion::pdf_substruct{5}
     # electron particles: r + z + vperp + vpa
-    electron::Union{pdf_substruct{4},Nothing}
+    electron::Union{electron_pdf_substruct{4},Nothing}
     #neutral particles: s + r + z + vzeta + vr + vz
     neutral::pdf_substruct{6}
 end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 4fb05ee7b..847368d82 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1844,7 +1844,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         # moment derivatives, because the timstep might need to be re-done with a smaller
         # dt, in which case scratch[t_params.n_rk_stages+1] will be reset to the values
         # from the beginning of the timestep here.
-        adaptive_timestep_update!(scratch, t, t_params, moments, fields, composition,
+        adaptive_timestep_update!(scratch, t, t_params, pdf, moments, fields, composition,
                                   collisions, geometry, external_source_settings,
                                   advect_objects, r, z, vperp, vpa, vzeta, vr, vz)
         # Re-do this in case adaptive_timestep_update re-arranged the `scratch` vector
@@ -1883,9 +1883,10 @@ end
 Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
-function adaptive_timestep_update!(scratch, t, t_params, moments, fields, composition,
-                                   collisions, geometry, external_source_settings,
-                                   advect_objects, r, z, vperp, vpa, vzeta, vr, vz)
+function adaptive_timestep_update!(scratch, t, t_params, pdf, moments, fields,
+                                   composition, collisions, geometry,
+                                   external_source_settings, advect_objects, r, z, vperp,
+                                   vpa, vzeta, vr, vz)
     #error_norm_method = "Linf"
     error_norm_method = "L2"
 
@@ -2105,6 +2106,26 @@ function adaptive_timestep_update!(scratch, t, t_params, moments, fields, compos
 
     adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
                                        total_points, current_dt, error_norm_method)
+    if composition.electron_physics == kinetic_electrons
+        if t_params.previous_dt[] == 0.0
+            # Reset electron pdf to its value at the beginning of this step.
+            begin_r_z_vperp_vpa_region()
+            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+                pdf.electron.norm[ivpa,ivperp,iz,ir] =
+                    pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir]
+                scratch[1].pdf_electron[ivpa,ivperp,iz,ir] =
+                    pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir]
+            end
+        else
+            # Store the current value, which will be the value at the beginning of the
+            # next step.
+            begin_r_z_vperp_vpa_region()
+            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+                pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir] =
+                    pdf.electron.norm[ivpa,ivperp,iz,ir]
+            end
+        end
+    end
 
     return nothing
 end

From 51e5f91cc3678ac377bdeb01e1b1aa25d5bf588b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Apr 2024 09:13:16 +0100
Subject: [PATCH 229/428] Fix communication of delta_phi when calculating phi
 from Epar

---
 moment_kinetics/src/em_fields.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 7c2a5f243..491601eb8 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -167,7 +167,7 @@ function calculate_phi_from_Epar!(phi, Epar, z)
         # Add contributions to integral along z from processes at smaller z-values than
         # this one.
         this_delta_phi = phi[end,:] .- phi[1,:]
-        for irank ∈ 1:z.nrank-1
+        for irank ∈ 0:z.nrank-2
             delta_phi = MPI.bcast(this_delta_phi, z.comm; root=irank)
             if z.irank > irank
                 @loop_r_z ir iz begin

From 15b02223a60b7d898edccad954621f4242314980 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Apr 2024 10:21:13 +0100
Subject: [PATCH 230/428] Make kinetic-electron debug test have 'failed'
 timesteps

Increasing the 'dt' setting (for both ions/neutrals and electrons) means
that the adaptive-timestep solver starts with a 'failed' timestep, so we
check more code branches.
---
 moment_kinetics/debug_test/kinetic_electron_inputs.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/debug_test/kinetic_electron_inputs.jl b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
index d5af2dd64..697d95ee7 100644
--- a/moment_kinetics/debug_test/kinetic_electron_inputs.jl
+++ b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
@@ -50,7 +50,7 @@ test_input = Dict("n_ion_species" => 1,
                   "constant_ionization_rate" => false,
                   "timestepping" => Dict{String,Any}("type" => "Fekete4(3)",
                                                      "nstep" => 3,
-                                                     "dt" => 1.0e-8,
+                                                     "dt" => 2.0e-8,
                                                      "minimum_dt" => 1.0e-8,
                                                      "CFL_prefactor" => 1.0,
                                                      "step_update_prefactor" => 0.4,
@@ -58,7 +58,7 @@ test_input = Dict("n_ion_species" => 1,
                                                      "split_operators" => false),
                   "electron_timestepping" => Dict{String,Any}("type" => "Fekete4(3)",
                                                               "nstep" => 10,
-                                                              "dt" => 2.0e-11,
+                                                              "dt" => 4.0e-11,
                                                               "minimum_dt" => 2.0e-11,
                                                               "initialization_residual_value" => 1.e10,
                                                               "converged_residual_value" => 1.e10,

From ec20a224db5db8a220facdc224c39483d3a68e98 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Apr 2024 11:01:02 +0100
Subject: [PATCH 231/428] Remove debug write-out of 'dt' in runge_kutta

Not totally sure this was safe when running in parallel.
---
 moment_kinetics/src/runge_kutta.jl | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 207366f04..8e985afdd 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -771,12 +771,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             # So far `error_norms` is the sum of squares of the errors. Now that summation
             # is finished, need to divide by total number of points and take square-root.
             error_norms .= sqrt.(error_norms ./ total_points)
-            open("debug$(global_size[]).txt", "a") do io
-                for e in error_norms
-                    print(io, e, " ")
-                end
-                println(io, t_params.dt[], " ;")
-            end
 
             # Weight the error from each variable equally by taking the mean, so the
             # larger number of points in the distribution functions does not mean that

From 35b3be81448cfc54f1e375574bb0289e960a447b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Apr 2024 20:22:02 +0100
Subject: [PATCH 232/428] Make coordinate arguments to define_io_coordinates!()
 optional

If `nothing` is passed for any of the arguments, just do not define that
coordinate.
---
 moment_kinetics/src/file_io.jl   | 67 +++++++++++++++++++-------------
 moment_kinetics/src/load_data.jl | 31 +++++++--------
 2 files changed, 56 insertions(+), 42 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 9d694b989..dc6e2428d 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -324,7 +324,7 @@ end
 """
 open output file to save the initial electron pressure and distribution function
 """
-function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, composition,
+function setup_initial_electron_io(io_input, vpa, vperp, z, r, composition,
                                    collisions, evolve_density, evolve_upar, evolve_ppar,
                                    external_source_settings, input_dict,
                                    restart_time_index, previous_runs_info)
@@ -364,7 +364,8 @@ function setup_initial_electron_io(io_input, vz, vr, vzeta, vpa, vperp, z, r, co
         write_input!(fid, input_dict, parallel_io)
 
         ### define coordinate dimensions ###
-        define_io_coordinates!(fid, vz, vr, vzeta, vpa, vperp, z, r, parallel_io)
+        define_io_coordinates!(fid, nothing, nothing, nothing, vpa, vperp, z, r,
+                               parallel_io)
 
         ### create variables for time-dependent quantities ###
         dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
@@ -667,12 +668,16 @@ function define_io_coordinates!(fid, vz, vr, vzeta, vpa, vperp, z, r, parallel_i
     @serial_region begin
         # create the "coords" group that will contain coordinate information
         coords = create_io_group(fid, "coords")
-        # create the "z" sub-group of "coords" that will contain z coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, z, "z", "spatial coordinate z", parallel_io)
-        # create the "r" sub-group of "coords" that will contain r coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, r, "r", "spatial coordinate r", parallel_io)
+        if z !== nothing
+            # create the "z" sub-group of "coords" that will contain z coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, z, "z", "spatial coordinate z", parallel_io)
+        end
+        if r !== nothing
+            # create the "r" sub-group of "coords" that will contain r coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, r, "r", "spatial coordinate r", parallel_io)
+        end
 
         if parallel_io
             # Parallel I/O produces a single file, so effectively a 'single block'
@@ -700,24 +705,34 @@ function define_io_coordinates!(fid, vz, vr, vzeta, vpa, vperp, z, r, parallel_i
                                 parallel_io=parallel_io, description="number of zr blocks")
         end
 
-        # create the "vz" sub-group of "coords" that will contain vz coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, vz, "vz", "velocity coordinate v_z", parallel_io)
-        # create the "vr" sub-group of "coords" that will contain vr coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, vr, "vr", "velocity coordinate v_r", parallel_io)
-        # create the "vzeta" sub-group of "coords" that will contain vzeta coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, vzeta, "vzeta", "velocity coordinate v_zeta",
-                              parallel_io)
-        # create the "vpa" sub-group of "coords" that will contain vpa coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, vpa, "vpa", "velocity coordinate v_parallel",
-                              parallel_io)
-        # create the "vperp" sub-group of "coords" that will contain vperp coordinate info,
-        # including total number of grid points and grid point locations
-        define_io_coordinate!(coords, vperp, "vperp", "velocity coordinate v_perp",
-                              parallel_io)
+        if vz !== nothing
+            # create the "vz" sub-group of "coords" that will contain vz coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, vz, "vz", "velocity coordinate v_z", parallel_io)
+        end
+        if vr !== nothing
+            # create the "vr" sub-group of "coords" that will contain vr coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, vr, "vr", "velocity coordinate v_r", parallel_io)
+        end
+        if vzeta !== nothing
+            # create the "vzeta" sub-group of "coords" that will contain vzeta coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, vzeta, "vzeta", "velocity coordinate v_zeta",
+                                  parallel_io)
+        end
+        if vpa !== nothing
+            # create the "vpa" sub-group of "coords" that will contain vpa coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, vpa, "vpa", "velocity coordinate v_parallel",
+                                  parallel_io)
+        end
+        if vperp !== nothing
+            # create the "vperp" sub-group of "coords" that will contain vperp coordinate info,
+            # including total number of grid points and grid point locations
+            define_io_coordinate!(coords, vperp, "vperp", "velocity coordinate v_perp",
+                                  parallel_io)
+        end
     end
 
     return nothing
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index be2e803db..dadd7bfbf 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -226,7 +226,12 @@ function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=no
     overview = get_group(fid, "overview")
     parallel_io = load_variable(overview, "parallel_io")
 
-    coord_group = get_group(get_group(fid, "coords"), name)
+    coords_group = get_group(fid, "coords")
+    if name ∈ keys(coords_group)
+        coord_group = get_group(coords_group, name)
+    else
+        return nothing, nothing, nothing
+    end
 
     ngrid = load_variable(coord_group, "ngrid")
     n_local = load_variable(coord_group, "n_local")
@@ -602,14 +607,18 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(
-                x.name => x.n != restart_x.n || !all(isapprox.(x.grid, restart_x.grid))
+                x.name => (restart_x !== nothing
+                           && (x.n != restart_x.n
+                               || !all(isapprox.(x.grid, restart_x.grid))))
                 for (x, restart_x) ∈ ((z, restart_z), (r, restart_r),
                                       (vperp, restart_vperp), (vpa, restart_vpa),
                                       (vzeta, restart_vzeta), (vr, restart_vr),
                                       (vz, restart_vz)))
 
             neutral_1V = (vzeta.n_global == 1 && vr.n_global == 1)
-            restart_neutral_1V = (restart_vzeta.n_global == 1 && restart_vr.n_global == 1)
+            restart_neutral_1V = ((restart_vzeta === nothing
+                                   || restart_vzeta.n_global == 1)
+                                  && (restart_vr === nothing || restart_vr.n_global == 1))
             if any(geometry.bzeta .!= 0.0) && ((neutral_1V && !restart_neutral_1V) ||
                                                (!neutral_1V && restart_neutral_1V))
                 # One but not the other of the run being restarted from and this run are
@@ -1011,18 +1020,6 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
                 for (x, restart_x) ∈ ((z, restart_z), (r, restart_r),
                                       (vperp, restart_vperp), (vpa, restart_vpa)))
 
-            neutral_1V = (vzeta.n_global == 1 && vr.n_global == 1)
-            restart_neutral_1V = (restart_vzeta.n_global == 1 && restart_vr.n_global == 1)
-            if geometry.bzeta != 0.0 && ((neutral_1V && !restart_neutral_1V) ||
-                                         (!neutral_1V && restart_neutral_1V))
-                # One but not the other of the run being restarted from and this run are
-                # 1V, but the interpolation below does not allow for vz and vpa being in
-                # different directions. Therefore interpolation between 1V and 3V cases
-                # only works (at the moment!) if bzeta=0.
-                error("Interpolation between 1V and 3V neutrals not yet supported when "
-                      * "bzeta!=0.")
-            end
-
             code_time = load_slice(dynamic, "time", time_index)
 
             r_range, z_range, vperp_range, vpa_range, vzeta_range, vr_range, vz_range =
@@ -1155,7 +1152,9 @@ function get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp, res
                            restart_vzeta, restart_vr, restart_vz)
     if parallel_io
         function get_range(coord)
-            if coord.irank == coord.nrank - 1
+            if coord === nothing
+                return 1:0
+            elseif coord.irank == coord.nrank - 1
                 return coord.global_io_range
             else
                 # Need to modify the range to load the end-point that is duplicated on

From 9ad08ba7b89282741b78d185e130083e19fd3c3b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Apr 2024 20:40:30 +0100
Subject: [PATCH 233/428] Add option for 'debug I/O' for electrons

When `debug_io = true` is passed in the `[electron_timestepping]`
section, at each call to `update_electron_pdf_with_time_advance!()`
during the ion/neutral time advance (not during initialisation, where
output can already be writen to the `*.initial_electron.*` file) a new
output file is created (overwriting any existing file, e.g. from
previous ion timesteps) with label `.electron_debug.` in which electron
quantities are saved from debugging.
---
 .../src/electron_kinetic_equation.jl          | 64 +++++++++++--------
 moment_kinetics/src/file_io.jl                | 32 +++++-----
 moment_kinetics/src/initial_conditions.jl     | 39 +++++------
 moment_kinetics/src/input_structs.jl          |  3 +-
 moment_kinetics/src/moment_kinetics_input.jl  |  1 +
 moment_kinetics/src/time_advance.jl           | 56 +++++++++++-----
 6 files changed, 118 insertions(+), 77 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index d5a33d8d0..661f74593 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -21,7 +21,7 @@ using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
 using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
 using ..external_sources: external_electron_source!
-using ..file_io: write_initial_electron_state, finish_initial_electron_io
+using ..file_io: setup_electron_io, write_electron_state, finish_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
 using ..runge_kutta: rk_update_variable!, rk_error_variable!, local_error_norm,
@@ -59,7 +59,7 @@ OUTPUT:
 function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_spectral,
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
+        max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
         initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
@@ -72,7 +72,7 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
             collisions, composition, r, z, vperp, vpa, z_spectral, vperp_spectral,
             vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
             external_source_settings, num_diss_params, max_electron_pdf_iterations;
-            io_initial_electron=io_initial_electron, initial_time=initial_time,
+            io_electron=io_electron, initial_time=initial_time,
             initial_output_counter=initial_output_counter,
             residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
@@ -129,7 +129,7 @@ The electron kinetic equation is:
     vpa_spectral = struct containing spectral information for the vpa-coordinate
     scratch_dummy = dummy arrays to be used for temporary storage
     max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
-    io_initial_electron = info struct for binary file I/O
+    io_electron = info struct for binary file I/O
     initial_time = initial value for the (pseudo-)time
 OUTPUT:
     pdf = updated (modified) electron pdf
@@ -137,7 +137,7 @@ OUTPUT:
 function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, collisions,
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_initial_electron=nothing, initial_time=0.0,
+        max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
         initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
 
     begin_r_z_region()
@@ -182,8 +182,25 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     # t_params are set relative to 0.0).
     moments_output_times = t_params.moments_output_times .+ initial_time
     dfns_output_times = t_params.dfns_output_times .+ initial_time
+    if io_electron === nothing && t_params.debug_io !== nothing
+        # Overwrite the debug output file with the output from this call to
+        # update_electron_pdf_with_time_advance!().
+        io_electron = setup_electron_io(t_params.debug_io[1], vpa, vperp, z, r,
+                                        composition, collisions,
+                                        moments.evolve_density, moments.evolve_upar,
+                                        moments.evolve_ppar, external_source_settings,
+                                        t_params.debug_io[2], -1, nothing,
+                                        "electron_debug")
+        do_debug_io = true
+        debug_io_nwrite = t_params.debug_io[3]
+    else
+        do_debug_io = false
+    end
     @serial_region begin
-        if io_initial_electron === nothing
+        if io_electron === nothing || do_debug_io
+            # When doing debug_io, we don't want to make the adaptive timestep adjust to
+            # output at specific times - instead we just output after a fixed number of
+            # steps, however long those steps were.
             t_params.next_output_time[] = Inf
         else
             t_params.next_output_time[] = dfns_output_times[1]
@@ -245,10 +262,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     begin_serial_region()
     output_counter += 1
     @serial_region begin
-        if io_initial_electron !== nothing
-            write_initial_electron_state(scratch[1].pdf_electron, moments, t_params, time,
-                                         io_initial_electron, output_counter, r, z, vperp,
-                                         vpa)
+        if io_electron !== nothing
+            write_electron_state(scratch[1].pdf_electron, moments, t_params, time,
+                                 io_electron, output_counter, r, z, vperp, vpa)
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
@@ -476,7 +492,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
         end
-        if (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)
+        if ((time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)
+            || (do_debug_io && (t_params.step_counter[] % debug_io_nwrite == 0)))
+
             begin_serial_region()
             @serial_region begin
                 if text_output
@@ -509,11 +527,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
             @serial_region begin
-                if io_initial_electron !== nothing
-                    write_initial_electron_state(scratch[t_params.n_rk_stages+1].pdf_electron,
-                                                 moments, t_params, time,
-                                                 io_initial_electron, output_counter, r,
-                                                 z, vperp, vpa)
+                if io_electron !== nothing
+                    write_electron_state(scratch[t_params.n_rk_stages+1].pdf_electron,
+                                         moments, t_params, time, io_electron,
+                                         output_counter, r, z, vperp, vpa)
                 end
             end
         end
@@ -549,20 +566,17 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             close(io_pdf)
             close(io_pdf_stages)
         end
-        if !electron_pdf_converged
-            # need to exit or handle this appropriately
-
-            if io_initial_electron !== nothing
+        if !electron_pdf_converged || do_debug_io
+            if io_electron !== nothing && io_electron !== true
                 output_counter += 1
-                write_initial_electron_state(final_scratch_pdf, moments, t_params, time,
-                                             io_initial_electron, output_counter, r, z,
-                                             vperp, vpa)
-                finish_initial_electron_io(io_initial_electron)
+                write_electron_state(final_scratch_pdf, moments, t_params, time,
+                                     io_electron, output_counter, r, z, vperp, vpa)
+                finish_electron_io(io_electron)
             end
-
         end
     end
     if !electron_pdf_converged
+        # need to exit or handle this appropriately
         error("!!!max number of iterations for electron pdf update exceeded!!!\n"
               * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
     end
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index dc6e2428d..1060dc34a 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -324,10 +324,10 @@ end
 """
 open output file to save the initial electron pressure and distribution function
 """
-function setup_initial_electron_io(io_input, vpa, vperp, z, r, composition,
-                                   collisions, evolve_density, evolve_upar, evolve_ppar,
-                                   external_source_settings, input_dict,
-                                   restart_time_index, previous_runs_info)
+function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
+                           evolve_density, evolve_upar, evolve_ppar,
+                           external_source_settings, input_dict, restart_time_index,
+                           previous_runs_info, prefix_label)
     begin_serial_region()
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -341,7 +341,7 @@ function setup_initial_electron_io(io_input, vpa, vperp, z, r, composition,
         parallel_io = io_input.parallel_io
         io_comm = comm_inter_block[]
 
-        electrons_prefix = string(out_prefix, ".initial_electron")
+        electrons_prefix = string(out_prefix, ".$prefix_label")
         if !parallel_io
             electrons_prefix *= ".$(iblock_index[])"
         end
@@ -394,8 +394,8 @@ function setup_initial_electron_io(io_input, vpa, vperp, z, r, composition,
 
         return file_info
     end
-    # For other processes in the block, return (nothing, nothing, nothing)
-    return nothing, nothing, nothing
+    # For other processes in the block, return nothing
+    return nothing
 end
 
 """
@@ -2199,15 +2199,13 @@ function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
 end
 
 """
-    write_initial_electron_state(pdf, moments, t_params, t,
-                                 io_initial_electron, t_idx, r, z,
-                                 vperp, vpa)
+    write_electron_state(pdf, moments, t_params, t, io_initial_electron,
+                         t_idx, r, z, vperp, vpa)
 
 Write the electron state to an output file.
 """
-function write_initial_electron_state(pdf, moments, t_params, t,
-                                      io_or_file_info_initial_electron, t_idx, r, z,
-                                      vperp, vpa)
+function write_electron_state(pdf, moments, t_params, t, io_or_file_info_initial_electron,
+                              t_idx, r, z, vperp, vpa)
 
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -2269,13 +2267,15 @@ end
 """
 close output files for electron initialization
 """
-function finish_initial_electron_io(
-        binary_initial_electron::Union{io_initial_electron_info,Tuple,Nothing})
+function finish_electron_io(
+        binary_initial_electron::Union{io_initial_electron_info,Tuple,Nothing,Bool})
 
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        if binary_initial_electron !== nothing && !isa(binary_initial_electron, Tuple)
+        if (binary_initial_electron !== nothing && !isa(binary_initial_electron, Tuple)
+            && !isa(binary_initial_electron, Bool))
+
             close(binary_initial_electron.fid)
         end
     end
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 7051aaf32..cc05e910b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -22,8 +22,7 @@ using ..external_sources
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
 using ..em_fields: update_phi!
-using ..file_io: setup_initial_electron_io, write_initial_electron_state,
-                 finish_initial_electron_io
+using ..file_io: setup_electron_io, write_electron_state, finish_electron_io
 using ..load_data: reload_electron_data!
 using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, electron_pdf_substruct,
                                  pdf_struct, moments_struct, boundary_distributions_struct
@@ -497,19 +496,22 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             # might have been loaded from a restart file).
             code_time = MPI.Bcast(code_time, 0, comm_block[])
         end
-        io_initial_electron = nothing
+        # Set to `true` rather than `nothing` so that processes that are not writing
+        # output (i.e. not rank-0 of their shared-memory block) know that 'initial
+        # electron output' is being written (so that they know not to activate 'debug
+        # I/O').
+        io_initial_electron = true
         @serial_region begin
             # Setup I/O for initial electron state
-            io_initial_electron = setup_initial_electron_io(io_input, vz, vr, vzeta, vpa,
-                                                            vperp, z, r, composition,
-                                                            collisions,
-                                                            moments.evolve_density,
-                                                            moments.evolve_upar,
-                                                            moments.evolve_ppar,
-                                                            external_source_settings,
-                                                            input_dict,
-                                                            restart_time_index,
-                                                            previous_runs_info)
+            io_initial_electron = setup_electron_io(io_input, vpa, vperp, z, r,
+                                                    composition, collisions,
+                                                    moments.evolve_density,
+                                                    moments.evolve_upar,
+                                                    moments.evolve_ppar,
+                                                    external_source_settings, input_dict,
+                                                    restart_time_index,
+                                                    previous_runs_info,
+                                                    "initial_electron")
 
             # update the electron pdf in the first scratch
             scratch[1].pdf_electron .= pdf.electron.norm
@@ -550,7 +552,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         t_params, collisions, composition,
                                         external_source_settings, num_diss_params,
                                         max_electron_pdf_iterations;
-                                        io_initial_electron=io_initial_electron,
+                                        io_electron=io_initial_electron,
                                         initial_time=code_time,
                                         residual_tolerance=t_input.initialization_residual_value,
                                         evolve_ppar=true)
@@ -567,7 +569,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         t_params, collisions, composition,
                                         external_source_settings, num_diss_params,
                                         max_electron_pdf_iterations;
-                                        io_initial_electron=io_initial_electron,
+                                        io_electron=io_initial_electron,
                                         initial_time=electron_pseudotime,
                                         initial_output_counter=n_debug_outputs)
 
@@ -580,10 +582,9 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
         t_idx = n_debug_outputs + 1
-        write_initial_electron_state(pdf.electron.norm, moments, t_params,
-                                     electron_pseudotime, io_initial_electron, t_idx, r,
-                                     z, vperp, vpa)
-        finish_initial_electron_io(io_initial_electron)
+        write_electron_state(pdf.electron.norm, moments, t_params, electron_pseudotime,
+                             io_initial_electron, t_idx, r, z, vperp, vpa)
+        finish_electron_io(io_initial_electron)
 
     end
     return nothing
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 75c05697f..40821898b 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -38,7 +38,7 @@ end
 an option but known at compile time when a `time_info` struct is passed as a function
 argument.
 """
-struct time_info{Terrorsum <: Real, T_electron}
+struct time_info{Terrorsum <: Real, T_debug_output, T_electron}
     nstep::mk_int
     end_time::mk_float
     dt::MPISharedArray{mk_float,1}
@@ -75,6 +75,7 @@ struct time_info{Terrorsum <: Real, T_electron}
     converged_residual_value::mk_float
     use_manufactured_solns_for_advance::Bool
     stopfile::String
+    debug_io::T_debug_output # Currently only used by electrons
     electron::T_electron
 end
 
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index ae2691f10..0edbd1ca3 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -245,6 +245,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         high_precision_error_sum=timestepping_section["high_precision_error_sum"],
         initialization_residual_value=1.0,
         no_restart=false,
+        debug_io=false,
        )
     if electron_timestepping_section["nwrite"] === nothing
         electron_timestepping_section["nwrite"] = electron_timestepping_section["nstep"]
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 847368d82..a55b0a405 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -274,7 +274,7 @@ If something is passed in `electron`, it is stored in the `electron_t_params` me
 the returned `time_info`.
 """
 function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_reload,
-                         manufactured_solns_input; electron=nothing)
+                         manufactured_solns_input, io_input, input_dict; electron=nothing)
     dt_shared = allocate_shared_float(1)
     previous_dt_shared = allocate_shared_float(1)
     next_output_time = allocate_shared_float(1)
@@ -320,19 +320,39 @@ function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_relo
     else
         error_sum_zero = 0.0
     end
+    if electron === nothing
+        # Setting up time_info for electrons.
+        # Store io_input as the debug_io variable so we can use it to open the debug
+        # output file.
+        if t_input.debug_io !== false
+            if !isa(t_input.debug_io, mk_int)
+                error("`debug_io` input should be an integer, giving the number of steps "
+                      * "between writes, if it is passed")
+            end
+            debug_io = (io_input, input_dict, t_input.debug_io)
+        else
+            debug_io = nothing
+        end
+        electron_t_params = nothing
+    elseif electron === false
+        debug_io = nothing
+        electron_t_params = nothing
+    else
+        debug_io = nothing
+        electron_t_params = electron
+    end
     return time_info(t_input.nstep, end_time, dt_shared, previous_dt_shared, next_output_time,
-                         dt_before_output, dt_before_last_fail, CFL_prefactor,
-                         step_to_output, Ref(0), Ref(0), mk_int[], mk_int[],
-                         moments_output_times, dfns_output_times, t_input.type, rk_coefs,
-                         n_rk_stages, rk_order, adaptive, low_storage, t_input.rtol,
-                         t_input.atol, t_input.atol_upar, t_input.step_update_prefactor,
-                         t_input.max_increase_factor,
-                         t_input.max_increase_factor_near_last_fail,
-                         t_input.last_fail_proximity_factor, t_input.minimum_dt,
-                         t_input.maximum_dt, error_sum_zero, t_input.split_operators,
-                         t_input.steady_state_residual, t_input.converged_residual_value,
-                         manufactured_solns_input.use_for_advance, t_input.stopfile_name,
-                         electron)
+                     dt_before_output, dt_before_last_fail, CFL_prefactor, step_to_output,
+                     Ref(0), Ref(0), mk_int[], mk_int[], moments_output_times,
+                     dfns_output_times, t_input.type, rk_coefs, n_rk_stages, rk_order,
+                     adaptive, low_storage, t_input.rtol, t_input.atol, t_input.atol_upar,
+                     t_input.step_update_prefactor, t_input.max_increase_factor,
+                     t_input.max_increase_factor_near_last_fail,
+                     t_input.last_fail_proximity_factor, t_input.minimum_dt,
+                     t_input.maximum_dt, error_sum_zero, t_input.split_operators,
+                     t_input.steady_state_residual, t_input.converged_residual_value,
+                     manufactured_solns_input.use_for_advance, t_input.stopfile_name,
+                     debug_io, electron_t_params)
 end
 
 """
@@ -363,7 +383,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
         electron_t_params = setup_time_info(t_input.electron_t_input, 0.0,
                                             electron_dt_reload,
                                             electron_dt_before_last_fail_reload,
-                                            manufactured_solns_input)
+                                            manufactured_solns_input, io_input,
+                                            input_dict)
         # Make Vectors that count which variable caused timestep limits and timestep failures
         # the right length. Do this setup even when not using adaptive timestepping, because
         # it is easier than modifying the file I/O according to whether we are using adaptive
@@ -380,10 +401,13 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
         # electron ppar
         push!(electron_t_params.failure_caused_by, 0)
     else
-        electron_t_params = nothing
+        # Pass `false` rather than `nothing` to `setup_time_info()` call for ions, which
+        # indicates that 'debug_io' should never be set up for ions.
+        electron_t_params = false
     end
     t_params = setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_reload,
-                               manufactured_solns_input; electron=electron_t_params)
+                               manufactured_solns_input, io_input, input_dict;
+                               electron=electron_t_params)
 
     # Make Vectors that count which variable caused timestep limits and timestep failures
     # the right length. Do this setup even when not using adaptive timestepping, because

From 0ae14c1c459a7b90f67aee5a01cfceec013f4089 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 9 Apr 2024 19:32:03 +0100
Subject: [PATCH 234/428] Add 'git' module in setup for Marconi

---
 machines/marconi/julia.env | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/machines/marconi/julia.env b/machines/marconi/julia.env
index 43b24569e..dfb006696 100644
--- a/machines/marconi/julia.env
+++ b/machines/marconi/julia.env
@@ -2,7 +2,7 @@
 
 module purge
 module load env-skl profile/base profile/advanced
-module load gnu/7.3.0 openmpi/3.1.4--gnu--7.3.0 intel/pe-xe-2018--binary python/3.9.4
+module load gnu/7.3.0 openmpi/3.1.4--gnu--7.3.0 intel/pe-xe-2018--binary python/3.9.4 git/2.17
 # Needed because Julia's Cairo library complains about libz<1.2.9
 module load zlib/1.2.11--intel--pe-xe-2018--binary
 

From 7bb15d6d559c141725a8c9171e730d462a3cd5e1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 9 Apr 2024 22:20:13 +0100
Subject: [PATCH 235/428] Update post-processing to work with 'electron_debug'
 files

---
 .../src/makie_post_processing.jl              | 35 +++++++++++++++---
 moment_kinetics/src/load_data.jl              | 36 +++++++++++++------
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index dfcf3933c..4b57b0116 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -562,11 +562,36 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         nz_min = 1
     end
     if dfns && has_run_info
-        nvperp_min = minimum(ri.vperp.n for ri in run_info if ri !== nothing)
-        nvpa_min = minimum(ri.vpa.n for ri in run_info if ri !== nothing)
-        nvzeta_min = minimum(ri.vzeta.n for ri in run_info if ri !== nothing)
-        nvr_min = minimum(ri.vr.n for ri in run_info if ri !== nothing)
-        nvz_min = minimum(ri.vz.n for ri in run_info if ri !== nothing)
+        if any(ri.vperp !== nothing for ri ∈ run_info)
+            nvperp_min = minimum(ri.vperp.n for ri in run_info
+                                 if ri !== nothing && ri.vperp !== nothing)
+        else
+            nvperp_min = 1
+        end
+        if any(ri.vpa !== nothing for ri ∈ run_info)
+            nvpa_min = minimum(ri.vpa.n for ri in run_info
+                               if ri !== nothing && ri.vpa !== nothing)
+        else
+            nvpa_min = 1
+        end
+        if any(ri.vzeta !== nothing for ri ∈ run_info)
+            nvzeta_min = minimum(ri.vzeta.n for ri in run_info
+                                 if ri !== nothing && ri.vzeta !== nothing)
+        else
+            nvzeta_min = 1
+        end
+        if any(ri.vr !== nothing for ri ∈ run_info)
+            nvr_min = minimum(ri.vr.n for ri in run_info
+                              if ri !== nothing && ri.vr !== nothing)
+        else
+            nvr_min = 1
+        end
+        if any(ri.vz !== nothing for ri ∈ run_info)
+            nvz_min = minimum(ri.vz.n for ri in run_info
+                              if ri !== nothing && ri.vz !== nothing)
+        else
+            nvz_min = 1
+        end
     else
         nvperp_min = 1
         nvpa_min = 1
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index dadd7bfbf..9d048a4fc 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3300,6 +3300,7 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
               * "`(String, Nothing)`. Got $run_dir")
     end
 
+    electron_debug = false
     if isfile(this_run_dir)
         # this_run_dir is actually a filename. Assume it is a moment_kinetics output file
         # and infer the directory and the run_name from the filename.
@@ -3313,6 +3314,9 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
             run_name = split(filename, ".dfns.")[1]
         elseif occursin(".initial_electron.", filename)
             run_name = split(filename, ".initial_electron.")[1]
+        elseif occursin(".electron_debug.", filename)
+            run_name = split(filename, ".electron_debug.")[1]
+            electron_debug = true
         else
             error("Cannot recognise '$this_run_dir/$filename' as a moment_kinetics output file")
         end
@@ -3356,7 +3360,11 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
     end
 
     if initial_electron
-        ext = "initial_electron"
+        if electron_debug
+            ext = "electron_debug"
+        else
+            ext = "initial_electron"
+        end
     elseif dfns
         ext = "dfns"
     else
@@ -3554,7 +3562,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     if ivperp === nothing
         if :vperp ∈ keys(run_info)
             # v-space coordinates only present if run_info contains distribution functions
-            nvperp = run_info.vperp.n
+            nvperp = run_info.vperp === nothing ? 1 : run_info.vperp.n
             ivperp = 1:nvperp
         else
             nvperp = nothing
@@ -3568,7 +3576,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     if ivpa === nothing
         if :vpa ∈ keys(run_info)
             # v-space coordinates only present if run_info contains distribution functions
-            nvpa = run_info.vpa.n
+            nvpa = run_info.vpa === nothing ? 1 : run_info.vpa.n
             ivpa = 1:nvpa
         else
             nvpa = nothing
@@ -3582,7 +3590,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     if ivzeta === nothing
         if :vzeta ∈ keys(run_info)
             # v-space coordinates only present if run_info contains distribution functions
-            nvzeta = run_info.vzeta.n
+            nvzeta = run_info.vzeta === nothing ? 1 : run_info.vzeta.n
             ivzeta = 1:nvzeta
         else
             nvzeta = nothing
@@ -3596,7 +3604,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     if ivr === nothing
         if :vr ∈ keys(run_info)
             # v-space coordinates only present if run_info contains distribution functions
-            nvr = run_info.vr.n
+            nvr = run_info.vr === nothing ? 1 : run_info.vr.n
             ivr = 1:nvr
         else
             nvr = nothing
@@ -3610,7 +3618,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     if ivz === nothing
         if :vz ∈ keys(run_info)
             # v-space coordinates only present if run_info contains distribution functions
-            nvz = run_info.vz.n
+            nvz = run_info.vz === nothing ? 1 : run_info.vz.n
             ivz = 1:nvz
         else
             nvz = nothing
@@ -4185,8 +4193,12 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
 
         setup_distributed_memory_MPI(1,1,1,1)
         setup_loop_ranges!(0, 1; s=run_info.n_ion_species, sn=run_info.n_neutral_species,
-                           r=nr, z=nz, vperp=nvperp, vpa=nvpa, vzeta=run_info.vzeta.n,
-                           vr=run_info.vr.n, vz=run_info.vz.n)
+                           r=nr, z=nz,
+                           vperp=(run_info.vperp === nothing ? 1 : run_info.vperp.n),
+                           vpa=(run_info.vpa === nothing ? 1 : run_info.vpa.n),
+                           vzeta=(run_info.vzeta === nothing ? 1 : run_info.vzeta.n),
+                           vr=(run_info.vr === nothing ? 1 : run_info.vr.n),
+                           vz=(run_info.vz === nothing ? 1 : run_info.vz.n))
         for it ∈ 1:nt
             begin_serial_region()
             # Only need some struct with a 'speed' variable
@@ -4251,8 +4263,12 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         speed=allocate_float(nvpa, nvperp, nz, nr, nt)
         setup_distributed_memory_MPI(1,1,1,1)
         setup_loop_ranges!(0, 1; s=run_info.n_ion_species, sn=run_info.n_neutral_species,
-                           r=nr, z=nz, vperp=nvperp, vpa=nvpa, vzeta=run_info.vzeta.n,
-                           vr=run_info.vr.n, vz=run_info.vz.n)
+                           r=nr, z=nz,
+                           vperp=(run_info.vperp === nothing ? 1 : run_info.vperp.n),
+                           vpa=(run_info.vpa === nothing ? 1 : run_info.vpa.n),
+                           vzeta=(run_info.vzeta === nothing ? 1 : run_info.vzeta.n),
+                           vr=(run_info.vr === nothing ? 1 : run_info.vr.n),
+                           vz=(run_info.vz === nothing ? 1 : run_info.vz.n))
         for it ∈ 1:nt
             begin_serial_region()
             # Only need some struct with a 'speed' variable

From 03a660eeba1ba695a3417846558c02af609d7038 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 10 Apr 2024 12:05:50 +0100
Subject: [PATCH 236/428] Print 'ion step' debug message

Useful when kinetic electrons are included, would be way too noisy when
running ion simulations...
---
 moment_kinetics/src/time_advance.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index a55b0a405..3c19c8f46 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2248,6 +2248,9 @@ function ssp_rk!(pdf, scratch, t, t_params, vz, vr, vzeta, vpa, vperp, gyrophase
     end
 
     for istage ∈ 1:n_rk_stages
+        if global_rank[] == 0
+            println("ion step ", t_params.step_counter[], ".", istage, " ", t)
+        end
         # do an Euler time advance, with scratch[2] containing the advanced quantities
         # and scratch[1] containing quantities at time level n
         update_solution_vector!(scratch, moments, istage, composition, vpa, vperp, z, r)

From ff230a8f95095aa00656fd2d704cd2fc6c34de71 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 Apr 2024 08:59:30 +0100
Subject: [PATCH 237/428] Move rk_update_variable!() for ppar later in loop

---
 .../src/electron_kinetic_equation.jl            | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 661f74593..4eabd20c7 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -301,14 +301,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                           evolve_ppar=evolve_ppar)
 
             rk_update_variable!(scratch, :pdf_electron, t_params, istage)
-            if evolve_ppar
-                rk_update_variable!(scratch, :electron_ppar, t_params, istage)
-                moments_struct_ppar = moments.electron.ppar
-                scratch_ppar = scratch[istage+1].electron_ppar
-                @loop_r_z ir iz begin
-                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
-                end
-            end
 
             latest_pdf = scratch[istage+1].pdf_electron
             begin_r_z_vperp_vpa_region()
@@ -415,6 +407,15 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     update_derived_moments_and_derivatives()
                 end
             end
+            if evolve_ppar
+                rk_update_variable!(scratch, :electron_ppar, t_params, istage)
+                moments_struct_ppar = moments.electron.ppar
+                scratch_ppar = scratch[istage+1].electron_ppar
+                @loop_r_z ir iz begin
+                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
+                end
+                _block_synchronize()
+            end
         end
 
         # update the time following the pdf update

From 6ab14027817bbc60010d246aaaa938b0f6ecd3ff Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 Apr 2024 09:01:15 +0100
Subject: [PATCH 238/428] Increment t_params.step_counter[] on iteration where
 electrons converge

---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 4eabd20c7..6057648fb 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -538,10 +538,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
         # check to see if the electron pdf satisfies the electron kinetic equation to within the specified tolerance
 
+        t_params.step_counter[] += 1
         if electron_pdf_converged
             break
         end
-        t_params.step_counter[] += 1
     end
     # Update the 'pdf' arrays with the final result
     begin_r_z_vperp_vpa_region()

From ebf7014666c7f593bc52e942b71b0dcc09374584 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 Apr 2024 13:34:46 +0100
Subject: [PATCH 239/428] Create output_dir at the earliest possible point

Useful to ensure that any function that needs to write into `output_dir`
is able to do so.
---
 moment_kinetics/src/moment_kinetics_input.jl | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 0edbd1ca3..836144b44 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -82,6 +82,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     # this is the directory where the simulation data will be stored
     base_directory = get(scan_input, "base_directory", "runs")
     output_dir = joinpath(base_directory, run_name)
+
     # if evolve_moments.density = true, evolve density via continuity eqn
     # and g = f/n via modified drift kinetic equation
     evolve_moments.density = get(scan_input, "evolve_moments_density", false)
@@ -556,6 +557,12 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     #irank_z = 0
     #nrank_z = 0
 
+    # Create output_dir if it does not exist.
+    if global_rank[] == 0
+        mkpath(output_dir)
+    end
+    _block_synchronize()
+
     # 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,
@@ -712,9 +719,6 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
 
     if global_rank[] == 0 && save_inputs_to_txt
         # Make file to log some information about inputs into.
-        # check to see if output_dir exists in the current directory
-        # if not, create it
-        isdir(output_dir) || mkpath(output_dir)
         io = open_ascii_output_file(string(output_dir,"/",run_name), "input")
     else
         io = devnull

From 1463468d875428e80349efa1facaf6312210e605 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 Apr 2024 12:50:38 +0100
Subject: [PATCH 240/428] Save/load FFTW wisdom from rank-0

...to ensure all processes have identical FFTW plans for consistency.
---
 moment_kinetics/src/chebyshev.jl             | 73 +++++++++++++++++---
 moment_kinetics/src/coordinates.jl           |  5 +-
 moment_kinetics/src/load_data.jl             | 55 ++++++++++-----
 moment_kinetics/src/moment_kinetics.jl       |  2 +-
 moment_kinetics/src/moment_kinetics_input.jl | 13 ++--
 5 files changed, 111 insertions(+), 37 deletions(-)

diff --git a/moment_kinetics/src/chebyshev.jl b/moment_kinetics/src/chebyshev.jl
index 98e2704d8..80a7da5c4 100644
--- a/moment_kinetics/src/chebyshev.jl
+++ b/moment_kinetics/src/chebyshev.jl
@@ -12,10 +12,12 @@ export chebyshev_info
 
 using LinearAlgebra: mul!
 using FFTW
+using MPI
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_complex
 using ..clenshaw_curtis: clenshawcurtisweights
 import ..calculus: elementwise_derivative!
+using ..communication
 import ..interpolation: interpolate_to_grid_1d!
 using ..moment_kinetics_structs: discretization_info
 
@@ -52,13 +54,64 @@ end
 create arrays needed for explicit Chebyshev pseudospectral treatment
 and create the plans for the forward and backward fast Fourier transforms
 """
-function setup_chebyshev_pseudospectral(coord)
-    lobatto = setup_chebyshev_pseudospectral_lobatto(coord)
-    radau = setup_chebyshev_pseudospectral_radau(coord)
+function setup_chebyshev_pseudospectral(coord, run_directory; ignore_MPI=false)
+    # First set up the FFTW plans on the (global) root process, then save the 'FFTW
+    # wisdom' and load it on all other processes, to ensure that we use the exact same
+    # FFT algorithms on all processes for consistency.
+    if run_directory === nothing
+        if global_size[] != 1
+            error("run_directory is required by setup_chebyshev_pseudospectral() when "
+                  * "running in parallel, in order to save FFTW wisdom.")
+        end
+        wisdom_filename = nothing
+    else
+        wisdom_filename = joinpath(run_directory, "fftw_wisdom.save")
+    end
+
+    if global_rank[] != 0
+        # Wait for rank-0
+        if !ignore_MPI
+            # Normal case, all processors are creating the coordinate
+            MPI.Barrier(comm_world)
+        elseif comm_inter_block[] != MPI.COMM_NULL
+            # ignore_MPI=true was passed, but non-null communicator exists. This happens
+            # in calls from load_restart_coordinates(), which is only called on
+            # block_rank[]==0.
+            MPI.Barrier(comm_inter_block[])
+        else
+            # Should be serial (e.g. used in post-processing), so no Barrier
+        end
+        # Load wisdom
+        FFTW.import_wisdom(wisdom_filename)
+        fftw_flags = FFTW.WISDOM_ONLY
+    else
+        fftw_flags = FFTW.MEASURE
+    end
+
+    lobatto = setup_chebyshev_pseudospectral_lobatto(coord, fftw_flags)
+    radau = setup_chebyshev_pseudospectral_radau(coord, fftw_flags)
+
+    if global_rank[] == 0
+        if wisdom_filename !== nothing
+            FFTW.export_wisdom(wisdom_filename)
+        end
+        if !ignore_MPI
+            # Normal case, all processors are creating the coordinate
+            MPI.Barrier(comm_world)
+        elseif comm_inter_block[] != MPI.COMM_NULL
+            # ignore_MPI=true was passed, but non-null communicator exists. This happens
+            # in calls from load_restart_coordinates(), which is only called on
+            # block_rank[]==0.
+            MPI.Barrier(comm_inter_block[])
+        else
+            # Should be serial (e.g. used in post-processing), so no Barrier
+        end
+    end
+
     return chebyshev_info(lobatto,radau)
 end
 
-function setup_chebyshev_pseudospectral_lobatto(coord)
+function setup_chebyshev_pseudospectral_lobatto(coord, fftw_flags)
     # ngrid_fft is the number of grid points in the extended domain
     # in z = cos(theta).  this is necessary to turn a cosine transform on [0,π]
     # into a complex transform on [0,2π], which is more efficient in FFTW
@@ -69,8 +122,8 @@ function setup_chebyshev_pseudospectral_lobatto(coord)
     fcheby = allocate_float(coord.ngrid, coord.nelement_local)
     dcheby = allocate_float(coord.ngrid)
     # setup the plans for the forward and backward Fourier transforms
-    forward_transform = plan_fft!(fext, flags=FFTW.MEASURE)
-    backward_transform = plan_ifft!(fext, flags=FFTW.MEASURE)
+    forward_transform = plan_fft!(fext, flags=fftw_flags)
+    backward_transform = plan_ifft!(fext, flags=fftw_flags)
     # create array for differentiation matrix 
     Dmat = allocate_float(coord.ngrid, coord.ngrid)
     cheb_derivative_matrix_elementwise!(Dmat,coord.ngrid)
@@ -81,7 +134,7 @@ function setup_chebyshev_pseudospectral_lobatto(coord)
     return chebyshev_base_info(fext, fcheby, dcheby, forward_transform, backward_transform, Dmat, D0)
 end
 
-function setup_chebyshev_pseudospectral_radau(coord)
+function setup_chebyshev_pseudospectral_radau(coord, fftw_flags)
         # ngrid_fft is the number of grid points in the extended domain
         # in z = cos(theta).  this is necessary to turn a cosine transform on [0,π]
         # into a complex transform on [0,2π], which is more efficient in FFTW
@@ -92,8 +145,8 @@ function setup_chebyshev_pseudospectral_radau(coord)
         fcheby = allocate_float(coord.ngrid, coord.nelement_local)
         dcheby = allocate_float(coord.ngrid)
         # setup the plans for the forward and backward Fourier transforms
-        forward_transform = plan_fft!(fext, flags=FFTW.MEASURE)
-        backward_transform = plan_ifft!(fext, flags=FFTW.MEASURE)
+        forward_transform = plan_fft!(fext, flags=fftw_flags)
+        backward_transform = plan_ifft!(fext, flags=fftw_flags)
         # create array for differentiation matrix 
         Dmat = allocate_float(coord.ngrid, coord.ngrid)
         cheb_derivative_matrix_elementwise_radau_by_FFT!(Dmat, coord, fcheby, dcheby, fext, forward_transform)
@@ -601,7 +654,7 @@ function chebyshev_radau_weights(moments::Array{mk_float,1}, n)
     # create array for moments on extended [0,2π] domain in theta = ArcCos[z]
     fext = allocate_complex(nfft)
     # make fft plan
-    forward_transform = plan_fft!(fext, flags=FFTW.MEASURE)
+    forward_transform = plan_fft!(fext, flags=FFTW.WISDOM_ONLY)
     # assign values of fext from moments 
     @inbounds begin
         for j ∈ 1:n
diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index 776f06033..9e2d472af 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -113,7 +113,8 @@ create arrays associated with a given coordinate,
 setup the coordinate grid, and populate the coordinate structure
 containing all of this information
 """
-function define_coordinate(input, parallel_io::Bool=false; ignore_MPI=false, init_YY::Bool=true)
+function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing,
+                           ignore_MPI=false, init_YY::Bool=true)
     # total number of grid points is ngrid for the first element
     # plus ngrid-1 unique points for each additional element due
     # to the repetition of a point at the element boundary
@@ -202,7 +203,7 @@ function define_coordinate(input, parallel_io::Bool=false; ignore_MPI=false, ini
         # create arrays needed for explicit Chebyshev pseudospectral treatment in this
         # coordinate and create the plans for the forward and backward fast Chebyshev
         # transforms
-        spectral = setup_chebyshev_pseudospectral(coord)
+        spectral = setup_chebyshev_pseudospectral(coord, run_directory; ignore_MPI=ignore_MPI)
         # obtain the local derivatives of the uniform grid with respect to the used grid
         derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
     elseif input.discretization == "gausslegendre_pseudospectral"
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 9d048a4fc..8cdf2407f 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -218,7 +218,7 @@ memory scratch arrays (`ignore_MPI=true` will be passed through to
 [`define_coordinate`](@ref)).
 """
 function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing,
-                              ignore_MPI=false)
+                              run_directory=nothing, ignore_MPI=false)
     if printout
         println("Loading $name coordinate data...")
     end
@@ -312,7 +312,8 @@ function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=no
                        advection_input("default", 0.0, 0.0, 0.0), MPI.COMM_NULL,
                        element_spacing_option)
 
-    coord, spectral = define_coordinate(input, parallel_io; ignore_MPI=ignore_MPI)
+    coord, spectral = define_coordinate(input, parallel_io; run_directory=run_directory,
+                                        ignore_MPI=ignore_MPI)
 
     return coord, spectral, chunk_size
 end
@@ -574,7 +575,7 @@ Reload pdf and moments from an existing output file.
 """
 function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                  restart_prefix_iblock, time_index, composition, geometry,
-                                 r, z, vpa, vperp, vzeta, vr, vz)
+                                 r, z, vpa, vperp, vzeta, vr, vz; run_directory=nothing)
     code_time = 0.0
     dt = nothing
     dt_before_last_fail = nothing
@@ -603,7 +604,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                 restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
                 restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
                 restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
-                                                               vzeta, vr, vz, parallel_io)
+                                                               vzeta, vr, vz, parallel_io;
+                                                               run_directory=run_directory)
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(
@@ -1078,32 +1080,47 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
     return code_time, previous_runs_info, time_index
 end
 
-function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io)
+function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io;
+                                  run_directory=nothing)
     if parallel_io
         restart_z, restart_z_spectral, _ =
-            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_r, restart_r_spectral, _ =
-            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_vpa, restart_vpa_spectral, _ =
-            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_vr, restart_vr_spectral, _ =
-            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
         restart_vz, restart_vz_spectral, _ =
-            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank, ignore_MPI=true)
+            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank,
+                                 run_directory=run_directory, ignore_MPI=true)
     else
-        restart_z, restart_z_spectral, _ = load_coordinate_data(fid, "z"; ignore_MPI=true)
-        restart_r, restart_r_spectral, _ = load_coordinate_data(fid, "r"; ignore_MPI=true)
+        restart_z, restart_z_spectral, _ =
+            load_coordinate_data(fid, "z"; run_directory=run_directory, ignore_MPI=true)
+        restart_r, restart_r_spectral, _ =
+            load_coordinate_data(fid, "r"; run_directory=run_directory, ignore_MPI=true)
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp"; ignore_MPI=true)
-        restart_vpa, restart_vpa_spectral, _ = load_coordinate_data(fid, "vpa"; ignore_MPI=true)
+            load_coordinate_data(fid, "vperp"; run_directory=run_directory,
+                                 ignore_MPI=true)
+        restart_vpa, restart_vpa_spectral, _ =
+            load_coordinate_data(fid, "vpa"; run_directory=run_directory, ignore_MPI=true)
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta"; ignore_MPI=true)
-        restart_vr, restart_vr_spectral, _ = load_coordinate_data(fid, "vr"; ignore_MPI=true)
-        restart_vz, restart_vz_spectral, _ = load_coordinate_data(fid, "vz"; ignore_MPI=true)
+            load_coordinate_data(fid, "vzeta"; run_directory=run_directory,
+                                 ignore_MPI=true)
+        restart_vr, restart_vr_spectral, _ =
+            load_coordinate_data(fid, "vr"; run_directory=run_directory, ignore_MPI=true)
+        restart_vz, restart_vz_spectral, _ =
+            load_coordinate_data(fid, "vz"; run_directory=run_directory, ignore_MPI=true)
 
         if restart_r.nrank != r.nrank
             error("Not using parallel I/O, and distributed MPI layout has "
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 2aabcafc5..256b8b898 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -308,7 +308,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
-                                    vz)
+                                    vz; run_directory=io_input.output_dir)
 
         begin_serial_region()
         @serial_region begin
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 836144b44..8463ef61b 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -690,28 +690,31 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
 
     # initialize z grid and write grid point locations to file
     z, z_spectral = define_coordinate(z_immutable, io_immutable.parallel_io;
-                                      ignore_MPI=ignore_MPI)
+                                      run_directory=output_dir, ignore_MPI=ignore_MPI)
     # initialize r grid and write grid point locations to file
     r, r_spectral = define_coordinate(r_immutable, io_immutable.parallel_io;
-                                      ignore_MPI=ignore_MPI)
+                                      run_directory=output_dir, ignore_MPI=ignore_MPI)
     # initialize vpa grid and write grid point locations to file
     vpa, vpa_spectral = define_coordinate(vpa_immutable, io_immutable.parallel_io;
-                                          ignore_MPI=ignore_MPI)
+                                          run_directory=output_dir, ignore_MPI=ignore_MPI)
     # initialize vperp grid and write grid point locations to file
     vperp, vperp_spectral = define_coordinate(vperp_immutable, io_immutable.parallel_io;
+                                              run_directory=output_dir,
                                               ignore_MPI=ignore_MPI)
     # initialize gyrophase grid and write grid point locations to file
     gyrophase, gyrophase_spectral = define_coordinate(gyrophase_immutable,
                                                       io_immutable.parallel_io;
+                                                      run_directory=output_dir,
                                                       ignore_MPI=ignore_MPI)
     # initialize vz grid and write grid point locations to file
     vz, vz_spectral = define_coordinate(vz_immutable, io_immutable.parallel_io;
-                                        ignore_MPI=ignore_MPI)
+                                        run_directory=output_dir, ignore_MPI=ignore_MPI)
     # initialize vr grid and write grid point locations to file
     vr, vr_spectral = define_coordinate(vr_immutable, io_immutable.parallel_io;
-                                        ignore_MPI=ignore_MPI)
+                                        run_directory=output_dir, ignore_MPI=ignore_MPI)
     # initialize vr grid and write grid point locations to file
     vzeta, vzeta_spectral = define_coordinate(vzeta_immutable, io_immutable.parallel_io;
+                                              run_directory=output_dir,
                                               ignore_MPI=ignore_MPI)
 
     external_source_settings = setup_external_sources!(scan_input, r, z,

From bbfa30747de08323334562475a907026687f16d2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 15 Apr 2024 12:50:38 +0100
Subject: [PATCH 241/428] Support "Save/load FFTW wisdom from rank-0" in
 electron functions

---
 moment_kinetics/src/initial_conditions.jl | 3 ++-
 moment_kinetics/src/load_data.jl          | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index cc05e910b..e19d18742 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -490,7 +490,8 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             # Reload pdf and moments from an existing output file
             code_time, previous_runs_info, restart_time_index =
                 reload_electron_data!(pdf, moments, t_params, backup_prefix_iblock, -1,
-                                      geometry, r, z, vpa, vperp, vzeta, vr, vz)
+                                      geometry, r, z, vpa, vperp, vzeta, vr, vz;
+                                      run_directory=io_input.output_dir)
 
             # Broadcast code_time from the root process of each shared-memory block (on which it
             # might have been loaded from a restart file).
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 8cdf2407f..fb1007f52 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -986,7 +986,7 @@ end
 Reload electron pdf and moments from an existing output file.
 """
 function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, time_index,
-                               geometry, r, z, vpa, vperp, vzeta, vr, vz)
+                               geometry, r, z, vpa, vperp, vzeta, vr, vz; run_directory=nothing)
     code_time = 0.0
     previous_runs_info = nothing
     begin_serial_region()
@@ -1014,7 +1014,8 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
                 restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
                 restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
                 restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
-                                                               vzeta, vr, vz, parallel_io)
+                                                               vzeta, vr, vz, parallel_io;
+                                                               run_directory=run_directory)
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(

From b8ce72b7862544c87da27831a6c288af2e54c8fd Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 12 Apr 2024 16:36:53 +0100
Subject: [PATCH 242/428] Fix output time logic

---
 moment_kinetics/src/runge_kutta.jl | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 8e985afdd..f5d19d333 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -783,6 +783,8 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
         error("Unrecognized error_norm_method '$method'")
     end
 
+    just_completed_output_step = false
+
     # Use current_dt instead of t_params.dt[] here because we are about to write to
     # the shared-memory variable t_params.dt[] below, and we do not want to add an extra
     # _block_synchronize() call after reading it here.
@@ -846,6 +848,8 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
                 if t_params.dt[] > CFL_limit
                     t_params.dt[] = CFL_limit
                 end
+
+                just_completed_output_step = true
             else
                 # Adjust timestep according to Fehlberg's suggestion
                 # (https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta%E2%80%93Fehlberg_method).
@@ -912,9 +916,12 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
     end
 
     @serial_region begin
-        if t + t_params.dt[] >= t_params.next_output_time[]
+        current_time = t + t_params.previous_dt[]
+        if (!just_completed_output_step
+            && (current_time + t_params.dt[] >= t_params.next_output_time[]))
+
             t_params.dt_before_output[] = t_params.dt[]
-            t_params.dt[] = t_params.next_output_time[] - t
+            t_params.dt[] = t_params.next_output_time[] - current_time
             t_params.step_to_output[] = true
         end
     end

From 560b4c6b3259f7823630715c4e5bb8f945914d96 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 16 Apr 2024 14:09:36 +0100
Subject: [PATCH 243/428] Speed up debug I/O

---
 .../src/electron_kinetic_equation.jl            |  9 ++-------
 moment_kinetics/src/file_io.jl                  | 17 +++++++++++++++++
 moment_kinetics/src/initial_conditions.jl       |  8 ++++++++
 moment_kinetics/src/time_advance.jl             | 10 +++++++++-
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 6057648fb..0bb718299 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -21,7 +21,7 @@ using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
 using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
 using ..external_sources: external_electron_source!
-using ..file_io: setup_electron_io, write_electron_state, finish_electron_io
+using ..file_io: get_electron_io_info, write_electron_state, finish_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
 using ..runge_kutta: rk_update_variable!, rk_error_variable!, local_error_norm,
@@ -185,12 +185,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     if io_electron === nothing && t_params.debug_io !== nothing
         # Overwrite the debug output file with the output from this call to
         # update_electron_pdf_with_time_advance!().
-        io_electron = setup_electron_io(t_params.debug_io[1], vpa, vperp, z, r,
-                                        composition, collisions,
-                                        moments.evolve_density, moments.evolve_upar,
-                                        moments.evolve_ppar, external_source_settings,
-                                        t_params.debug_io[2], -1, nothing,
-                                        "electron_debug")
+        io_electron = get_electron_io_info(t_params.debug_io[1], "electron_debug")
         do_debug_io = true
         debug_io_nwrite = t_params.debug_io[3]
     else
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 1060dc34a..bc4ea8834 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -398,6 +398,23 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
     return nothing
 end
 
+"""
+Get the `file_info` for an existing electron I/O file
+"""
+function get_electron_io_info(io_input, prefix_label)
+    out_prefix = joinpath(io_input.output_dir, io_input.run_name)
+    electrons_prefix = string(out_prefix, ".$prefix_label")
+    if io_input.binary_format == hdf5
+        filename = string(electrons_prefix, ".h5")
+    elseif io_input.binary_format == netcdf
+        filename = string(electrons_prefix, ".cdf")
+    else
+        error("Unrecognized binary_format=$(io_input.binary_format)")
+    end
+
+    return (filename, io_input.parallel_io, comm_inter_block[])
+end
+
 """
 Reopen an existing initial electron output file to append more data
 """
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index e19d18742..965735134 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -546,6 +546,14 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         if global_rank[] == 0
             println("Initializing electrons - evolving both pdf_electron and electron_ppar")
         end
+        if t_params.debug_io !== nothing
+            io_electron = setup_electron_io(t_params.debug_io[1], vpa, vperp, z, r,
+                                            composition, collisions,
+                                            moments.evolve_density, moments.evolve_upar,
+                                            moments.evolve_ppar, external_source_settings,
+                                            t_params.debug_io[2], -1, nothing,
+                                            "electron_debug")
+        end
         electron_pseudotime, n_debug_outputs =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 3c19c8f46..1bbd749e2 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -14,7 +14,8 @@ using ..array_allocation: allocate_float, allocate_shared_float, allocate_shared
 using ..communication
 using ..communication: _block_synchronize
 using ..debugging
-using ..file_io: write_data_to_ascii, write_all_moments_data_to_binary, write_all_dfns_data_to_binary, debug_dump
+using ..file_io: write_data_to_ascii, write_all_moments_data_to_binary,
+                 write_all_dfns_data_to_binary, debug_dump, setup_electron_io
 using ..initial_conditions: initialize_electrons!
 using ..looping
 using ..moment_kinetics_structs: scratch_pdf
@@ -492,6 +493,13 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                               external_source_settings, scratch_dummy, scratch, t_params,
                               t_input, num_diss_params, advection_structs, io_input,
                               input_dict; restart_from_Boltzmann_electrons=restarting)
+    elseif restarting && t_params.electron.debug_io !== nothing
+        io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z, r,
+                                        composition, collisions, moments.evolve_density,
+                                        moments.evolve_upar, moments.evolve_ppar,
+                                        external_source_settings,
+                                        t_params.electron.debug_io[2], -1, nothing,
+                                        "electron_debug")
     end
 
     # update the derivatives of the electron moments as these may be needed when

From 0936b5a23ff67197f294c69168fdc2b577df13fc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 Apr 2024 20:56:53 +0100
Subject: [PATCH 244/428] Fix duplicated calculations in steady state residual
 check

---
 moment_kinetics/src/analysis.jl | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/analysis.jl b/moment_kinetics/src/analysis.jl
index 7e3bd1600..fc342e050 100644
--- a/moment_kinetics/src/analysis.jl
+++ b/moment_kinetics/src/analysis.jl
@@ -647,17 +647,22 @@ function steady_state_square_residuals(variable, variable_at_previous_time, dt;
 
             if only_max_abs
                 absolute_residual =
-                    _steady_state_absolute_residual(variable, variable_at_previous_time,
+                    _steady_state_absolute_residual(this_slice, this_slice_previous_time,
                                                     reshaped_dt)
                 # Need to wrap the maximum(...) in a call to vec(...) so that we return a
                 # Vector, not an N-dimensional array where the first (N-1) dimensions all
                 # have size 1.
-                local_max_absolute = max.(local_max_absolute,
-                                          vec(maximum(absolute_residual,
-                                                      dims=tuple((1:t_dim-1)...))))
+                this_dims = tuple((1:t_dim-3)...)
+                if this_dims === ()
+                    local_max_absolute = max.(local_max_absolute, [absolute_residual])
+                else
+                    local_max_absolute = max.(local_max_absolute,
+                                              vec(maximum(absolute_residual,
+                                                          dims=this_dims)))
+                end
             else
                 absolute_square_residual, relative_square_residual =
-                    _steady_state_square_residual(variable, variable_at_previous_time,
+                    _steady_state_square_residual(this_slice, this_slice_previous_time,
                                                   reshaped_dt, epsilon, variable_max)
                 # Need to wrap the sum(...) or maximum(...) in a call to vec(...) so that
                 # we return a Vector, not an N-dimensional array where the first (N-1)

From bac49f2cec14a4958c4e52b6e3b04ff4bacb5482 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 17 Apr 2024 11:58:17 +0100
Subject: [PATCH 245/428] Enforce sheath-edge bc on electron qpar when using
 Braginskii electrons

This should allow `braginskii_fluid`-electron simulations to work again.
Previous boundary conditions (setting T_e equal to ion temperature or
wall temperature) had been removed (this commit replaces them, using a
bc on qpar instead), which would have broken `braginskii_fluid`
simulations.
---
 .../src/electron_fluid_equations.jl           | 44 +++++++++++++++++++
 moment_kinetics/src/time_advance.jl           |  6 ++-
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 33300b6e9..f1a7a5e2b 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -386,4 +386,48 @@ function update_electron_vth_temperature!(moments, ppar, dens)
     return nothing
 end
 
+"""
+    electron_fluid_qpar_boundary_condition!(electron_moments, z)
+
+Impose fluid approximation to electron sheath boundary condition on the parallel heat
+flux. See Stangeby textbook, equations (2.89) and (2.90).
+"""
+function electron_fluid_qpar_boundary_condition!(electron_moments, z)
+    begin_r_region()
+
+    if z.irank == 0 && (z.irank == z.nrank - 1)
+        z_indices = (1, z.n)
+    elseif z.irank == 0
+        z_indices = (1,)
+    elseif z.irank == z.nrank - 1
+        z_indices = (z.n,)
+    else
+        return nothing
+    end
+
+    @loop_r ir begin
+        for iz ∈ z_indices
+            ppar = electron_moments.ppar[iz,ir]
+            upar = electron_moments.upar[iz,ir]
+            dens = electron_moments.dens[iz,ir]
+            particle_flux = dens * upar
+            T_e = electron_moments.temp[iz,ir]
+
+            # Stangeby (2.90)
+            gamma_e = 5.5
+
+            # Stangeby (2.89)
+            total_heat_flux = gamma_e * T_e * particle_flux
+
+            # E.g. Helander&Sigmar (2.14), neglecting electron viscosity and kinetic
+            # energy fluxes due to small mass ratio
+            conductive_heat_flux = total_heat_flux - 2.5 * ppar * upar
+
+            electron_moments.qpar[iz,ir] = conductive_heat_flux
+        end
+    end
+
+    return nothing
+end
+
 end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 1bbd749e2..a535430ab 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -69,7 +69,7 @@ using ..utils: to_minutes, get_minimum_CFL_z, get_minimum_CFL_vpa,
                get_minimum_CFL_neutral_z, get_minimum_CFL_neutral_vz
 using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
-using ..electron_fluid_equations: calculate_electron_qpar!
+using ..electron_fluid_equations: calculate_electron_qpar!, electron_fluid_qpar_boundary_condition!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..electron_fluid_equations: electron_energy_equation!, update_electron_vth_temperature!
 using ..input_structs: braginskii_fluid
@@ -1763,7 +1763,9 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         calculate_electron_qpar!(moments.electron, new_scratch.pdf_electron,
             new_scratch.electron_ppar, new_scratch.electron_upar, new_scratch.upar,
             collisions.nu_ei, composition.me_over_mi, composition.electron_physics, vpa)
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics == braginskii_fluid
+            electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        elseif composition.electron_physics == kinetic_electrons
             max_electron_pdf_iterations = 100000
 
             # Copy ion and electron moments from `scratch` into `moments` to be used in

From 5baf19a1b115040df6d0cd33c80ec5b284a97e17 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 Apr 2024 14:40:55 +0100
Subject: [PATCH 246/428] Try-catch in calculate_steady_state_residual()
 post-proc function

---
 .../src/makie_post_processing.jl              | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 4b57b0116..ebc96994d 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -3662,24 +3662,28 @@ function calculate_steady_state_residual end
 function calculate_steady_state_residual(run_info::Tuple, variable_name; is=1,
                                          data=nothing, plot_prefix=nothing,
                                          fig_axes=nothing)
-    n_runs = length(run_info)
-    if data === nothing
-        data = Tuple(nothing for _ ∈ 1:n_runs)
-    end
-    if fig_axes === nothing
-        fig_axes = _get_steady_state_residual_fig_axes(length(run_info))
-    end
+    try
+        n_runs = length(run_info)
+        if data === nothing
+            data = Tuple(nothing for _ ∈ 1:n_runs)
+        end
+        if fig_axes === nothing
+            fig_axes = _get_steady_state_residual_fig_axes(length(run_info))
+        end
 
-    for (i, (ri, d)) ∈ enumerate(zip(run_info, data))
-        calculate_steady_state_residual(ri, variable_name; is=is, data=d,
-                                        fig_axes=fig_axes, i_run=i)
-    end
+        for (i, (ri, d)) ∈ enumerate(zip(run_info, data))
+            calculate_steady_state_residual(ri, variable_name; is=is, data=d,
+                                            fig_axes=fig_axes, i_run=i)
+        end
 
-    if plot_prefix !== nothing
-        _save_residual_plots(fig_axes, plot_prefix)
-    end
+        if plot_prefix !== nothing
+            _save_residual_plots(fig_axes, plot_prefix)
+        end
 
-    return fig_axes
+        return fig_axes
+    catch e
+        println("Error in calculate_steady_state_residual(). Error was ", e)
+    end
 end
 
 function calculate_steady_state_residual(run_info, variable_name; is=1, data=nothing,

From 8b474e1ee02fee1037873c4f8a9c1cea3d7cf576 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 Apr 2024 14:43:39 +0100
Subject: [PATCH 247/428] Fix restarting using fixed (non-adaptive) timestep

`dt` was being reloaded from the restart file. This is correct to give a
good initial guess for `dt` when using an adaptive timestep, but when
using a fixed timestep we should always use the value from the input
file.
---
 moment_kinetics/src/time_advance.jl | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index a535430ab..3c12497cd 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -276,6 +276,17 @@ the returned `time_info`.
 """
 function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_reload,
                          manufactured_solns_input, io_input, input_dict; electron=nothing)
+    rk_coefs, n_rk_stages, rk_order, adaptive, low_storage, CFL_prefactor =
+        setup_runge_kutta_coefficients!(t_input.type,
+                                        t_input.CFL_prefactor,
+                                        t_input.split_operators)
+
+    if !adaptive
+        # No adaptive timestep, want to use the value from the input file even when we are
+        # restarting
+        dt_reload = nothing
+    end
+
     dt_shared = allocate_shared_float(1)
     previous_dt_shared = allocate_shared_float(1)
     next_output_time = allocate_shared_float(1)
@@ -312,10 +323,7 @@ function setup_time_info(t_input, code_time, dt_reload, dt_before_last_fail_relo
     if dfns_output_times[end] < end_time - epsilon
         push!(dfns_output_times, end_time)
     end
-    rk_coefs, n_rk_stages, rk_order, adaptive, low_storage, CFL_prefactor =
-        setup_runge_kutta_coefficients!(t_input.type,
-                                        t_input.CFL_prefactor,
-                                        t_input.split_operators)
+
     if t_input.high_precision_error_sum
         error_sum_zero = Float128(0.0)
     else

From e93276e6e27075c40a60917515fac60e53750956 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 Apr 2024 14:48:39 +0100
Subject: [PATCH 248/428] Fix synchronization when not using MPI in mk_input()

All synchronization calls should be within `if !ignore_MPI` blocks, to
avoid breaking post-processing and tests.
---
 moment_kinetics/src/moment_kinetics_input.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 8463ef61b..eede8eb66 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -558,10 +558,12 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     #nrank_z = 0
 
     # Create output_dir if it does not exist.
-    if global_rank[] == 0
-        mkpath(output_dir)
+    if !ignore_MPI
+        if global_rank[] == 0
+            mkpath(output_dir)
+        end
+        _block_synchronize()
     end
-    _block_synchronize()
 
     # replace mutable structures with immutable ones to optimize performance
     # and avoid possible misunderstandings	

From ec176c1416aa4c75d89b2d7b202d79262a754ee2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 21 Apr 2024 14:50:02 +0100
Subject: [PATCH 249/428] Fix 'braginskii_fluid' electron option

Boundary conditions had been deleted while updating kinetic electrons.

The boundary condition is now imposed on `moments.electron.qpar`, using
a sheath transmission coefficient from Stangeby's textbook. Fixed
temperature boundary conditions (as were used previously) could also be
useful - these should perhaps be re-implemented in future.

Also adds several missing factors of me_over_mi.
---
 .../src/electron_fluid_equations.jl           |  7 +-
 moment_kinetics/src/initial_conditions.jl     | 80 ++++++++++++++++++-
 moment_kinetics/src/moment_kinetics.jl        |  3 +-
 moment_kinetics/src/time_advance.jl           | 15 ++--
 4 files changed, 94 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index f1a7a5e2b..2359af6e1 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -372,14 +372,15 @@ end
 #    end
 #end
 
-function update_electron_vth_temperature!(moments, ppar, dens)
+function update_electron_vth_temperature!(moments, ppar, dens, composition)
     begin_r_z_region()
 
     temp = moments.electron.temp
     vth = moments.electron.vth
     @loop_r_z ir iz begin
-        temp[iz,ir] = 2 * ppar[iz,ir] / dens[iz,ir]
-        vth[iz,ir] = sqrt(temp[iz,ir])
+        p = max(ppar[iz,ir], 0.0)
+        temp[iz,ir] = 2 * p[iz,ir] / dens[iz,ir]
+        vth[iz,ir] = sqrt(temp[iz,ir] / composition.me_over_mi)
     end
     moments.electron.temp_updated[] = true
 
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 965735134..ae1e05864 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -35,10 +35,10 @@ using ..velocity_moments: update_neutral_uz!, update_neutral_ur!, update_neutral
 using ..velocity_moments: update_ppar!, update_upar!, update_density!, update_pperp!, update_vth!, reset_moments_status!
 using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
-using ..electron_fluid_equations: calculate_electron_qpar!
+using ..electron_fluid_equations: calculate_electron_qpar!, electron_fluid_qpar_boundary_condition!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..electron_kinetic_equation: update_electron_pdf!, enforce_boundary_condition_on_electron_pdf!
-using ..input_structs: boltzmann_electron_response_with_simple_sheath, kinetic_electrons
+using ..input_structs: boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons
 using ..derivatives: derivative_z!
 using ..utils: get_default_restart_filename, get_prefix_iblock_and_move_existing_file
 
@@ -328,10 +328,85 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
+    if composition.electron_physics == braginskii_fluid
+        electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        if restart_from_Boltzmann_electrons
+            # Restarting from Boltzmann. If we use an exactly constant T_e profile,
+            # qpar for the electrons will be non-zero only at the boundary points,
+            # which will crash the code unless the timestep is insanely small. Give
+            # T_e a cubic profile with gradients at the boundaries that match the
+            # boundary qpar. The boundary values are both the constant 'Boltzmann
+            # electron' temperature
+            begin_serial_region()
+            @serial_region begin
+                # Ignore the 0.71*p_e*(u_i-u_e) term here as it would vanish when u_i=u_e
+                # and this is only a rough initial condition anyway
+                #
+                # q at the boundaries tells us dTe/dz for Braginskii electrons
+                if z.irank == 0
+                    dTe_dz_lower = @. -moments.electron.qpar[1,:] * 2.0 / 3.16 /
+                                       moments.electron.ppar[1,:] *
+                                       composition.me_over_mi * collisions.nu_ei
+                else
+                    dTe_dz_lower = nothing
+                end
+                dTe_dz_lower = MPI.bcast(dTe_dz_lower, z.comm; root=0)
+
+                if z.irank == z.nrank - 1
+                    dTe_dz_upper = @. -moments.electron.qpar[end,:] * 2.0 / 3.16 /
+                                       moments.electron.ppar[end,:] *
+                                       composition.me_over_mi * collisions.nu_ei
+                else
+                    dTe_dz_upper = nothing
+                end
+                dTe_dz_upper = MPI.bcast(dTe_dz_upper, z.comm; root=(z.nrank - 1))
+
+                # The temperature should already be equal to the 'Boltzmann electron'
+                # Te, so we just need to add a cubic that vanishes at ±Lz/2
+                # δT = A + B*z + C*z^2 + D*z^3
+                # ⇒ A - B*Lz/2 + C*Lz^2/4 - D*Lz^3/8 = 0
+                #   A + B*Lz/2 + C*Lz^2/4 + D*Lz^3/8 = 0
+                #   B - C*Lz + 3*D*Lz^2/4 = dTe/dz_lower
+                #   B + C*Lz + 3*D*Lz^2/4 = dTe/dz_upper
+                #
+                # Adding the first pair together, and subtracting the second pair:
+                #   A + C*Lz^2/4 = 0
+                #   2*C*Lz = dT/dz_upper - dT/dz_lower
+                #
+                # Subtracting the first pair and adding the second instead:
+                #   B*Lz/2 + D*Lz^3/8 = 0  ⇒  D*Lz^2/2 = -2*B
+                #   2*B + 3*D*Lz^2/2 = dTe/dz_upper - dTe/dz_lower
+                #
+                #   2*B - 3*2*B = -4*B = dTe/dz_upper + dTe/dz_lower
+                Lz = z.L
+                zg = z.grid
+                C = @. (dTe_dz_upper - dTe_dz_lower) / 2.0 / Lz
+                A = @. -C * Lz^2 / 4
+                B = @. -(dTe_dz_lower + dTe_dz_upper) / 4.0
+                D = @. -4.0 * B / Lz^2
+                @loop_r ir begin
+                    @. moments.electron.temp[:,ir] += A[ir] + B[ir]*zg + C[ir]*zg^2 +
+                                                      D[ir]*zg^3
+                end
+
+                @. moments.electron.vth = sqrt(moments.electron.temp /
+                                               composition.me_over_mi)
+                @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
+                @views derivative_z!(moments.electron.dT_dz, moments.electron.temp,
+                                     scratch_dummy.buffer_rs_1[:,1],
+                                     scratch_dummy.buffer_rs_2[:,1],
+                                     scratch_dummy.buffer_rs_3[:,1],
+                                     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
+            end
+        end
+    end
     moments.electron.qpar_updated[] = false
     calculate_electron_qpar!(moments.electron, pdf.electron, moments.electron.ppar,
         moments.electron.upar, moments.ion.upar, collisions.nu_ei, composition.me_over_mi,
         composition.electron_physics, vpa)
+    if composition.electron_physics == braginskii_fluid
+        electron_fluid_qpar_boundary_condition!(moments.electron, z)
+    end
     # calculate the zed derivative of the initial electron parallel heat flux
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
         scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
@@ -353,6 +428,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
             scratch[1].electron_upar .= moments.electron.upar
             scratch[1].electron_ppar .= moments.electron.ppar
             scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
+            scratch[1].electron_temp .= moments.electron.temp
         end
         # get the initial electrostatic potential and parallel electric field
         update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 256b8b898..9a6ac26be 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -312,7 +312,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
 
         begin_serial_region()
         @serial_region begin
-            @. moments.electron.temp = moments.electron.vth^2
+            @. moments.electron.temp = composition.me_over_mi * moments.electron.vth^2
+            @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
         end
         if composition.electron_physics == kinetic_electrons
             begin_r_z_vperp_vpa_region()
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 3c12497cd..0ae6c73ec 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -493,8 +493,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
 
     # Now that `t_params` and `scratch` have been created, initialize electrons if
     # necessary
-    if !restarting || (composition.electron_physics == kinetic_electrons &&
-                       !restart_had_kinetic_electrons)
+    if !restarting ||
+        (composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons) &&
+         !restart_had_kinetic_electrons)
         initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions,
@@ -718,7 +719,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
         # compute the updated electron temperature
         # NB: not currently necessary, as initial vth is not directly dependent on ion quantities
         @serial_region begin
-            @. moments.electron.temp = moments.electron.vth^2
+            @. moments.electron.temp = composition.me_over_mi * moments.electron.vth^2
         end
         # as the electron temperature has now been updated, set the appropriate flag
         moments.electron.temp_updated[] = true
@@ -738,6 +739,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
         calculate_electron_qpar!(moments.electron, pdf.electron, moments.electron.ppar,
             moments.electron.upar, moments.ion.upar, collisions.nu_ei,
             composition.me_over_mi, composition.electron_physics, vpa)
+        if composition.electron_physics == braginskii_fluid
+            electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        end
         # calculate the electron-ion parallel friction force
         calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
             moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
@@ -1755,14 +1759,15 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
         else
             begin_r_z_region()
             @loop_r_z ir iz begin
-                new_scratch.electron_ppar[iz,ir] = 0.5 * new_scratch.electron_density[iz,ir] *
+                new_scratch.electron_ppar[iz,ir] = 0.5 * composition.me_over_mi *
+                                                   new_scratch.electron_density[iz,ir] *
                                                    moments.electron.vth[iz,ir]^2
             end
         end
         # regardless of electron model, electron ppar is now updated
         moments.electron.ppar_updated[] = true
         update_electron_vth_temperature!(moments, new_scratch.electron_ppar,
-                                         new_scratch.electron_density)
+                                         new_scratch.electron_density, composition)
         # calculate the corresponding zed derivatives of the moments
         calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z, z_spectral,
                                                num_diss_params.electron.moment_dissipation_coefficient, 

From 1473a7ae2046ef2de1fb0e031acc5e1492db7040 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Apr 2024 19:22:01 +0100
Subject: [PATCH 250/428] Utility function to get an enum value from a String

---
 moment_kinetics/src/utils.jl | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl
index 052f71687..1007f7425 100644
--- a/moment_kinetics/src/utils.jl
+++ b/moment_kinetics/src/utils.jl
@@ -277,6 +277,23 @@ function get_prefix_iblock_and_move_existing_file(restart_filename, output_dir)
     return backup_prefix_iblock
 end
 
+"""
+    enum_from_string(enum_type, name)
+
+Get an the value of `enum_type`, whose name is given by the String (or Symbol) `name`.
+
+Returns `nothing` if the name is not found.
+"""
+function enum_from_string(enum_type, name)
+    name = Symbol(name)
+    for e ∈ instances(enum_type)
+        if name == Symbol(e)
+            return e
+        end
+    end
+    return nothing
+end
+
 # Utility functions for timestepping
 
 """

From 8d4610c068a3a573425d3f0d421824214825bb90 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Apr 2024 19:43:59 +0100
Subject: [PATCH 251/428] Tidy up electron_physics handling when restarting

Get the `electron_physics` setting that was used for the run being
restarted from (as the `restart_electron_physics` variable) and use that
to decide how to initialize the electrons. More general than just
`restart_from_Boltzmann_electrons` and `restart_had_kinetic_electrons`,
so allows slightly tidier code.
---
 moment_kinetics/src/initial_conditions.jl | 55 ++++++++++++-----------
 moment_kinetics/src/load_data.jl          | 22 +++++----
 moment_kinetics/src/moment_kinetics.jl    |  6 +--
 moment_kinetics/src/time_advance.jl       | 15 ++++---
 4 files changed, 54 insertions(+), 44 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index ae1e05864..b1a9f157c 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -257,7 +257,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                                vperp_spectral, vpa_spectral, collisions,
                                external_source_settings, scratch_dummy, scratch, t_params,
                                t_input, num_diss_params, advection_structs, io_input,
-                               input_dict; restart_from_Boltzmann_electrons=false)
+                               input_dict; restart_electron_physics)
     
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
@@ -268,18 +268,9 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         moments.ion.upar, moments.ion.dens, composition.electron_physics, r, z)
     # different choices for initialization of electron temperature/pressure/vth depending on whether
     # we are restarting from a previous simulation with Boltzmann electrons or not
-    if restart_from_Boltzmann_electrons
-        begin_serial_region()
-        @serial_region begin
-            # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
-            # that the electron parallel temperature is constant along the field line and equal to T_e
-            moments.electron.temp .= composition.T_e
-            # the thermal speed is related to the temperature by vth_e / v_ref = sqrt((T_e/T_ref) / (m_e/m_ref))
-            moments.electron.vth .= sqrt(composition.T_e / composition.me_over_mi)
-            # ppar = 0.5 * n * T, so we can calculate the parallel pressure from the density and T_e
-            moments.electron.ppar .= 0.5 * moments.electron.dens * composition.T_e
-        end
-    else
+    if restart_electron_physics === nothing
+        # Not restarting, so create initial profiles
+
         # initialise the electron thermal speed profile
         init_electron_vth!(moments.electron.vth, moments.ion.vth, composition.T_e, composition.me_over_mi, z.grid)
         begin_r_z_region()
@@ -291,7 +282,21 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         @loop_r_z ir iz begin
             moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
         end
-    end
+    elseif restart_electron_physics ∉ (braginskii_fluid, kinetic_electrons)
+        # Restarting from Boltzmann electron simulation, so start with constant
+        # electron temperature
+        begin_serial_region()
+        @serial_region begin
+            # if restarting from a simulations where Boltzmann electrons were used, then the assumption is
+            # that the electron parallel temperature is constant along the field line and equal to T_e
+            moments.electron.temp .= composition.T_e
+            # the thermal speed is related to the temperature by vth_e / v_ref = sqrt((T_e/T_ref) / (m_e/m_ref))
+            moments.electron.vth .= sqrt(composition.T_e / composition.me_over_mi)
+            # ppar = 0.5 * n * T, so we can calculate the parallel pressure from the density and T_e
+            moments.electron.ppar .= 0.5 * moments.electron.dens * composition.T_e
+        end
+    end # else, we are restarting from `braginskii_fluid` or `kinetic_electrons`, so keep the reloaded electron pressure/temperature profiles.
+
     # the electron temperature has now been updated
     moments.electron.temp_updated[] = true
     # the electron parallel pressure now been updated
@@ -330,7 +335,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
     if composition.electron_physics == braginskii_fluid
         electron_fluid_qpar_boundary_condition!(moments.electron, z)
-        if restart_from_Boltzmann_electrons
+        if restart_electron_physics ∉ (nothing, braginskii_fluid, kinetic_electrons)
             # Restarting from Boltzmann. If we use an exactly constant T_e profile,
             # qpar for the electrons will be non-zero only at the boundary points,
             # which will crash the code unless the timestep is insanely small. Give
@@ -421,18 +426,16 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # arrays need to exist and be otherwise initialised in order to compute the initial
     # electron pdf. The electron arrays will be updated as necessary by
     # initialize_electron_pdf!().
-    if !restart_from_Boltzmann_electrons
-        begin_serial_region()
-        @serial_region begin
-            scratch[1].electron_density .= moments.electron.dens
-            scratch[1].electron_upar .= moments.electron.upar
-            scratch[1].electron_ppar .= moments.electron.ppar
-            scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
-            scratch[1].electron_temp .= moments.electron.temp
-        end
-        # get the initial electrostatic potential and parallel electric field
-        update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
+    begin_serial_region()
+    @serial_region begin
+        scratch[1].electron_density .= moments.electron.dens
+        scratch[1].electron_upar .= moments.electron.upar
+        scratch[1].electron_ppar .= moments.electron.ppar
+        scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
+        scratch[1].electron_temp .= moments.electron.temp
     end
+    # get the initial electrostatic potential and parallel electric field
+    update_phi!(fields, scratch[1], z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy)
 
     # initialize the electron pdf that satisfies the electron kinetic equation
     initialize_electron_pdf!(scratch, pdf, moments, fields.phi, r, z, vpa, vperp,
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index fb1007f52..afc8137bc 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -20,7 +20,7 @@ using ..coordinates: coordinate, define_coordinate
 using ..electron_vpa_advection: update_electron_speed_vpa!
 using ..electron_z_advection: update_electron_speed_z!
 using ..file_io: check_io_implementation, get_group, get_subgroup_keys, get_variable_keys
-using ..input_structs: advection_input, grid_input, hdf5, netcdf
+using ..input_structs
 using ..interpolation: interpolate_to_grid_1d!
 using ..krook_collisions
 using ..looping
@@ -29,7 +29,7 @@ using ..neutral_vz_advection: update_speed_neutral_vz!
 using ..neutral_z_advection: update_speed_neutral_z!
 using ..type_definitions: mk_float, mk_int
 using ..utils: get_CFL!, get_minimum_CFL_z, get_minimum_CFL_vpa, get_minimum_CFL_neutral_z,
-               get_minimum_CFL_neutral_vz
+               get_minimum_CFL_neutral_vz, enum_from_string
 using ..vpa_advection: update_speed_vpa!
 using ..z_advection: update_speed_z!
 
@@ -582,7 +582,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
     electron_dt = nothing
     electron_dt_before_last_fail = nothing
     previous_runs_info = nothing
-    restart_had_kinetic_electrons = false
+    restart_electron_physics = nothing
     begin_serial_region()
     @serial_region begin
         fid = open_readonly_output_file(restart_prefix_iblock[1], "dfns";
@@ -597,6 +597,8 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
             restart_evolve_density, restart_evolve_upar, restart_evolve_ppar =
                 load_mk_options(fid)
 
+            restart_input = load_input(fid)
+
             previous_runs_info = load_run_info_history(fid)
 
             restart_n_ion_species, restart_n_neutral_species = load_species_data(fid)
@@ -823,8 +825,13 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                 restart_electron_evolve_ppar = true, true, true
             electron_evolve_density, electron_evolve_upar, electron_evolve_ppar =
                 true, true, true
-            restart_had_kinetic_electrons = ("f_electron" ∈ keys(dynamic))
-            if pdf.electron !== nothing && restart_had_kinetic_electrons
+            if "electron_physics" ∈ keys(restart_input)
+                restart_electron_physics = enum_from_string(electron_physics_type,
+                                                            restart_input["electron_physics"])
+            else
+                restart_electron_physics = boltzmann_electron_response
+            end
+            if pdf.electron !== nothing && restart_electron_physics == kinetic_electrons
                 pdf.electron.norm .=
                     reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
                                         r_range, z_range, vperp_range, vpa_range,
@@ -975,11 +982,10 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
         end
     end
 
-    restart_had_kinetic_electrons = MPI.Bcast(restart_had_kinetic_electrons, 0,
-                                              comm_block[])
+    restart_electron_physics = MPI.bcast(restart_electron_physics, 0, comm_block[])
 
     return code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
-           previous_runs_info, time_index, restart_had_kinetic_electrons
+           previous_runs_info, time_index, restart_electron_physics
 end
 
 """
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 9a6ac26be..be55a1119 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -289,7 +289,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         electron_dt = nothing
         electron_dt_before_last_fail = nothing
         previous_runs_info = nothing
-        restart_had_kinetic_electrons = false
+        restart_electron_physics = nothing
     else
         restarting = true
 
@@ -304,7 +304,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
 
         # Reload pdf and moments from an existing output file
         code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
-        previous_runs_info, restart_time_index, restart_had_kinetic_electrons =
+        previous_runs_info, restart_time_index, restart_electron_physics =
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
@@ -346,7 +346,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             dt_before_last_fail, electron_dt, electron_dt_before_last_fail, collisions,
             species, geometry, boundary_distributions, external_source_settings,
             num_diss_params, manufactured_solns_input, advection_structs, scratch_dummy,
-            io_input, restarting, restart_had_kinetic_electrons, input_dict)
+            io_input, restarting, restart_electron_physics, input_dict)
 
     # This is the closest we can get to the end time of the setup before writing it to the
     # output file
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 0ae6c73ec..2a1b4fd28 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -379,8 +379,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
                              collisions, species, geometry, boundary_distributions,
                              external_source_settings, num_diss_params,
                              manufactured_solns_input, advection_structs, scratch_dummy,
-                             io_input, restarting, restart_had_kinetic_electrons,
-                             input_dict)
+                             io_input, restarting, restart_electron_physics, input_dict)
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
     n_neutral_species = composition.n_neutral_species
@@ -493,16 +492,18 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
 
     # Now that `t_params` and `scratch` have been created, initialize electrons if
     # necessary
-    if !restarting ||
-        (composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons) &&
-         !restart_had_kinetic_electrons)
+    if composition.electron_physics != restart_electron_physics
         initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions,
                               external_source_settings, scratch_dummy, scratch, t_params,
                               t_input, num_diss_params, advection_structs, io_input,
-                              input_dict; restart_from_Boltzmann_electrons=restarting)
-    elseif restarting && t_params.electron.debug_io !== nothing
+                              input_dict;
+                              restart_electron_physics=restart_electron_physics)
+    elseif restarting && composition.electron_physics == kinetic_electrons &&
+           t_params.electron.debug_io !== nothing
+        # Create *.electron_debug.h5 file so that it can be re-opened in
+        # update_electron_pdf!().
         io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z, r,
                                         composition, collisions, moments.evolve_density,
                                         moments.evolve_upar, moments.evolve_ppar,

From 1781f1a6c23006485f1a580a56715a641cd5c92b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Apr 2024 20:19:21 +0100
Subject: [PATCH 252/428] Reset dt when restarting braginskii_fluid from other
 electron physics

`braginskii_fluid` has a very restrictive CFL condition due to the
diffusive heat conduction, so when using adaptive timestepping, reset to
the (presumably small) initial timestep set by the input file, rather
than trying to continue with the adaptive timestep from the run being
restarted from.
---
 moment_kinetics/src/time_advance.jl | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 2a1b4fd28..48545ff00 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -387,6 +387,17 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, vz_sp
     electron_mom_diss_coeff = num_diss_params.electron.moment_dissipation_coefficient
     neutral_mom_diss_coeff = num_diss_params.neutral.moment_dissipation_coefficient
 
+    if composition.electron_physics == braginskii_fluid &&
+            restart_electron_physics !== braginskii_fluid
+
+        # When restarting braginskii_fluid from a different electron physics type, and
+        # using an adaptive timestep, do not want to keep the `dt` from the previous
+        # simulation, as the diffusive behaviour of braginskii fluid means the CFL
+        # condition is likely to be very restrictive, so we probably need to reset to a
+        # small initial timestep.
+        dt_reload = nothing
+    end
+
     if composition.electron_physics == kinetic_electrons
         electron_t_params = setup_time_info(t_input.electron_t_input, 0.0,
                                             electron_dt_reload,

From a0b941ea7533c479948e78335b8e7d481190b1a6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Apr 2024 10:23:24 +0100
Subject: [PATCH 253/428] Remove unnecessary arguments in load_data

As `ignore_MPI=true` now by default, can remove `ignore_MPI` and
`run_directory` arguments in many places.
---
 moment_kinetics/src/initial_conditions.jl |  3 +-
 moment_kinetics/src/load_data.jl          | 74 ++++++++++-------------
 moment_kinetics/src/moment_kinetics.jl    |  2 +-
 3 files changed, 33 insertions(+), 46 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index b1a9f157c..659c5e9ab 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -569,8 +569,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             # Reload pdf and moments from an existing output file
             code_time, previous_runs_info, restart_time_index =
                 reload_electron_data!(pdf, moments, t_params, backup_prefix_iblock, -1,
-                                      geometry, r, z, vpa, vperp, vzeta, vr, vz;
-                                      run_directory=io_input.output_dir)
+                                      geometry, r, z, vpa, vperp, vzeta, vr, vz)
 
             # Broadcast code_time from the root process of each shared-memory block (on which it
             # might have been loaded from a restart file).
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 6ca8bab2b..5be2a7550 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -199,7 +199,8 @@ function load_input(fid)
 end
 
 """
-    load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing)
+    load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing,
+                         ignore_MPI=true)
 
 Load data for the coordinate `name` from a file-handle `fid`.
 
@@ -213,12 +214,12 @@ If `irank` and `nrank` are passed, then the `coord` and `spectral` objects retur
 be set up for the parallelisation specified by `irank` and `nrank`, rather than the one
 implied by the output file.
 
-If `ignore_MPI=true` is passed, the returned coordinates will be created without shared
-memory scratch arrays (`ignore_MPI=true` will be passed through to
+Unless `ignore_MPI=false` is passed, the returned coordinates will be created without
+shared memory scratch arrays (`ignore_MPI=true` will be passed through to
 [`define_coordinate`](@ref)).
 """
 function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=nothing,
-                              run_directory=nothing, ignore_MPI=true)
+                              ignore_MPI=true)
     if printout
         println("Loading $name coordinate data...")
     end
@@ -574,7 +575,7 @@ Reload pdf and moments from an existing output file.
 """
 function reload_evolving_fields!(pdf, moments, boundary_distributions,
                                  restart_prefix_iblock, time_index, composition, geometry,
-                                 r, z, vpa, vperp, vzeta, vr, vz; run_directory=nothing)
+                                 r, z, vpa, vperp, vzeta, vr, vz)
     code_time = 0.0
     dt = nothing
     dt_before_last_fail = nothing
@@ -605,8 +606,7 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                 restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
                 restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
                 restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
-                                                               vzeta, vr, vz, parallel_io;
-                                                               run_directory=run_directory)
+                                                               vzeta, vr, vz, parallel_io)
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(
@@ -991,7 +991,7 @@ end
 Reload electron pdf and moments from an existing output file.
 """
 function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, time_index,
-                               geometry, r, z, vpa, vperp, vzeta, vr, vz; run_directory=nothing)
+                               geometry, r, z, vpa, vperp, vzeta, vr, vz)
     code_time = 0.0
     previous_runs_info = nothing
     begin_serial_region()
@@ -1019,8 +1019,7 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
                 restart_vperp_spectral, restart_vpa, restart_vpa_spectral, restart_vzeta,
                 restart_vzeta_spectral, restart_vr,restart_vr_spectral, restart_vz,
                 restart_vz_spectral = load_restart_coordinates(fid, r, z, vperp, vpa,
-                                                               vzeta, vr, vz, parallel_io;
-                                                               run_directory=run_directory)
+                                                               vzeta, vr, vz, parallel_io)
 
             # Test whether any interpolation is needed
             interpolation_needed = Dict(
@@ -1086,47 +1085,37 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
     return code_time, previous_runs_info, time_index
 end
 
-function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io;
-                                  run_directory=nothing)
+function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io)
     if parallel_io
         restart_z, restart_z_spectral, _ =
-            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "z"; irank=z.irank, nrank=z.nrank)
         restart_r, restart_r_spectral, _ =
-            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "r"; irank=r.irank, nrank=r.nrank)
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vperp"; irank=vperp.irank, nrank=vperp.nrank)
         restart_vpa, restart_vpa_spectral, _ =
-            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vpa"; irank=vpa.irank, nrank=vpa.nrank)
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vzeta"; irank=vzeta.irank, nrank=vzeta.nrank)
         restart_vr, restart_vr_spectral, _ =
-            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vr"; irank=vr.irank, nrank=vr.nrank)
         restart_vz, restart_vz_spectral, _ =
-            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank,
-                                 run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vz"; irank=vz.irank, nrank=vz.nrank)
     else
         restart_z, restart_z_spectral, _ =
-            load_coordinate_data(fid, "z"; run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "z")
         restart_r, restart_r_spectral, _ =
-            load_coordinate_data(fid, "r"; run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "r")
         restart_vperp, restart_vperp_spectral, _ =
-            load_coordinate_data(fid, "vperp"; run_directory=run_directory,
-                                 ignore_MPI=true)
+            load_coordinate_data(fid, "vperp")
         restart_vpa, restart_vpa_spectral, _ =
-            load_coordinate_data(fid, "vpa"; run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vpa")
         restart_vzeta, restart_vzeta_spectral, _ =
-            load_coordinate_data(fid, "vzeta"; run_directory=run_directory,
-                                 ignore_MPI=true)
+            load_coordinate_data(fid, "vzeta")
         restart_vr, restart_vr_spectral, _ =
-            load_coordinate_data(fid, "vr"; run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vr")
         restart_vz, restart_vz_spectral, _ =
-            load_coordinate_data(fid, "vz"; run_directory=run_directory, ignore_MPI=true)
+            load_coordinate_data(fid, "vz")
 
         if restart_r.nrank != r.nrank
             error("Not using parallel I/O, and distributed MPI layout has "
@@ -3437,24 +3426,23 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
     evolve_density, evolve_upar, evolve_ppar = load_mk_options(file_final_restart)
 
     z_local, z_local_spectral, z_chunk_size =
-        load_coordinate_data(file_final_restart, "z"; ignore_MPI=true)
+        load_coordinate_data(file_final_restart, "z")
     r_local, r_local_spectral, r_chunk_size =
         load_coordinate_data(file_final_restart, "r")
-    r, r_spectral, z, z_spectral = construct_global_zr_coords(r_local, z_local;
-                                                              ignore_MPI=true)
+    r, r_spectral, z, z_spectral = construct_global_zr_coords(r_local, z_local)
 
     vperp, vperp_spectral, vperp_chunk_size =
-        load_coordinate_data(file_final_restart, "vperp"; ignore_MPI=true)
+        load_coordinate_data(file_final_restart, "vperp")
     vpa, vpa_spectral, vpa_chunk_size =
-        load_coordinate_data(file_final_restart, "vpa"; ignore_MPI=true)
+        load_coordinate_data(file_final_restart, "vpa")
 
     if n_neutral_species > 0
         vzeta, vzeta_spectral, vzeta_chunk_size =
-            load_coordinate_data(file_final_restart, "vzeta"; ignore_MPI=true)
+            load_coordinate_data(file_final_restart, "vzeta")
         vr, vr_spectral, vr_chunk_size =
-            load_coordinate_data(file_final_restart, "vr"; ignore_MPI=true)
+            load_coordinate_data(file_final_restart, "vr")
         vz, vz_spectral, vz_chunk_size =
-            load_coordinate_data(file_final_restart, "vz"; ignore_MPI=true)
+            load_coordinate_data(file_final_restart, "vz")
     else
         dummy_adv_input = advection_input("default", 1.0, 0.0, 0.0)
         dummy_comm = MPI.COMM_NULL
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index be55a1119..b773e2f7e 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -308,7 +308,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             reload_evolving_fields!(pdf, moments, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
-                                    vz; run_directory=io_input.output_dir)
+                                    vz)
 
         begin_serial_region()
         @serial_region begin

From 47b646f71c79a213ddb71d3409a979a0790839aa Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Apr 2024 10:50:48 +0100
Subject: [PATCH 254/428] Don't recalculate electron pdf after failed ion
 timestep

---
 moment_kinetics/src/time_advance.jl | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 48545ff00..8693adae5 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1754,7 +1754,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     end
     update_derived_ion_moments_and_derivatives()
 
-    function update_derived_electron_moments_and_derivatives()
+    function update_electrons(update_pdf=true)
         # update the lowest three electron moments (density, upar and ppar)
         calculate_electron_density!(new_scratch.electron_density, moments.electron.dens_updated, new_scratch.density)
         calculate_electron_upar_from_charge_conservation!(new_scratch.electron_upar, moments.electron.upar_updated,
@@ -1810,18 +1810,21 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
                 moments.electron.ppar[iz,ir] = new_scratch.electron_ppar[iz,ir]
             end
 
-            update_electron_pdf!(scratch, pdf.electron.norm, moments, fields.phi, r, z, vperp,
-                                 vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
-                                 vpa_advect, scratch_dummy, t_params.electron, collisions,
-                                 composition, external_source_settings, num_diss_params,
-                                 max_electron_pdf_iterations)
+            if update_pdf
+                update_electron_pdf!(scratch, pdf.electron.norm, moments, fields.phi, r,
+                                     z, vperp, vpa, z_spectral, vperp_spectral,
+                                     vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                     t_params.electron, collisions, composition,
+                                     external_source_settings, num_diss_params,
+                                     max_electron_pdf_iterations)
+            end
         end
         # update the electron parallel friction force
         calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, new_scratch.electron_density,
             new_scratch.electron_upar, new_scratch.upar, moments.electron.dT_dz, composition.me_over_mi,
             collisions.nu_ei, composition.electron_physics)
     end
-    update_derived_electron_moments_and_derivatives()
+    update_electrons()
 
     if composition.n_neutral_species > 0
         ##
@@ -1915,7 +1918,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
             # pdf These need to be re-calculated because `new_scratch` was swapped with
             # the beginning of the timestep, because the timestep failed
             update_derived_ion_moments_and_derivatives()
-            update_derived_electron_moments_and_derivatives()
+            update_electrons(false)
             if composition.n_neutral_species > 0
                 update_derived_neutral_moments_and_derivatives()
             end

From ec354e84934745e1669ad9050f21ad0870127190 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 25 Apr 2024 21:59:55 +0100
Subject: [PATCH 255/428] Fix merge of Krook collisions - update ee and ei
 collision setup

---
 ...wall+sheath-bc_kinetic_krook_loworder.toml |  4 +-
 ...fraction0.5_split3_boltzmann-vpadiss0.toml |  4 +-
 ...c_recyclefraction0.5_split3_boltzmann.toml |  4 +-
 ...lefraction0.5_split3_kinetic-vpadiss0.toml |  4 +-
 .../src/electron_kinetic_equation.jl          |  2 +-
 moment_kinetics/src/fokker_planck.jl          |  4 +-
 moment_kinetics/src/input_structs.jl          |  8 ++-
 moment_kinetics/src/krook_collisions.jl       | 68 ++++++++++++-------
 moment_kinetics/src/moment_kinetics_input.jl  |  8 ++-
 moment_kinetics/src/reference_parameters.jl   | 66 +++++++++++++++---
 10 files changed, 127 insertions(+), 45 deletions(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
index 363955faf..da7a1bf63 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic_krook_loworder.toml
@@ -43,7 +43,6 @@ vpa_IC_temperature_phase2 = 0.0
 charge_exchange_frequency = 2.0
 electron_charge_exchange_frequency = 0.0
 nu_ei = 0.0
-krook_collisions_option = "reference_parameters"
 ionization_frequency = 2.0
 electron_ionization_frequency = 2.0
 ionization_energy = 1.0
@@ -81,6 +80,9 @@ ascii_output = true
 #vpa_dissipation_coefficient = 2.0
 vpa_dissipation_coefficient = 20.0
 
+[krook_collisions]
+use_krook = true
+
 [timestepping]
 nstep = 40000
 #nstep = 1
diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
index 97e534663..098618577 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = true
 evolve_moments_parallel_pressure = true
 evolve_moments_conservation = true
 recycling_fraction = 0.5
-krook_collisions_option = "reference_parameters"
 T_e = 0.2 # 1.0
 T_wall = 0.1
 initial_density1 = 1.0
@@ -93,3 +92,6 @@ force_minimum_pdf_value = 0.0
 #vz_dissipation_coefficient = 1.0e-2
 #vz_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true
diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
index 23ba4265d..4b36701ae 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = true
 evolve_moments_parallel_pressure = true
 evolve_moments_conservation = true
 recycling_fraction = 0.5
-krook_collisions_option = "reference_parameters"
 T_e = 0.2 # 1.0
 T_wall = 0.1
 initial_density1 = 1.0
@@ -93,3 +92,6 @@ vz_dissipation_coefficient = 1.0e-1
 #vz_dissipation_coefficient = 1.0e-2
 #vz_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true
diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
index d4b93cd94..ceaf3a8c8 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = true
 evolve_moments_parallel_pressure = true
 evolve_moments_conservation = true
 recycling_fraction = 0.5
-krook_collisions_option = "reference_parameters"
 T_e = 0.2 # 1.0
 T_wall = 0.1
 initial_density1 = 1.0
@@ -110,3 +109,6 @@ force_minimum_pdf_value = 0.0
 #vz_dissipation_coefficient = 1.0e-2
 #vz_dissipation_coefficient = 1.0e-3
 force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 0bb718299..48db01646 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1545,7 +1545,7 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
     add_dissipation_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, scratch_dummy,
                           z_spectral, z, vpa, vpa_spectral, num_diss_params, dt)
 
-    if collisions.krook_collision_frequency_prefactor_ee > 0.0
+    if collisions.krook.nuee0 > 0.0 || collisions.krook.nuei0 > 0.0
         # Add a Krook collision operator
         # Set dt=-1 as we update the residual here rather than adding an update to
         # 'fvec_out'.
diff --git a/moment_kinetics/src/fokker_planck.jl b/moment_kinetics/src/fokker_planck.jl
index 52b2572ed..012183483 100644
--- a/moment_kinetics/src/fokker_planck.jl
+++ b/moment_kinetics/src/fokker_planck.jl
@@ -51,7 +51,7 @@ using ..velocity_moments: integrate_over_vspace
 using ..velocity_moments: get_density, get_upar, get_ppar, get_pperp, get_qpar, get_pressure, get_rmom
 using ..looping
 using ..input_structs: fkpl_collisions_input, set_defaults_and_check_section!
-using ..reference_parameters: get_reference_collision_frequency
+using ..reference_parameters: get_reference_collision_frequency_ii
 using ..fokker_planck_calculus: init_Rosenbluth_potential_integration_weights!
 using ..fokker_planck_calculus: init_Rosenbluth_potential_boundary_integration_weights!
 using ..fokker_planck_calculus: allocate_boundary_integration_weights
@@ -81,7 +81,7 @@ frequency_option = "manual"
 """
 function setup_fkpl_collisions_input(toml_input::Dict, reference_params)
     # get reference collision frequency (note factor of 1/2 due to definition choices)
-    nuii_fkpl_default = 0.5*get_reference_collision_frequency(reference_params)
+    nuii_fkpl_default = 0.5*get_reference_collision_frequency_ii(reference_params)
     # read the input toml and specify a sensible default
     input_section = set_defaults_and_check_section!(toml_input, "fokker_planck_collisions",
        # begin default inputs (as kwargs)
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index c6e5feb68..a7cd7a0a6 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -356,8 +356,12 @@ end
 """
 Base.@kwdef struct krook_collisions_input
     use_krook::Bool
-    # Coulomb collision rate at the reference density and temperature
-    krook_collision_frequency_prefactor::mk_float#
+    # Ion-ion Coulomb collision rate at the reference density and temperature
+    nuii0::mk_float
+    # Electron-electron Coulomb collision rate at the reference density and temperature
+    nuee0::mk_float
+    # Electron-ion Coulomb collision rate at the reference density and temperature
+    nuei0::mk_float
     # Setting to switch between different options for Krook collision operator
     frequency_option::String # "reference_parameters" # "manual", 
 end
diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index 39acf71fe..f7865a3f0 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -8,7 +8,9 @@ export setup_krook_collisions_input, get_collision_frequency_ii, get_collision_f
 using ..constants
 using ..looping
 using ..input_structs: krook_collisions_input, set_defaults_and_check_section!
-using ..reference_parameters: get_reference_collision_frequency
+using ..reference_parameters: get_reference_collision_frequency_ii,
+                              get_reference_collision_frequency_ee,
+                              get_reference_collision_frequency_ei
 
 
 """
@@ -22,18 +24,24 @@ frequency_option = "manual"
 """
 function setup_krook_collisions_input(toml_input::Dict, reference_params)
     # get reference collision frequency
-    nuii_krook_default = get_reference_collision_frequency(reference_params)
+    nuii_krook_default = get_reference_collision_frequency_ii(reference_params)
+    nuee_krook_default = get_reference_collision_frequency_ee(reference_params)
+    nuei_krook_default = get_reference_collision_frequency_ei(reference_params)
     # read the input toml and specify a sensible default    
     input_section = input_section = set_defaults_and_check_section!(toml_input, "krook_collisions",
        # begin default inputs (as kwargs)
        use_krook = false,
-       krook_collision_frequency_prefactor = -1.0,
+       nuii0 = -1.0,
+       nuee0 = -1.0,
+       nuei0 = -1.0,
        frequency_option = "reference_parameters")
        
     # ensure that the collision frequency is consistent with the input option
     frequency_option = input_section["frequency_option"]
     if frequency_option == "reference_parameters"
-        input_section["krook_collision_frequency_prefactor"] = nuii_krook_default
+        input_section["nuii0"] = nuii_krook_default
+        input_section["nuee0"] = nuee_krook_default
+        input_section["nuei0"] = nuei_krook_default
     elseif frequency_option == "manual" 
         # use the frequency from the input file
         # do nothing
@@ -44,7 +52,9 @@ function setup_krook_collisions_input(toml_input::Dict, reference_params)
     # finally, ensure prefactor < 0 if use_krook is false
     # so that prefactor > 0 is the only check required in the rest of the code
     if !input_section["use_krook"]
-        input_section["krook_collision_frequency_prefactor"] = -1.0
+        input_section["nuii0"] = -1.0
+        input_section["nuee0"] = -1.0
+        input_section["nuei0"] = -1.0
     end
     input = Dict(Symbol(k)=>v for (k,v) in input_section)
     #println(input)
@@ -61,16 +71,20 @@ Calculate the ion-ion collision frequency, depending on the settings/parameters
 together.
 """
 function get_collision_frequency_ii(collisions, n, vth)
-    if collisions.krook_collisions_option == "reference_parameters"
-        return @. collisions.krook_collision_frequency_prefactor_ii * n * vth^(-3)
-    elseif collisions.krook_collisions_option == "manual"
+    # extract krook options from collisions struct
+    colk = collisions.krook
+    nuii0 = colk.nuii0
+    frequency_option = colk.frequency_option
+    if frequency_option == "reference_parameters"
+        return @. nuii0 * n * vth^(-3)
+    elseif frequency_option == "manual"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
         # returning a scalar from this branch but an array from the "reference_parameters"
         # branch).
-        return @. collisions.krook_collision_frequency_prefactor_ii + 0.0 * n
-    elseif collisions.krook_collisions_option == "none"
+        return @. nuii0 + 0.0 * n
+    elseif frequency_option == "none"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
@@ -78,8 +92,8 @@ function get_collision_frequency_ii(collisions, n, vth)
         # branch).
         return @. 0.0 * n
     else
-        error("Unrecognised option "
-              * "krook_collisions_option=$(collisions.krook_collisions_option)")
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_option)")
     end
 end
 
@@ -93,16 +107,20 @@ in `collisions`, for the given density `n` and electron thermal speed `vthe`.
 together.
 """
 function get_collision_frequency_ee(collisions, n, vthe)
-    if collisions.krook_collisions_option == "reference_parameters"
-        return @. collisions.krook_collision_frequency_prefactor_ee * n * vthe^(-3)
-    elseif collisions.krook_collisions_option == "manual"
+    # extract krook options from collisions struct
+    colk = collisions.krook
+    nuee0 = colk.nuee0
+    frequency_option = colk.frequency_option
+    if frequency_option == "reference_parameters"
+        return @. nuee0 * n * vthe^(-3)
+    elseif frequency_option == "manual"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
         # returning a scalar from this branch but an array from the "reference_parameters"
         # branch).
-        return @. collisions.krook_collision_frequency_prefactor_ee + 0.0 * n
-    elseif collisions.krook_collisions_option == "none"
+        return @. nuee0 + 0.0 * n
+    elseif frequency_option == "none"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
@@ -110,8 +128,8 @@ function get_collision_frequency_ee(collisions, n, vthe)
         # branch).
         return @. 0.0 * n
     else
-        error("Unrecognised option "
-              * "krook_collisions_option=$(collisions.krook_collisions_option)")
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_option)")
     end
 end
 
@@ -127,18 +145,18 @@ together.
 function get_collision_frequency_ei(collisions, n, vthe)
     # extract krook options from collisions struct
     colk = collisions.krook
-    krook_collision_frequency_prefactor = colk.krook_collision_frequency_prefactor
+    nuei0 = colk.nuei0
     frequency_option = colk.frequency_option
     if frequency_option == "reference_parameters"
-        return @. krook_collision_frequency_prefactor_ei * n * vthe^(-3)
+        return @. nuei0 * n * vthe^(-3)
     elseif frequency_option == "manual"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
         # returning a scalar from this branch but an array from the "reference_parameters"
         # branch).
-        return @. collisions.krook_collision_frequency_prefactor_ei + 0.0 * n
-    elseif collisions.krook_collisions_option == "none"
+        return @. nuei0 + 0.0 * n
+    elseif frequency_option == "none"
         # Include 0.0*n so that the result gets promoted to an array if n is an array,
         # which hopefully means this function will have a fixed return type given the
         # types of the arguments (we don't want to be 'type unstable' for array inputs by
@@ -146,8 +164,8 @@ function get_collision_frequency_ei(collisions, n, vthe)
         # branch).
         return @. 0.0 * n
     else
-        error("Unrecognised option "
-              * "krook_collisions_option=$(collisions.krook_collisions_option)")
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_option)")
     end
 end
 
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 57df47308..6e1f36273 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -118,6 +118,9 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     # normalised values used in the code.
     reference_params = setup_reference_parameters(scan_input)
 
+    # Set me_over_mi here so we can use reference_params
+    composition.me_over_mi = reference_params.me / reference_params.mref
+
     ## set geometry_input
     geometry_in = setup_geometry_input(scan_input, get_default_rhostar(reference_params))
     
@@ -1109,8 +1112,9 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     Er_constant = 0.0
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0
-    # ratio of the electron particle mass to the ion particle mass
-    me_over_mi = 1.0/1836.0/2.0
+    # ratio of the electron particle mass to the ion particle mass - value set later using
+    # reference_params
+    me_over_mi = NaN
     # The ion flux reaching the wall that is recycled as neutrals is reduced by
     # `recycling_fraction` to account for ions absorbed by the wall.
     recycling_fraction = 1.0
diff --git a/moment_kinetics/src/reference_parameters.jl b/moment_kinetics/src/reference_parameters.jl
index b33967b21..97b771228 100644
--- a/moment_kinetics/src/reference_parameters.jl
+++ b/moment_kinetics/src/reference_parameters.jl
@@ -8,7 +8,8 @@ physical units of the simulation, and are needed for a few specific steps during
 module reference_parameters
 
 export setup_reference_parameters
-export get_reference_collision_frequency
+export get_reference_collision_frequency_ii, get_reference_collision_frequency_ee,
+       get_reference_collision_frequency_ei
 
 using ..constants
 using ..input_structs
@@ -32,6 +33,8 @@ function setup_reference_parameters(input_dict)
     Nref_per_cm3 = reference_parameter_section["Nref"] * 1.0e-6
     Tref = reference_parameter_section["Tref"]
 
+    reference_parameter_section["me"] = electron_mass
+
     # Coulomb logarithm at reference parameters for same-species, singly-charged ion-ion
     # collisions, using NRL formulary. Formula given for n in units of cm^-3 and T in
     # units of eV.
@@ -58,22 +61,17 @@ function setup_reference_parameters(input_dict)
 end
 
 """
-Calculate normalized collision frequency at reference parameters for Coulomb collisions.
+Calculate normalized ion-ion collision frequency at reference parameters for Coulomb collisions.
 
 Currently valid only for hydrogenic ions (Z=1)
 """
-function get_reference_collision_frequency(reference_params)
+function get_reference_collision_frequency_ii(reference_params)
     Nref = reference_params.Nref
     Tref = reference_params.Tref
     mref = reference_params.mref
     timeref = reference_params.timeref
     cref = reference_params.cref
-
-    Nref_per_cm3 = Nref * 1.0e-6
-
-    # Coulomb logarithm at reference parameters for same-species ion-ion collisions, using
-    # NRL formulary. Formula given for n in units of cm^-3 and T in units of eV.
-    logLambda_ii = 23.0 - log(sqrt(2.0*Nref_per_cm3) / Tref^1.5)
+    logLambda_ii = reference_params.logLambda_ii
 
     # Collision frequency, using \hat{\nu} from Appendix, p. 277 of Helander "Collisional
     # Transport in Magnetized Plasmas" (2002).
@@ -84,4 +82,54 @@ function get_reference_collision_frequency(reference_params)
     return nu_ii0
 end
 
+"""
+Calculate normalized electron-electron collision frequency at reference parameters for Coulomb collisions.
+"""
+function get_reference_collision_frequency_ee(reference_params)
+    Nref = reference_params.Nref
+    Tref = reference_params.Tref
+    me = reference_params.me
+    timeref = reference_params.timeref
+    cref = reference_params.cref
+    logLambda_ee = reference_params.logLambda_ee
+
+    # Collision frequency, using \hat{\nu} from Appendix, p. 277 of Helander "Collisional
+    # Transport in Magnetized Plasmas" (2002).
+    # Note the electron thermal speed used in the code is normalised to cref, so we use
+    # cref in these two formulas rather than a reference electron thermal speed, so that
+    # when multiplied by the normalised electron thermal speed we get the correct
+    # normalised collision frequency.
+    nu_ee0_per_s = Nref * proton_charge^4 * logLambda_ee  /
+                   (4.0 * π * epsilon0^2 * me^2 * cref^3) # s^-1
+    nu_ee0 = nu_ee0_per_s * timeref
+
+    return nu_ee0
+end
+
+"""
+Calculate normalized electron-ion collision frequency at reference parameters for Coulomb collisions.
+
+Currently valid only for hydrogenic ions (Z=1)
+"""
+function get_reference_collision_frequency_ei(reference_params)
+    Nref = reference_params.Nref
+    Tref = reference_params.Tref
+    me = reference_params.me
+    timeref = reference_params.timeref
+    cref = reference_params.cref
+    logLambda_ei = reference_params.logLambda_ei
+
+    # Collision frequency, using \hat{\nu} from Appendix, p. 277 of Helander "Collisional
+    # Transport in Magnetized Plasmas" (2002).
+    # Note the electron thermal speed used in the code is normalised to cref, so we use
+    # cref in these two formulas rather than a reference electron thermal speed, so that
+    # when multiplied by the normalised electron thermal speed we get the correct
+    # normalised collision frequency.
+    nu_ei0_per_s = Nref * proton_charge^4 * logLambda_ei  /
+                   (4.0 * π * epsilon0^2 * me^2 * cref^3) # s^-1
+    nu_ei0 = nu_ei0_per_s * timeref
+
+    return nu_ei0
+end
+
 end

From 7b4d8f93fb2cba85890bed67243603767c23934c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Apr 2024 09:30:35 +0100
Subject: [PATCH 256/428] Fix pressure flooring in
 update_electron_vth_temperature!()

---
 moment_kinetics/src/electron_fluid_equations.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 2359af6e1..357ec0e3b 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -379,7 +379,7 @@ function update_electron_vth_temperature!(moments, ppar, dens, composition)
     vth = moments.electron.vth
     @loop_r_z ir iz begin
         p = max(ppar[iz,ir], 0.0)
-        temp[iz,ir] = 2 * p[iz,ir] / dens[iz,ir]
+        temp[iz,ir] = 2 * p / dens[iz,ir]
         vth[iz,ir] = sqrt(temp[iz,ir] / composition.me_over_mi)
     end
     moments.electron.temp_updated[] = true

From 30fcfe97f448109d0e56c69952117fda6318058d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Apr 2024 09:46:56 +0100
Subject: [PATCH 257/428] Fix wall+sheath-bc_kinetic.toml

---
 examples/kinetic-electrons/wall+sheath-bc_kinetic.toml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
index 9fde4911c..a0a928479 100644
--- a/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
+++ b/examples/kinetic-electrons/wall+sheath-bc_kinetic.toml
@@ -64,7 +64,8 @@ vpa_bc = "zero"
 vpa_discretization = "chebyshev_pseudospectral"
 #vpa_discretization = "gausslegendre_pseudospectral"
 vz_ngrid = 17
-vz_nelement = 10
+#vz_nelement = 10
+vz_nelement = 20
 vz_L = 8.0
 vz_bc = "zero"
 vz_discretization = "chebyshev_pseudospectral"

From 30a3d5f62c62c2eb6227ae202761b69a21e58632 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Apr 2024 10:56:48 +0100
Subject: [PATCH 258/428] Fix coordinate loading from NetCDF files

Use `get_subgroup_keys()` instead of just `keys()` to support NetCDF as
well as HDF5.
---
 moment_kinetics/src/load_data.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 5be2a7550..787dcc6ee 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -228,7 +228,7 @@ function load_coordinate_data(fid, name; printout=false, irank=nothing, nrank=no
     parallel_io = load_variable(overview, "parallel_io")
 
     coords_group = get_group(fid, "coords")
-    if name ∈ keys(coords_group)
+    if name ∈ get_subgroup_keys(coords_group)
         coord_group = get_group(coords_group, name)
     else
         return nothing, nothing, nothing

From 9be23e78460e46defa5f02d60a4f626bb6f97f59 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Apr 2024 16:48:27 +0100
Subject: [PATCH 259/428] Fix formatting in assign_endpoint!()

---
 moment_kinetics/src/calculus.jl | 82 ++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl
index 1a10376db..405567e52 100644
--- a/moment_kinetics/src/calculus.jl
+++ b/moment_kinetics/src/calculus.jl
@@ -347,47 +347,47 @@ in the main code
 
 function assign_endpoint!(df1d::AbstractArray{mk_float,Ndims},
  receive_buffer::AbstractArray{mk_float,Mdims},key::String,coord) where {Ndims,Mdims}
-	if key == "lower"
-		j = 1
-	elseif key == "upper"
-		j = coord.n
-	else
-		println("ERROR: invalid key in assign_endpoint!")
-	end
-        # test against coord name -- make sure to use exact string delimiters e.g. "x" not 'x'
-        # test against Ndims (autodetermined) to choose which array slices to use in assigning endpoints
-        #println("DEBUG MESSAGE: coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
-        if coord.name == "z" && Ndims==2
-            df1d[j,:] .= receive_buffer[:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "z" && Ndims==3
-            df1d[j,:,:] .= receive_buffer[:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "z" && Ndims==4
-            df1d[:,:,j,:] .= receive_buffer[:,:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "z" && Ndims==5
-            df1d[:,:,j,:,:] .= receive_buffer[:,:,:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "z" && Ndims==6
-            df1d[:,:,:,j,:,:] .= receive_buffer[:,:,:,:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "r" && Ndims==2
-            df1d[:,j] .= receive_buffer[:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "r" && Ndims==3
-            df1d[:,j,:] .= receive_buffer[:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "r" && Ndims==4
-            df1d[:,:,:,j] .= receive_buffer[:,:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "r" && Ndims==5
-            df1d[:,:,:,j,:] .= receive_buffer[:,:,:,:]
-            #println("ASSIGNING DATA")
-        elseif coord.name == "r" && Ndims==6
-            df1d[:,:,:,:,j,:] .= receive_buffer[:,:,:,:,:]
-            #println("ASSIGNING DATA")
-        else
+    if key == "lower"
+            j = 1
+    elseif key == "upper"
+            j = coord.n
+    else
+            println("ERROR: invalid key in assign_endpoint!")
+    end
+    # test against coord name -- make sure to use exact string delimiters e.g. "x" not 'x'
+    # test against Ndims (autodetermined) to choose which array slices to use in assigning endpoints
+    #println("DEBUG MESSAGE: coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
+    if coord.name == "z" && Ndims==2
+        df1d[j,:] .= receive_buffer[:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "z" && Ndims==3
+        df1d[j,:,:] .= receive_buffer[:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "z" && Ndims==4
+        df1d[:,:,j,:] .= receive_buffer[:,:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "z" && Ndims==5
+        df1d[:,:,j,:,:] .= receive_buffer[:,:,:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "z" && Ndims==6
+        df1d[:,:,:,j,:,:] .= receive_buffer[:,:,:,:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==2
+        df1d[:,j] .= receive_buffer[:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==3
+        df1d[:,j,:] .= receive_buffer[:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==4
+        df1d[:,:,:,j] .= receive_buffer[:,:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==5
+        df1d[:,:,:,j,:] .= receive_buffer[:,:,:,:]
+        #println("ASSIGNING DATA")
+    elseif coord.name == "r" && Ndims==6
+        df1d[:,:,:,:,j,:] .= receive_buffer[:,:,:,:,:]
+        #println("ASSIGNING DATA")
+    else
         error("ERROR: failure to assign endpoints in reconcile_element_boundaries_MPI! (centered): coord.name: ",coord.name," Ndims: ",Ndims," key: ",key)
     end
 end

From 7f9327869e3aef3d53fdbbba9df31625b3008a56 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 30 May 2024 11:46:04 +0100
Subject: [PATCH 260/428] Better error message for negative timestep from step
 to output

---
 moment_kinetics/src/runge_kutta.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 60fe6d380..0bf317f00 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1251,6 +1251,11 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             t_params.dt_before_output[] = t_params.dt[]
             t_params.dt[] = t_params.next_output_time[] - current_time
             t_params.step_to_output[] = true
+
+            if t_params.dt[] < 0.0
+                error("When trying to step to next output time, made negative timestep "
+                      * "dt=$(t_params.dt[])")
+            end
         end
     end
 

From 25bba757d5ed512b560f345830e621d700c76a70 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 31 May 2024 14:44:19 +0100
Subject: [PATCH 261/428] Support fixed-step-count output for kinetic electrons

---
 .../src/electron_kinetic_equation.jl          | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index bf37f8a07..e9269579b 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -192,13 +192,12 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         do_debug_io = false
     end
     @serial_region begin
-        if io_electron === nothing || do_debug_io
-            # When doing debug_io, we don't want to make the adaptive timestep adjust to
-            # output at specific times - instead we just output after a fixed number of
-            # steps, however long those steps were.
-            t_params.next_output_time[] = Inf
-        else
+        if t_params.adaptive && !t_params.write_after_fixed_step_count && io_electron !== nothing
             t_params.next_output_time[] = dfns_output_times[1]
+        else
+            # Using fixed output step count, so we don't want to make the adaptive
+            # timestep adjust to output at specific times.
+            t_params.next_output_time[] = Inf
         end
     end
 
@@ -467,7 +466,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
 
         if text_output
-            if (mod(t_params.step_counter[] - initial_step_counter, t_params.nwrite)==1)
+            if (mod(t_params.step_counter[] - initial_step_counter, t_params.nwrite_moments)==1)
                 begin_serial_region()
                 @serial_region begin
                     @loop_vpa ivpa begin
@@ -488,7 +487,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
         end
-        if ((time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon)
+        if ((t_params.adaptive && (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon))
+            || (!t_params.adaptive && t_params.step_counter[] % t_params.nwrite_moments == 0)
             || (do_debug_io && (t_params.step_counter[] % debug_io_nwrite == 0)))
 
             begin_serial_region()
@@ -516,7 +516,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
             output_counter += 1
-            if output_counter - initial_output_counter ≤ length(dfns_output_times)
+            if t_params.adaptive && !t_params.write_after_fixed_step_count &&
+                    (output_counter - initial_output_counter ≤ length(dfns_output_times))
                 @serial_region begin
                     t_params.next_output_time[] =
                         dfns_output_times[output_counter - initial_output_counter]

From 5a7d1936e486a8f67e4afc9e8fce672957067fe7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 30 May 2024 09:14:58 +0100
Subject: [PATCH 262/428] Separate `scratch_electron` struct for electron
 pseudo-timestepping

More logical than having electrons share the `scratch_pdf` structs, and
makes it easier to fit the `update_electron_pdf!()` call into refactored
ion/neutral timestepping functions.
---
 moment_kinetics/src/initial_conditions.jl     |  18 +-
 moment_kinetics/src/moment_kinetics.jl        |  17 +-
 .../src/moment_kinetics_structs.jl            |  12 +-
 moment_kinetics/src/time_advance.jl           | 432 ++++++++++--------
 4 files changed, 261 insertions(+), 218 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 0d6e7397e..dee497b0b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -257,9 +257,10 @@ end
 function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                                vperp_spectral, vpa_spectral, collisions, gyroavs,
-                               external_source_settings, scratch_dummy, scratch, t_params,
-                               t_input, num_diss_params, advection_structs, io_input,
-                               input_dict; restart_electron_physics)
+                               external_source_settings, scratch_dummy, scratch,
+                               scratch_electron, t_params, t_input, num_diss_params,
+                               advection_structs, io_input, input_dict;
+                               restart_electron_physics)
     
     moments.electron.dens_updated[] = false
     # initialise the electron density profile
@@ -436,12 +437,15 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
         scratch[1].electron_temp .= moments.electron.temp
     end
-    # get the initial electrostatic potential and parallel electric field
-    update_phi!(fields, scratch[1], vperp, z, r, composition, collisions, moments,
-                z_spectral, r_spectral, scratch_dummy, gyroavs)
+    if scratch_electron !== nothing
+        begin_serial_region()
+        @serial_region begin
+            scratch_electron[1].electron_ppar .= moments.electron.ppar
+        end
+    end
 
     # initialize the electron pdf that satisfies the electron kinetic equation
-    initialize_electron_pdf!(scratch, pdf, moments, fields.phi, r, z, vpa, vperp,
+    initialize_electron_pdf!(scratch_electron, pdf, moments, fields.phi, r, z, vpa, vperp,
                              vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
                              advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index dd72994c8..f76935e85 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -342,8 +342,9 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
-    moments, spectral_objects, scratch, scratch_implicit, advance, advance_implicit,
-    t_params, fp_arrays, gyroavs, manufactured_source_list, nl_solver_params =
+    moments, spectral_objects, scratch, scratch_implicit, scratch_electron, advance,
+    advance_implicit, t_params, fp_arrays, gyroavs, manufactured_source_list,
+    nl_solver_params =
         setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrophase,
             vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
             z_spectral, r_spectral, composition, moments, t_input, code_time, dt,
@@ -376,12 +377,12 @@ function setup_moment_kinetics(input_dict::AbstractDict;
 
     begin_s_r_z_vperp_region()
 
-    return pdf, scratch, scratch_implicit, code_time, t_params, vz, vr, vzeta, vpa, vperp,
-           gyrophase, z, r, moments, fields, spectral_objects, advection_structs,
-           composition, collisions, geometry, gyroavs, boundary_distributions,
-           external_source_settings, num_diss_params, nl_solver_params, advance,
-           advance_implicit, fp_arrays, scratch_dummy, manufactured_source_list, ascii_io,
-           io_moments, io_dfns
+    return pdf, scratch, scratch_implicit, scratch_electron, code_time, t_params, vz, vr,
+           vzeta, vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
+           advection_structs, composition, collisions, geometry, gyroavs,
+           boundary_distributions, external_source_settings, num_diss_params,
+           nl_solver_params, advance, advance_implicit, fp_arrays, scratch_dummy,
+           manufactured_source_list, ascii_io, io_moments, io_dfns
 end
 
 """
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index b64e14a12..acb1a378e 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -9,8 +9,7 @@ using ..type_definitions: mk_float
 
 """
 """
-struct scratch_pdf{n_distribution_ion, n_moment, 
-                   n_distribution_electron, n_moment_electron, 
+struct scratch_pdf{n_distribution_ion, n_moment, n_moment_electron,
                    n_distribution_neutral, n_moment_neutral}
     # ions
     pdf::MPISharedArray{mk_float, n_distribution_ion}
@@ -20,7 +19,6 @@ struct scratch_pdf{n_distribution_ion, n_moment,
     pperp::MPISharedArray{mk_float, n_moment}
     temp_z_s::MPISharedArray{mk_float, n_moment}
     # electrons
-    pdf_electron::MPISharedArray{mk_float, n_distribution_electron}
     electron_density::MPISharedArray{mk_float, n_moment_electron}
     electron_upar::MPISharedArray{mk_float, n_moment_electron}
     electron_ppar::MPISharedArray{mk_float, n_moment_electron}
@@ -33,6 +31,14 @@ struct scratch_pdf{n_distribution_ion, n_moment,
     pz_neutral::MPISharedArray{mk_float, n_moment_neutral}
 end
 
+"""
+"""
+struct scratch_electron_pdf{n_distribution_electron, n_moment_electron}
+    # electrons
+    pdf_electron::MPISharedArray{mk_float, n_distribution_electron}
+    electron_ppar::MPISharedArray{mk_float, n_moment_electron}
+end
+
 """
 """
 struct em_fields_struct
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index ff1fa8ec0..f341b6078 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -18,7 +18,7 @@ using ..file_io: write_data_to_ascii, write_all_moments_data_to_binary,
                  write_all_dfns_data_to_binary, debug_dump, setup_electron_io
 using ..initial_conditions: initialize_electrons!
 using ..looping
-using ..moment_kinetics_structs: scratch_pdf
+using ..moment_kinetics_structs: scratch_pdf, scratch_electron_pdf
 using ..velocity_moments: update_moments!, update_moments_neutral!, reset_moments_status!, update_derived_moments!, update_derived_moments_neutral!
 using ..velocity_moments: update_density!, update_upar!, update_ppar!, update_pperp!, update_qpar!, update_vth!
 using ..velocity_moments: update_neutral_density!, update_neutral_qz!
@@ -625,16 +625,18 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
 
     # create an array of structs containing scratch arrays for the pdf and low-order moments
     # that may be evolved separately via fluid equations
-    n_rk_stages = t_params.n_rk_stages
-    if t_params.electron !== nothing
-        n_rk_stages = max(n_rk_stages, t_params.electron.n_rk_stages)
-    end
-    scratch = setup_scratch_arrays(moments, pdf, n_rk_stages + 1)
+    scratch = setup_scratch_arrays(moments, pdf, t_params.n_rk_stages + 1)
     if t_params.rk_coefs_implicit !== nothing
-        scratch_implicit = setup_scratch_arrays(moments, pdf, n_rk_stages)
+        scratch_implicit = setup_scratch_arrays(moments, pdf, t_params.n_rk_stages)
     else
         scratch_implicit = nothing
     end
+    if composition.electron_physics == kinetic_electrons
+        scratch_electron = setup_electron_scratch_arrays(moments, pdf,
+                                                         t_params.electron.n_rk_stages+1)
+    else
+        scratch_electron = nothing
+    end
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     # create arrays for Fokker-Planck collisions 
@@ -652,9 +654,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions, gyroavs,
-                              external_source_settings, scratch_dummy, scratch, t_params,
-                              t_input, num_diss_params, advection_structs, io_input,
-                              input_dict;
+                              external_source_settings, scratch_dummy, scratch,
+                              scratch_electron, t_params, t_input, num_diss_params,
+                              advection_structs, io_input, input_dict;
                               restart_electron_physics=restart_electron_physics)
     elseif restarting && composition.electron_physics == kinetic_electrons &&
            t_params.electron.debug_io !== nothing
@@ -933,9 +935,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     # Ensure all processes are synchronized at the end of the setup
     _block_synchronize()
 
-    return moments, spectral_objects, scratch, scratch_implicit, advance,
-           advance_implicit, t_params, fp_arrays, gyroavs, manufactured_source_list,
-           nl_solver_params
+    return moments, spectral_objects, scratch, scratch_implicit, scratch_electron,
+           advance, advance_implicit, t_params, fp_arrays, gyroavs,
+           manufactured_source_list, nl_solver_params
 end
 
 """
@@ -1406,21 +1408,14 @@ create an array of structs containing scratch arrays for the normalised pdf and
 that may be evolved separately via fluid equations
 """
 function setup_scratch_arrays(moments, pdf, n)
-    # will create n_rk_stages+1 structs, each of which will contain one pdf,
+    # will create n structs, each of which will contain one pdf,
     # density, parallel flow, parallel pressure, and perpendicular pressure array for ions
     # (possibly) the same for electrons, and the same for neutrals. The actual array will
     # be created at the end of the first step of the loop below, once we have a
     # `scratch_pdf` object of the correct type.
-    scratch = Vector{scratch_pdf{5,3,4,2,6,3}}(undef, n)
+    scratch = Vector{scratch_pdf{5,3,2,6,3}}(undef, n)
     pdf_dims = size(pdf.ion.norm)
     moment_dims = size(moments.ion.dens)
-    if pdf.electron !== nothing
-        using_electrons = true
-        pdf_electron_dims = size(pdf.electron.norm)
-    else
-        using_electrons = false
-        pdf_electron_dims = (0,0,0,0)
-    end
     moment_electron_dims = size(moments.electron.dens)
 
     pdf_neutral_dims = size(pdf.neutral.norm)
@@ -1436,7 +1431,6 @@ function setup_scratch_arrays(moments, pdf, n)
         pperp_array = allocate_shared_float(moment_dims...)
         temp_array = allocate_shared_float(moment_dims...)
 
-        pdf_electron_array = allocate_shared_float(pdf_electron_dims...)
         density_electron_array = allocate_shared_float(moment_electron_dims...)
         upar_electron_array = allocate_shared_float(moment_electron_dims...)
         ppar_electron_array = allocate_shared_float(moment_electron_dims...)
@@ -1451,11 +1445,11 @@ function setup_scratch_arrays(moments, pdf, n)
 
         scratch[istage] = scratch_pdf(pdf_array, density_array, upar_array,
                                       ppar_array, pperp_array, temp_array,
-                                      pdf_electron_array, density_electron_array,
-                                      upar_electron_array, ppar_electron_array,
-                                      pperp_electron_array, temp_electron_array,
-                                      pdf_neutral_array, density_neutral_array,
-                                      uz_neutral_array, pz_neutral_array)
+                                      density_electron_array, upar_electron_array,
+                                      ppar_electron_array, pperp_electron_array,
+                                      temp_electron_array, pdf_neutral_array,
+                                      density_neutral_array, uz_neutral_array,
+                                      pz_neutral_array)
         @serial_region begin
             scratch[istage].pdf .= pdf.ion.norm
             scratch[istage].density .= moments.ion.dens
@@ -1463,9 +1457,6 @@ function setup_scratch_arrays(moments, pdf, n)
             scratch[istage].ppar .= moments.ion.ppar
             scratch[istage].pperp .= moments.ion.pperp
 
-            if using_electrons
-                scratch[istage].pdf_electron .= pdf.electron.norm
-            end
             scratch[istage].electron_density .= moments.electron.dens
             scratch[istage].electron_upar .= moments.electron.upar
             scratch[istage].electron_ppar .= moments.electron.ppar
@@ -1480,6 +1471,31 @@ function setup_scratch_arrays(moments, pdf, n)
     return scratch
 end
 
+function setup_electron_scratch_arrays(moments, pdf, n)
+    # will create n structs, each of which will contain one pdf, and parallel pressure
+    # array for electrons.
+    # The actual array will be created at the end of the first step of the loop below,
+    # once we have a `scratch_electron_pdf` object of the correct type.
+    scratch = Vector{scratch_electron_pdf{4,2}}(undef, n)
+    pdf_dims = size(pdf.electron.norm)
+    moment_dims = size(moments.electron.dens)
+
+    # populate each of the structs
+    for istage ∈ 1:n
+        # Allocate arrays in temporary variables so that we can identify them
+        # by source line when using @debug_shared_array
+        pdf_array = allocate_shared_float(pdf_dims...)
+        ppar_array = allocate_shared_float(moment_dims...)
+
+        scratch[istage] = scratch_electron_pdf(pdf_array, ppar_array)
+        @serial_region begin
+            scratch[istage].pdf_electron .= pdf.electron.norm
+            scratch[istage].electron_ppar .= moments.electron.ppar
+        end
+    end
+    return scratch
+end
+
 """
 solve ∂f/∂t + v(z,t)⋅∂f/∂z + dvpa/dt ⋅ ∂f/∂vpa= 0
 define approximate characteristic velocity
@@ -1488,13 +1504,13 @@ df/dt + δv⋅∂f/∂z = 0, with δv(z,t)=v(z,t)-v₀(z)
 for prudent choice of v₀, expect δv≪v so that explicit
 time integrator can be used without severe CFL condition
 """
-function time_advance!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa,
-                       vperp, gyrophase, z, r, moments, fields, spectral_objects,
-                       advect_objects, composition, collisions, geometry, gyroavs,
-                       boundary_distributions, external_source_settings, num_diss_params,
-                       nl_solver_params, advance, advance_implicit, fp_arrays,
-                       scratch_dummy, manufactured_source_list, ascii_io, io_moments,
-                       io_dfns)
+function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz,
+                       vr, vzeta, vpa, vperp, gyrophase, z, r, moments, fields,
+                       spectral_objects, advect_objects, composition, collisions,
+                       geometry, gyroavs, boundary_distributions,
+                       external_source_settings, num_diss_params, nl_solver_params,
+                       advance, advance_implicit, fp_arrays, scratch_dummy,
+                       manufactured_source_list, ascii_io, io_moments, io_dfns)
 
     @debug_detect_redundant_block_synchronize begin
         # Only want to check for redundant _block_synchronize() calls during the
@@ -1563,22 +1579,23 @@ function time_advance!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzet
         
         if t_params.split_operators
             # MRH NOT SUPPORTED
-            time_advance_split_operators!(pdf, scratch, scratch_implicit, t, t_params,
-                                          vpa, z, vpa_spectral, z_spectral, moments,
-                                          fields, vpa_advect, z_advect, composition,
-                                          collisions, external_source_settings,
-                                          num_diss_params, nl_solver_params, advance,
-                                          advance_implicit, t_params.step_counter[])
+            time_advance_split_operators!(pdf, scratch, scratch_implicit,
+                                          scratch_electron, t, t_params, vpa, z,
+                                          vpa_spectral, z_spectral, moments, fields,
+                                          vpa_advect, z_advect, composition, collisions,
+                                          external_source_settings, num_diss_params,
+                                          nl_solver_params, advance, advance_implicit,
+                                          t_params.step_counter[])
         else
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vz,
-                                       vr, vzeta, vpa, vperp, gyrophase, z, r, moments,
-                                       fields, spectral_objects, advect_objects,
-                                       composition, collisions, geometry, gyroavs,
-                                       boundary_distributions, external_source_settings,
-                                       num_diss_params, nl_solver_params, advance,
-                                       advance_implicit, fp_arrays, scratch_dummy,
-                                       manufactured_source_list, diagnostic_checks,
-                                       t_params.step_counter[])
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                                       t, t_params, vz, vr, vzeta, vpa, vperp, gyrophase,
+                                       z, r, moments, fields, spectral_objects,
+                                       advect_objects, composition, collisions, geometry,
+                                       gyroavs, boundary_distributions,
+                                       external_source_settings, num_diss_params,
+                                       nl_solver_params, advance, advance_implicit,
+                                       fp_arrays, scratch_dummy, manufactured_source_list,
+                                       diagnostic_checks, t_params.step_counter[])
         end
         # update the time
         t += t_params.previous_dt[]
@@ -1814,11 +1831,12 @@ end
 
 """
 """
-function time_advance_split_operators!(pdf, scratch, scratch_implicit, t, t_params, vpa,
-                                       z, vpa_spectral, z_spectral, moments, fields,
-                                       vpa_advect, z_advect, composition, collisions,
-                                       external_source_settings, num_diss_params,
-                                       nl_solver_params, advance, advance_implicit, istep)
+function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_electron,
+                                       t, t_params, vpa, z, vpa_spectral, z_spectral,
+                                       moments, fields, vpa_advect, z_advect, composition,
+                                       collisions, external_source_settings,
+                                       num_diss_params, nl_solver_params, advance,
+                                       advance_implicit, istep)
 
     # define some abbreviated variables for tidiness
     n_ion_species = composition.n_ion_species
@@ -1831,178 +1849,186 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, t, t_para
         # advance the operator-split 1D advection equation in vpa
         # vpa-advection only applies for ion species
         advance.vpa_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-            vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-            composition, collisions, external_source_settings, num_diss_params,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+            t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
+            z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
         advance.vpa_advection = false
         # z_advection! advances the operator-split 1D advection equation in z
         # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-            vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-            composition, collisions, external_source_settings, num_diss_params,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+            t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
+            z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
         advance.z_advection = false
         # account for charge exchange collisions between ions and neutrals
         if composition.n_neutral_species > 0
             if collisions.charge_exchange > 0.0
                 advance.ion_cx_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                    vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    moments, fields, vpa_advect, z_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.ion_cx_collisions = false
                 advance.neutral_cx_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                    vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    moments, fields, vpa_advect, z_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.neutral_cx_collisions = false
             end
             if collisions.ionization > 0.0
                 advance.ion_ionization_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, z, vpa,
-                    z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    moments, fields, z_advect, vpa_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.ion_ionization_collisions = false
                 advance.neutral_ionization_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, z, vpa,
-                    z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    moments, fields, z_advect, vpa_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.neutral_ionization_collisions = false
             end
         end
         if collisions.krook_collision_frequency_prefactor  > 0.0
             advance.krook_collisions_ii = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, z, vpa,
-                z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
-                z_SL, vpa_SL, composition, collisions, sources, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, z, vpa, z_spectral, vpa_spectral, moments, fields, z_advect,
+                vpa_advect, z_SL, vpa_SL, composition, collisions, sources,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.krook_collisions_ii = false
         end
         # and add the source terms associated with redefining g = pdf/density or pdf*vth/density
         # to the kinetic equation
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
             advance.source_terms = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.source_terms = false
         end
         # use the continuity equation to update the density
         if moments.evolve_density
             advance.continuity = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.continuity = false
         end
         # use force balance to update the parallel flow
         if moments.evolve_upar
             advance.force_balance = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.force_balance = false
         end
         # use the energy equation to update the parallel pressure
         if moments.evolve_ppar
             advance.energy = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.energy = false
         end
     else
         # use the energy equation to update the parallel pressure
         if moments.evolve_ppar
             advance.energy = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.energy = false
         end
         # use force balance to update the parallel flow
         if moments.evolve_upar
             advance.force_balance = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.force_balance = false
         end
         # use the continuity equation to update the density
         if moments.evolve_density
             advance.continuity = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.continuity = false
         end
         # and add the source terms associated with redefining g = pdf/density or pdf*vth/density
         # to the kinetic equation
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
             advance.source_terms = true
-            time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                composition, collisions, external_source_settings, num_diss_params,
-                nl_solver_params, advance, advance_implicit, istep)
+            time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
+                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                vpa_advect, z_advect, composition, collisions, external_source_settings,
+                num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.source_terms = false
         end
         # account for charge exchange collisions between ions and neutrals
         if composition.n_neutral_species > 0
             if collisions.ionization > 0.0
                 advance.neutral_ionization = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, z, vpa,
-                    z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    moments, fields, z_advect, vpa_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.neutral_ionization = false
                 advance.ion_ionization = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, z, vpa,
-                    z_spectral, vpa_spectral, moments, fields, z_advect, vpa_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    moments, fields, z_advect, vpa_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.ion_ionization = false
             end
             if collisions.charge_exchange > 0.0
                 advance.neutral_cx_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                    vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    moments, fields, vpa_advect, z_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.neutral_cx_collisions = false
                 advance.ion_cx_collisions = true
-                time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-                    vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-                    composition, collisions, external_source_settings, num_diss_params,
-                    nl_solver_params, advance, advance_implicit, istep)
+                time_advance_no_splitting!(pdf, scratch, scratch_implicit,
+                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    moments, fields, vpa_advect, z_advect, composition, collisions,
+                    external_source_settings, num_diss_params, nl_solver_params, advance,
+                    advance_implicit, istep)
                 advance.ion_cx_collisions = false
             end
         end
         # z_advection! advances the operator-split 1D advection equation in z
         # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-            vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-            composition, collisions, external_source_settings, num_diss_params,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+            t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
+            z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
         advance.z_advection = false
         # advance the operator-split 1D advection equation in vpa
         # vpa-advection only applies for ion species
         advance.vpa_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vpa, z,
-            vpa_spectral, z_spectral, moments, fields, vpa_advect, z_advect,
-            composition, collisions, external_source_settings, num_diss_params,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+            t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
+            z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
         advance.vpa_advection = false
     end
@@ -2011,21 +2037,21 @@ end
 
 """
 """
-function time_advance_no_splitting!(pdf, scratch, scratch_implicit, t, t_params, vz, vr,
-                                    vzeta, vpa, vperp, gyrophase, z, r, moments, fields,
-                                    spectral_objects, advect_objects, composition,
-                                    collisions, geometry, gyroavs, boundary_distributions,
-                                    external_source_settings, num_diss_params,
-                                    nl_solver_params, advance, advance_implicit,
-                                    fp_arrays, scratch_dummy, manufactured_source_list,
-                                    diagnostic_checks, istep)
-
-    ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa, vperp,
-            gyrophase, z, r, moments, fields, spectral_objects, advect_objects,
-            composition, collisions, geometry, gyroavs, boundary_distributions,
-            external_source_settings, num_diss_params, nl_solver_params, advance,
-            advance_implicit, fp_arrays, scratch_dummy, manufactured_source_list,
-            diagnostic_checks, istep)
+function time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+                                    t_params, vz, vr, vzeta, vpa, vperp, gyrophase, z, r,
+                                    moments, fields, spectral_objects, advect_objects,
+                                    composition, collisions, geometry, gyroavs,
+                                    boundary_distributions, external_source_settings,
+                                    num_diss_params, nl_solver_params, advance,
+                                    advance_implicit, fp_arrays, scratch_dummy,
+                                    manufactured_source_list, diagnostic_checks, istep)
+
+    ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz, vr, vzeta,
+            vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
+            advect_objects, composition, collisions, geometry, gyroavs,
+            boundary_distributions, external_source_settings, num_diss_params,
+            nl_solver_params, advance, advance_implicit, fp_arrays, scratch_dummy,
+            manufactured_source_list, diagnostic_checks, istep)
 
     return nothing
 end
@@ -2069,10 +2095,10 @@ Apply boundary conditions and moment constraints to updated pdfs and calculate d
 moments and moment derivatives
 """
 function apply_all_bcs_constraints_update_moments!(
-        this_scratch, moments, fields, boundary_distributions, vz, vr, vzeta, vpa, vperp,
-        z, r, spectral_objects, advect_objects, composition, collisions, geometry, gyroavs,
-        num_diss_params, advance, scratch_dummy, diagnostic_moments; pdf_bc_constraints=true,
-        update_electron_pdf=true)
+        this_scratch, moments, fields, boundary_distributions, scratch_electron, vz, vr,
+        vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+        collisions, geometry, gyroavs, num_diss_params, advance, scratch_dummy,
+        diagnostic_moments; pdf_bc_constraints=true)
 
     begin_s_r_z_region()
 
@@ -2149,9 +2175,13 @@ function apply_all_bcs_constraints_update_moments!(
             moments.electron.ppar[iz,ir] = this_scratch.electron_ppar[iz,ir]
         end
 
-        if update_electron_pdf
-            update_electron_pdf!(this_scratch, pdf.electron.norm, moments, fields.phi, r,
-                                 z, vperp, vpa, z_spectral, vperp_spectral,
+        # When we do not need to apply bc's and constraints to the ion/neutral pdf
+        # (because this function is being called after a failed timestep, to reset to the
+        # state at the beginning of the step), we also do not need to update the
+        # electrons.
+        if pdf_bc_constraints
+            update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, fields.phi,
+                                 r, z, vperp, vpa, z_spectral, vperp_spectral,
                                  vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                  t_params.electron, collisions, composition,
                                  external_source_settings, num_diss_params,
@@ -2228,8 +2258,9 @@ function apply_all_bcs_constraints_update_moments!(
 end
 
 """
-    adaptive_timestep_update!(scratch, scratch_implicit, t, t_params, moments,
-                              fields, composition, collisions, geometry,
+    adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
+                              t, t_params, moments, fields,
+                              composition, collisions, geometry,
                               external_source_settings, spectral_objects,
                               advect_objects, gyroavs, num_diss_params, advance,
                               scratch_dummy, r, z, vperp, vpa, vzeta, vr, vz,
@@ -2238,12 +2269,13 @@ end
 Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
-function adaptive_timestep_update!(scratch, scratch_implicit, t, t_params, moments,
-                                   fields, boundary_distributions, composition,
-                                   collisions, geometry, external_source_settings,
-                                   spectral_objects, advect_objects, gyroavs,
-                                   num_diss_params, advance, scratch_dummy, r, z, vperp,
-                                   vpa, vzeta, vr, vz, success, nl_max_its_fraction)
+function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron, t,
+                                   t_params, moments, fields, boundary_distributions,
+                                   composition, collisions, geometry,
+                                   external_source_settings, spectral_objects,
+                                   advect_objects, gyroavs, num_diss_params, advance,
+                                   scratch_dummy, r, z, vperp, vpa, vzeta, vr, vz,
+                                   success, nl_max_its_fraction)
     #error_norm_method = "Linf"
     error_norm_method = "L2"
 
@@ -2510,10 +2542,10 @@ function adaptive_timestep_update!(scratch, scratch_implicit, t, t_params, momen
         # pdf These need to be re-calculated because `scratch[istage+1]` is now the
         # state at the beginning of the timestep, because the timestep failed
         apply_all_bcs_constraints_update_moments!(
-            scratch[t_params.n_rk_stages+1], moments, fields, nothing, vz, vr, vzeta,
-            vpa, vperp, z, r, spectral_objects, advect_objects, composition, collisions,
-            geometry, gyroavs, num_diss_params, advance, scratch_dummy, false;
-            pdf_bc_constraints=false, update_electron_pdf=false)
+            scratch[t_params.n_rk_stages+1], moments, fields, nothing, nothing, vz, vr,
+            vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+            collisions, geometry, gyroavs, num_diss_params, advance, scratch_dummy, false;
+            pdf_bc_constraints=false)
     end
 
     if composition.electron_physics == kinetic_electrons
@@ -2523,7 +2555,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, t, t_params, momen
             @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
                 pdf.electron.norm[ivpa,ivperp,iz,ir] =
                     pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir]
-                scratch[1].pdf_electron[ivpa,ivperp,iz,ir] =
+                scratch_electron[1].pdf_electron[ivpa,ivperp,iz,ir] =
                     pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir]
             end
         else
@@ -2542,11 +2574,11 @@ end
 
 """
 """
-function ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa, vperp,
-                 gyrophase, z, r, moments, fields, spectral_objects, advect_objects,
-                 composition, collisions, geometry, gyroavs, boundary_distributions,
-                 external_source_settings, num_diss_params, nl_solver_params, advance,
-                 advance_implicit, fp_arrays, scratch_dummy,
+function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz, vr,
+                 vzeta, vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
+                 advect_objects, composition, collisions, geometry, gyroavs,
+                 boundary_distributions, external_source_settings, num_diss_params,
+                 nl_solver_params, advance, advance_implicit, fp_arrays, scratch_dummy,
                  manufactured_source_list,diagnostic_checks, istep)
 
     begin_s_r_z_region()
@@ -2637,10 +2669,10 @@ function ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa
                 # which is used as input to the explicit part of the IMEX time step.
                 old_scratch = scratch_implicit[istage]
                 apply_all_bcs_constraints_update_moments!(
-                    scratch_implicit[istage], moments, fields, boundary_distributions, vz,
-                    vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects,
-                    composition, collisions, geometry, gyroavs, num_diss_params, advance,
-                    scratch_dummy, false)
+                    scratch_implicit[istage], moments, fields, boundary_distributions,
+                    scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
+                    advect_objects, composition, collisions, geometry, gyroavs,
+                    num_diss_params, advance, scratch_dummy, false)
             end
         else
             # Fully explicit method starts the forward-Euler step with the result from the
@@ -2672,10 +2704,10 @@ function ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa
                                 || t_params.implicit_coefficient_is_zero[istage+1])
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
         apply_all_bcs_constraints_update_moments!(
-            scratch[istage+1], moments, fields, boundary_distributions, vz, vr, vzeta,
-            vpa, vperp, z, r, spectral_objects, advect_objects, composition, collisions,
-            geometry, gyroavs, num_diss_params, advance, scratch_dummy,
-            diagnostic_moments; pdf_bc_constraints=apply_bc_constraints)
+            scratch[istage+1], moments, fields, boundary_distributions, scratch_electron,
+            vz, vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects,
+            composition, collisions, geometry, gyroavs, num_diss_params, advance,
+            scratch_dummy, diagnostic_moments; pdf_bc_constraints=apply_bc_constraints)
     end
 
     if t_params.adaptive
@@ -2687,12 +2719,13 @@ function ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa
                         nl_max_its_fraction)
             end
         end
-        adaptive_timestep_update!(scratch, scratch_implicit, t, t_params, moments, fields,
+        adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
+                                  t, t_params, pdf, moments, fields,
                                   boundary_distributions, composition, collisions,
                                   geometry, external_source_settings, spectral_objects,
-                                  advect_objects, gyroavs, num_diss_params, advance,
-                                  scratch_dummy, r, z, vperp, vpa, vzeta, vr, vz, success,
-                                  nl_max_its_fraction)
+                                  advect_objects, gyroavs, num_diss_params,
+                                  advance, scratch_dummy, r, z, vperp, vpa,
+                                  vzeta, vr, vz, success, nl_max_its_fraction)
     elseif !success
         error("Implicit part of timestep failed")
     end
@@ -3253,20 +3286,19 @@ function implicit_ion_advance!(fvec_out, fvec_in, pdf, fields, moments, advect_o
         # scratch_pdf struct containing the array passed as f_new
         new_scratch = scratch_pdf(f_new, fvec_out.density, fvec_out.upar, fvec_out.ppar,
                                   fvec_out.pperp, fvec_out.temp_z_s,
-                                  fvec_out.pdf_electron, fvec_out.electron_density,
-                                  fvec_out.electron_upar, fvec_out.electron_ppar,
-                                  fvec_out.electron_pperp, fvec_out.electron_temp,
-                                  fvec_out.pdf_neutral, fvec_out.density_neutral,
-                                  fvec_out.uz_neutral,
+                                  fvec_out.electron_density, fvec_out.electron_upar,
+                                  fvec_out.electron_ppar, fvec_out.electron_pperp,
+                                  fvec_out.electron_temp, fvec_out.pdf_neutral,
+                                  fvec_out.density_neutral, fvec_out.uz_neutral,
                                   fvec_out.pz_neutral)
         # scratch_pdf struct containing the array passed as residual
         residual_scratch = scratch_pdf(residual, fvec_out.density, fvec_out.upar,
                                        fvec_out.ppar, fvec_out.pperp, fvec_out.temp_z_s,
-                                       fvec_out.pdf_electron, fvec_out.electron_density,
-                                       fvec_out.electron_upar, fvec_out.electron_ppar,
-                                       fvec_out.electron_pperp, fvec_out.electron_temp,
-                                       fvec_out.pdf_neutral, fvec_out.density_neutral,
-                                       fvec_out.uz_neutral, fvec_out.pz_neutral)
+                                       fvec_out.electron_density, fvec_out.electron_upar,
+                                       fvec_out.electron_ppar, fvec_out.electron_pperp,
+                                       fvec_out.electron_temp, fvec_out.pdf_neutral,
+                                       fvec_out.density_neutral, fvec_out.uz_neutral,
+                                       fvec_out.pz_neutral)
 
         # Ensure moments are consistent with f_new
         update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,

From 881a7489727da411df518871957b6a3e446565ca Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 30 May 2024 14:14:14 +0100
Subject: [PATCH 263/428] Fix merge of 'imex' and 'electrons'

---
 .../src/electron_fluid_equations.jl           | 10 ++--
 moment_kinetics/src/time_advance.jl           | 58 +++++++++++--------
 2 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index dadef51ec..108145a02 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -113,7 +113,8 @@ function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens
     return nothing
 end
 
-function calculate_electron_moments!(scratch, moments, composition, collisions, r, z, vpa)
+function calculate_electron_moments!(scratch, pdf, moments, composition, collisions, r, z,
+                                     vpa)
     calculate_electron_density!(scratch.electron_density, moments.electron.dens_updated,
                                 scratch.density)
     calculate_electron_upar_from_charge_conservation!(
@@ -130,10 +131,9 @@ function calculate_electron_moments!(scratch, moments, composition, collisions,
     end
     update_electron_vth_temperature!(moments, scratch.electron_ppar,
                                      scratch.electron_density, composition)
-    calculate_electron_qpar!(moments.electron, scratch.pdf_electron,
-                             scratch.electron_ppar, scratch.electron_upar, scratch.upar,
-                             collisions.nu_ei, composition.me_over_mi,
-                             composition.electron_physics, vpa)
+    calculate_electron_qpar!(moments.electron, pdf.electron, scratch.electron_ppar,
+                             scratch.electron_upar, scratch.upar, collisions.nu_ei,
+                             composition.me_over_mi, composition.electron_physics, vpa)
     if composition.electron_physics == braginskii_fluid
         electron_fluid_qpar_boundary_condition!(moments.electron, z)
     end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index f341b6078..7fbcb8112 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2095,10 +2095,10 @@ Apply boundary conditions and moment constraints to updated pdfs and calculate d
 moments and moment derivatives
 """
 function apply_all_bcs_constraints_update_moments!(
-        this_scratch, moments, fields, boundary_distributions, scratch_electron, vz, vr,
-        vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
-        collisions, geometry, gyroavs, num_diss_params, advance, scratch_dummy,
-        diagnostic_moments; pdf_bc_constraints=true)
+        this_scratch, pdf, moments, fields, boundary_distributions, scratch_electron, vz,
+        vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+        collisions, geometry, gyroavs, external_source_settings, num_diss_params,
+        t_params, advance, scratch_dummy, diagnostic_moments; pdf_bc_constraints=true)
 
     begin_s_r_z_region()
 
@@ -2149,7 +2149,7 @@ function apply_all_bcs_constraints_update_moments!(
     calculate_ion_moment_derivatives!(moments, this_scratch, scratch_dummy, z, z_spectral,
                                       num_diss_params.ion.moment_dissipation_coefficient)
 
-    calculate_electron_moments!(this_scratch, moments, composition, collisions, r, z,
+    calculate_electron_moments!(this_scratch, pdf, moments, composition, collisions, r, z,
                                 vpa)
     calculate_electron_moment_derivatives!(moments, this_scratch, scratch_dummy, z,
                                            z_spectral,
@@ -2270,7 +2270,7 @@ Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
 function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron, t,
-                                   t_params, moments, fields, boundary_distributions,
+                                   t_params, pdf, moments, fields, boundary_distributions,
                                    composition, collisions, geometry,
                                    external_source_settings, spectral_objects,
                                    advect_objects, gyroavs, num_diss_params, advance,
@@ -2361,6 +2361,10 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         begin_s_r_z_region()
         rk_loworder_solution!(scratch, scratch_implicit, :ppar, t_params)
     end
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        begin_r_z_region()
+        rk_loworder_solution!(scratch, scratch_implicit, :electron_ppar, t_params)
+    end
     if n_neutral_species > 0
         begin_sn_r_z_vzeta_vr_region()
         rk_loworder_solution!(scratch, scratch_implicit, :pdf_neutral, t_params; neutrals=true)
@@ -2380,16 +2384,19 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
 
     # Apply boundary conditions and constraints
     apply_all_bcs_constraints_update_moments!(
-        scratch[2], moments, fields, boundary_distributions, vz, vr, vzeta,
-        vpa, vperp, z, r, spectral_objects, advect_objects, composition, geometry,
-        gyroavs, num_diss_params, advance, scratch_dummy, false)
+        scratch[2], pdf, moments, fields, boundary_distributions, scratch_electron, vz,
+        vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+        collisions, geometry, gyroavs, external_source_settings, num_diss_params,
+        t_params, advance, scratch_dummy, false)
 
     # Re-calculate moment derivatives in the `moments` struct, in case they were changed
     # by the previous call
     apply_all_bcs_constraints_update_moments!(
-        scratch[t_params.n_rk_stages+1], moments, fields, boundary_distributions, vz, vr,
-        vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition, geometry,
-        gyroavs, num_diss_params, advance, scratch_dummy, false; pdf_bc_constraints=false)
+        scratch[t_params.n_rk_stages+1], pdf, moments, fields, boundary_distributions,
+        scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
+        advect_objects, composition, collisions, geometry, gyroavs,
+        external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
+        false; pdf_bc_constraints=false)
 
     # Calculate the timstep error estimates
     ion_pdf_error = local_error_norm(scratch[2].pdf, scratch[t_params.n_rk_stages+1].pdf,
@@ -2436,7 +2443,6 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
 
     if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
         begin_r_z_region()
-        rk_error_variable!(scratch, scratch_implicit, :electron_ppar, t_params)
         electron_p_err = local_error_norm(scratch[2].electron_ppar,
                                           scratch[t_params.n_rk_stages+1].electron_ppar,
                                           t_params.rtol, t_params.atol;
@@ -2542,10 +2548,10 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         # pdf These need to be re-calculated because `scratch[istage+1]` is now the
         # state at the beginning of the timestep, because the timestep failed
         apply_all_bcs_constraints_update_moments!(
-            scratch[t_params.n_rk_stages+1], moments, fields, nothing, nothing, vz, vr,
-            vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
-            collisions, geometry, gyroavs, num_diss_params, advance, scratch_dummy, false;
-            pdf_bc_constraints=false)
+            scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
+            vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+            collisions, geometry, gyroavs, external_source_settings, num_diss_params,
+            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false)
     end
 
     if composition.electron_physics == kinetic_electrons
@@ -2669,10 +2675,11 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # which is used as input to the explicit part of the IMEX time step.
                 old_scratch = scratch_implicit[istage]
                 apply_all_bcs_constraints_update_moments!(
-                    scratch_implicit[istage], moments, fields, boundary_distributions,
-                    scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
-                    advect_objects, composition, collisions, geometry, gyroavs,
-                    num_diss_params, advance, scratch_dummy, false)
+                    scratch_implicit[istage], pdf, moments, fields,
+                    boundary_distributions, scratch_electron, vz, vr, vzeta, vpa, vperp,
+                    z, r, spectral_objects, advect_objects, composition, collisions,
+                    geometry, gyroavs, external_source_settings, num_diss_params,
+                    t_params, advance, scratch_dummy, false)
             end
         else
             # Fully explicit method starts the forward-Euler step with the result from the
@@ -2704,10 +2711,11 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                                 || t_params.implicit_coefficient_is_zero[istage+1])
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
         apply_all_bcs_constraints_update_moments!(
-            scratch[istage+1], moments, fields, boundary_distributions, scratch_electron,
-            vz, vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects,
-            composition, collisions, geometry, gyroavs, num_diss_params, advance,
-            scratch_dummy, diagnostic_moments; pdf_bc_constraints=apply_bc_constraints)
+            scratch[istage+1], pdf, moments, fields, boundary_distributions,
+            scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
+            advect_objects, composition, collisions, geometry, gyroavs,
+            external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
+            diagnostic_moments; pdf_bc_constraints=apply_bc_constraints)
     end
 
     if t_params.adaptive

From 11dd922330e064d4c80dd6ab88d621494005c9c6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 29 May 2024 15:18:06 +0100
Subject: [PATCH 264/428] Fix electron dt when restarting from a run without
 kinetic electrons

When the run being restarted from did not use kinetic electrons, there
is no electron_dt to reload, and the value defaults to 0.0. In this case
ignore the 'reloaded' value - instead keep the value that was set by the
input options.
---
 moment_kinetics/src/load_data.jl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index f1dd97dfe..b67983f47 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -1075,7 +1075,13 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
                                     restart_evolve_density, restart_evolve_upar,
                                     restart_evolve_ppar)
 
-            t_params.dt[] = load_slice(dynamic, "electron_dt", time_index)
+            new_dt = load_slice(dynamic, "electron_dt", time_index)
+            if new_dt > 0.0
+                # if the reloaded electron_dt was 0.0, then the previous run would not
+                # have been using kinetic electrons, so we only use the value if it is
+                # positive
+                t_params.dt[] = new_dt
+            end
             t_params.previous_dt[] = t_params.dt[]
             t_params.dt_before_output[] = t_params.dt[]
             t_params.dt_before_last_fail[] =

From d8cfa262509abd88aea31472e63b58a223fe7b0b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 29 May 2024 13:02:58 +0100
Subject: [PATCH 265/428] New method to apply boundary conditions with
 constraints to electron pdf

New method defined so that once the boundary conditions (and
constraints) are applied once, further calls to the function would leave
the pdf unchanged (so that the pdf would not drift in the limit that the
timestep goes to zero).
---
 .../src/electron_kinetic_equation.jl          | 813 +++++++-----------
 moment_kinetics/src/initial_conditions.jl     |   6 +-
 moment_kinetics/src/time_advance.jl           |   8 +
 3 files changed, 312 insertions(+), 515 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index e9269579b..bb01ec011 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -309,7 +309,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                                         vpa, vperp_spectral, vpa_spectral,
                                                         vpa_advect, moments,
                                                         num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                        composition.me_over_mi)
+                                                        composition.me_over_mi, 0.1 * t_params.rtol)
 
             begin_r_z_region()
             A = moments.electron.constraints_A_coefficient
@@ -625,7 +625,7 @@ end
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, moments, vpa_diffusion,
-                                                     me_over_mi)
+                                                     me_over_mi, newton_tol)
     # Enforce velocity-space boundary conditions
     if vpa.n > 1
         begin_r_z_vperp_region()
@@ -652,28 +652,154 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
     begin_r_region()
 
-    # pdf_adjustment_option determines the velocity-dependent pre-factor for the
-    # corrections to the pdf needed to ensure moment constraints are satisfied
-    #pdf_adjustment_option = "vpa4"
-    #pdf_adjustment_option = "vpa4_gaussian"
-    pdf_adjustment_option = "no1st_vpa2"
-
-    cutoff_step_width = 0.1
-
-    # wpa_values will be used to store the wpa = (vpa - upar)/vthe values corresponding to a vpa grid symmetric about vpa=0
-    #wpa_values = vpa.scratch
-    # interpolated_pdf will be used to store the pdf interpolated onto the vpa-symmetric grid
-    #interpolated_pdf = vpa.scratch2
+    newton_max_its = 100
     reversed_pdf = vpa.scratch
 
-    # ivpa_zero is the index of the interpolated_pdf corresponding to vpa = 0
-    #ivpa_zero = (vpa.n+1)÷2
+    function get_residual_and_coefficients_for_bc(a1, a1prime, a2, a2prime, b1, b1prime,
+                                                  c1, c1prime, c2, c2prime, d1, d1prime,
+                                                  e1, e1prime, e2, e2prime, u_over_vt)
+        alpha = a1 + 2.0 * a2
+        alphaprime = a1prime + 2.0 * a2prime
+        beta = c1 + 2.0 * c2
+        betaprime = c1prime + 2.0 * c2prime
+        gamma = u_over_vt^2 * alpha - 2.0 * u_over_vt * b1 + beta
+        gammaprime = u_over_vt^2 * alphaprime - 2.0 * u_over_vt * b1prime + betaprime
+        delta = u_over_vt^2 * beta - 2.0 * u_over_vt * d1 + e1 + 2.0 * e2
+        deltaprime = u_over_vt^2 * betaprime - 2.0 * u_over_vt * d1prime + e1prime + 2.0 * e2prime
+
+        A = (0.5 * beta - delta) / (beta * gamma - alpha * delta)
+        Aprime = (0.5 * betaprime - deltaprime
+                  - (0.5 * beta - delta) * (gamma * betaprime + beta * gammaprime - delta * alphaprime - alpha * deltaprime)
+                    / (beta * gamma - alpha * delta)
+                 ) / (beta * gamma - alpha * delta)
+        C = (1.0 - alpha * A) / beta
+        Cprime = -(A * alphaprime + alpha * Aprime) / beta - (1.0 - alpha * A) * betaprime / beta^2
+
+        epsilon = A * b1 + C * d1 - u_over_vt
+        epsilonprime = b1 * Aprime + A * b1prime + d1 * Cprime + C * d1prime
+
+        return epsilon, epsilonprime, A, C
+    end
 
     if z.irank == 0
+        if z.bc != "wall"
+            error("Options other than wall bc not implemented yet for electrons")
+        end
         @loop_r ir begin
-            # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-            #@. wpa_values = vpa.grid #- upar[1,ir] / vthe[1,ir]
-            #wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid - upar[1,ir] / vthe[1,ir]
+            # Impose sheath-edge boundary condition, while also imposing moment
+            # constraints and determining the cut-off velocity (and therefore the sheath
+            # potential).
+
+            # Delete the upar contribution here if ignoring the 'upar shift'
+            vpa_unnorm = @. vpa.scratch2 = vthe[1,ir] * vpa.grid + upar[1,ir]
+
+            u_over_vt = upar[1,ir] / vthe[1,ir]
+
+            # Initial guess for cut-off velocity is result from previous RK stage (which
+            # might be the previous timestep if this is the first stage). Recalculate this
+            # value from phi.
+            vmax = sqrt(phi[1,ir] / me_over_mi)
+
+            # -vmax is between vmax_ind-1 and vmax_ind
+            vmax_ind = searchsortedfirst(vpa_unnorm, -vmax)
+
+            # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
+            # shift'
+            sigma = -u_over_vt
+
+            # sigma is between sigma_ind-1 and sigma_ind
+            sigma_ind = searchsortedfirst(vpa_unnorm, 0.0)
+
+            # sigma_fraction is the fraction of the distance between sigma_ind-1 and
+            # sigma_ind where sigma is.
+            sigma_fraction = (sigma - vpa_unnorm[sigma_ind-1]) / (vpa_unnorm[sigma_ind] - vpa_unnorm[sigma_ind-1])
+
+            # Per-grid-point contributions to moment integrals
+            # Note that we need to include the normalisation factor of 1/sqrt(pi) that
+            # would be factored in by integrate_over_vspace(). This will need to
+            # change/adapt when we support 2V as well as 1V.
+            density_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.wgts / sqrt(pi)
+            flow_integral_pieces = @views @. vpa.scratch4 = density_integral_pieces * vpa_unnorm / vthe[1,ir]
+            energy_integral_pieces = @views @. vpa.scratch5 = flow_integral_pieces * vpa_unnorm / vthe[1,ir]
+            cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[1,ir]
+            quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[1,ir]
+
+            function get_integrals_and_derivatives(vmax, vmax_ind)
+                # vmax_fraction is the fraction of the distance between vmax_ind-1 and
+                # vmax_ind where -vmax is.
+                vmax_fraction = (-vmax - vpa_unnorm[vmax_ind-1]) / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind-1])
+
+                function get_for_one_moment(integral_pieces, skip_part2=false)
+                    # Integral contribution from the cell containing vmax
+                    integral_vmax_cell = (0.5 * integral_pieces[vmax_ind-1] + 0.5 * integral_pieces[vmax_ind])
+
+                    part1 = sum(integral_pieces[1:vmax_ind-2])
+                    part1 += 0.5 * integral_pieces[vmax_ind-1] + vmax_fraction * integral_vmax_cell
+                    # part1prime is d(part1)/d(vmax)
+                    part1prime = -integral_vmax_cell / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind-1])
+
+                    if skip_part2
+                        part2 = nothing
+                        part2prime = nothing
+                    else
+                        # Integral contribution from the cell containing sigma
+                        integral_sigma_cell = (0.5 * integral_pieces[sigma_ind-1] + 0.5 * integral_pieces[sigma_ind])
+
+                        part2 = sum(integral_pieces[vmax_ind+1:sigma_ind-2])
+                        part2 += (1.0 - vmax_fraction) * integral_vmax_cell + 0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[sigma_ind-1] + sigma_fraction * integral_sigma_cell
+                        # part2prime is d(part2)/d(vmax)
+                        part2prime = -part1prime
+                    end
+
+                    return part1, part1prime, part2, part2prime
+                end
+                a1, a1prime, a2, a2prime = get_for_one_moment(density_integral_pieces)
+                b1, b1prime, _, _ = get_for_one_moment(flow_integral_pieces, true)
+                c1, c1prime, c2, c2prime = get_for_one_moment(energy_integral_pieces)
+                d1, d1prime, _, _ = get_for_one_moment(cubic_integral_pieces, true)
+                e1, e1prime, e2, e2prime = get_for_one_moment(quartic_integral_pieces)
+
+                return get_residual_and_coefficients_for_bc(a1, a1prime, a2, a2prime, b1,
+                                                            b1prime, c1, c1prime, c2,
+                                                            c2prime, d1, d1prime, e1,
+                                                            e1prime, e2, e2prime,
+                                                            u_over_vt)
+            end
+
+            counter = 1
+            A = 1.0
+            C = 0.0
+            # Always do at least one update of vmax
+            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+            while true
+                # Newton iteration update. Note that primes denote derivatives with
+                # respect to vmax
+                delta_v = - epsilon / epsilonprime
+
+                # Prevent the step size from getting too big, to make Newton iteration
+                # more robust.
+                delta_v = min(delta_v, 0.1 * vthe[1,ir])
+                delta_v = max(delta_v, -0.1 * vthe[1,ir])
+
+                vmax = vmax + delta_v
+                vmax_ind = searchsortedfirst(vpa_unnorm, -vmax)
+
+                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+
+                if abs(epsilon) < newton_tol * abs(u_over_vt)
+                    break
+                end
+
+                if counter ≥ newton_max_its
+                    error("Newton iteration for electron lower-z boundary failed to "
+                          * "converge after $counter iterations")
+                end
+                counter += 1
+            end
+
+            # Adjust pdf so that after reflecting and cutting off tail, it will obey the
+            # constraints.
+            @. pdf[:,1,1,ir] *= A + C * vpa_unnorm^2 / vthe[1,ir]^2
 
             # Want to construct the w-grid corresponding to -vpa.
             #   wpa(vpa) = (vpa - upar)/vth
@@ -686,220 +812,29 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
             # Need to reverse vpa.grid because the grid passed as the second argument of
             # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[1,ir] / vthe[1,ir]
-            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[1,ir] / vthe[1,ir]
-            reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
+            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
+            reverse!(reversed_wpa_of_minus_vpa)
+
             # interpolate the pdf onto this grid
             #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
             @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
             reverse!(reversed_pdf)
-            # fill in the vpa > 0 points of the pdf by mirroring the vpa < 0 points
-            #@. interpolated_pdf[ivpa_zero+1:end] = interpolated_pdf[ivpa_zero-1:-1:1]
-            # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-            #@. wpa_values = vpa.grid #+ upar[1,ir] / vthe[1,ir]
-            # interpolate back onto the original wpa grid
-            #@views interpolate_to_grid_1d!(pdf[:,1,1,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-            # construct wpa * pdf
-            #@. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid
-            # calculate the first moment of the normalised pdf
-            #first_vspace_moment = 0.0
-            #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # the initial max vpa index is the largest one possible; this will be reduced if the first moment is positive
-            ivpa = 0
-            ivpa_max = vpa.n + 1
-            # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-            # if the first moment is positive, then the cutoff speed needs to be reduced
-            upar_over_vth = upar[1,ir] / vthe[1,ir]
-            #println("upar=", upar[1,ir], " vthe=", vthe[1,ir])
-            #println("$first_vspace_moment, u/vth=$upar_over_vth")
-            vpa_unnorm = @. vpa.scratch3 = vthe[1,ir] * vpa.grid + upar[1,ir]
-            upar_integral = 0.0
-            #while first_vspace_moment > upar_over_vth # > 0.0
-            #    # zero out the pdf at the current cutoff velocity
-            #    pdf[ivpa_max,1,1,ir] = 0.0
-            #    # update wpa * pdf
-            #    vpa.scratch3[ivpa_max] = 0.0
-            #    # calculate the updated first moment of the normalised pdf
-            #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            #    #println("truncating pdf $ivpa_max, $first_vspace_moment, u/vth=$upar_over_vth")
-            #    if first_vspace_moment > upar_over_vth #0.0
-            #        ivpa_max -= 1
-            #    end
-            #end
-            upar0 = upar[1,ir]
-            if upar0 >= 0.0
-                error("In lower-z boundary condition, upar0=$upar0 has the wrong sign.")
-            elseif isnan(upar0)
-                error("In lower-z boundary condition, upar0=$upar0.")
-            end
-            #println("before pdf left ", pdf[:,1,1,ir])
-            while upar_integral > upar0 && ivpa_max > 1
-                ivpa += 1
-                ivpa_max -= 1
-                # zero out the reversed pdf at the current cutoff velocity
-                #reversed_pdf[ivpa_max] = 0.0
-                # calculate the updated first moment of the normalised pdf
-                upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,1,ir] * vpa.wgts[ivpa]
-                #println("left ", ivpa, " ", upar_integral, " ", upar0)
-            end
-            if ivpa ≥ vpa.n
-                error("In lower-z boundary condition, upar_integral failed to reach upar0")
-            elseif ivpa ≤ 1
-                error("In lower-z boundary condition, ivpa=$ivpa ≤ 1")
-            end
-            integral_excess = upar_integral - upar0
-            fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,1,ir]
-            #println("fraction_of_pdf=", fraction_of_pdf)
-
-            # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-            # integral from the point at ivpa is required to make upar_integral<upar0) the
-            # cut-off velocity is half way between ivpa and ivpa+1, while when
-            # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-            # ivpa is required to make upar_integral<upar0) the cut-off is half way between
-            # ivpa-1 and ivpa.
-            vmax = 0.5 * (vpa_unnorm[ivpa] + vpa_unnorm[ivpa+1]) +
-                   0.5 * fraction_of_pdf*(vpa_unnorm[ivpa-1] - vpa_unnorm[ivpa+1])
-
-            #println("vmax=$vmax, v-no-interp=", vpa_unnorm[ivpa])
-            #wmax = (-vmax - upar[1,ir]) / vthe[1,ir]
-            #println("wmax=$wmax, w-no-interp", (vpa_unnorm[ivpa] - upar0)/vthe[1,ir])
-            #@loop_vpa ivpa begin
-            #    reversed_pdf[ivpa] *= 0.5*(1.0 - tanh((vpa.grid[ivpa] - wmax) / cutoff_step_width))
-            #end
-            reversed_pdf[ivpa_max+1:end] .= 0
-            reversed_pdf[ivpa_max] *= fraction_of_pdf
-            #println("first_vspace_moment=$first_vspace_moment, ivpa_max=$ivpa_max")
-            #println("done first cutoff loop")
+
+            ivpa_max = searchsortedlast(vpa_unnorm, vmax)
+            reversed_pdf[ivpa_max+1:end] .= 0.0
+            # vmax_fraction is the fraction of the distance between ivpa_max and
+            # ivpa_max+1 where vmax is.
+            vmax_fraction = (vmax - vpa_unnorm[ivpa_max]) / (vpa_unnorm[ivpa_max+1] - vpa_unnorm[ivpa_max])
+            reversed_pdf[ivpa_max] *= vmax_fraction
+
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-            #phi[1,ir] = me_over_mi * vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
             phi[1,ir] = me_over_mi * vmax^2
-            iv0 = findfirst(x -> x>0.0, vpa_unnorm)
-            if iv0 === nothing
-                error("All unnormalised vpa values at lower-z sheath entrance are negative. "
-                      * "Cannot apply electron boundary condition.")
-            end
-            pdf[iv0:end,1,1,ir] .= reversed_pdf[iv0:end]
-            #println("check reversed change ", reversed_pdf[iv0:end])
-            #println("reversed_pdf ", reversed_pdf)
-            #println("after pdf left ", pdf[:,1,1,ir])
-            # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-            # of the modified pdf are 1, 0 and 1/2 respectively
-            # will need vpa / vthe = wpa + upar/vthe
-            @. vpa.scratch2 = vpa.grid + upar[1,ir] / vthe[1,ir]
-            # first need to calculate int dwpa pdf = zeroth_moment
-            zeroth_moment = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
-            # calculate int dwpa wpa^2 * pdf = wpa2_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.grid^2
-            wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2
-            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,1,ir]
-            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,1,ir]
-            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            
-            normalisation_constant_A = 0.0
-            normalisation_constant_B = 0.0
-            normalisation_constant_C = 0.0
-            if pdf_adjustment_option == "absvpa"
-                # calculate int dwpa |vpa| * pdf = absvpa_moment
-                @. vpa.scratch3 = pdf[:,1,1,ir] * abs(vpa.scratch2)
-                absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-                # calculate the 'B' normalisation constant
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                    (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                    + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-                # calculate the 'A' normalisation constant
-                normalisation_constant_A = (1 + normalisation_constant_B * 
-                    (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-                # calculate the 'C' normalisation constant
-                normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-                # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-                @. pdf[:,1,1,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                    + vpa.scratch2^2 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "vpa4"
-                # calculate int dwpa vpa^4 * pdf = vpa4_moment
-                @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,1,ir]
-                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-                normalisation_constant_A = (1 + normalisation_constant_B 
-                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-                @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                    + vpa.scratch2^4 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "vpa4_gaussian"
-                afac = 0.1
-                bfac = 0.2
-                # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-                @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
-                vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-                @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,1,ir]
-                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-                normalisation_constant_A = (1 + normalisation_constant_B 
-                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-                #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                #                           (vpa2_wpa2_moment
-                #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-                #normalisation_constant_A = (1 - normalisation_constant_B 
-                #                                * vpa2_moment) / zeroth_moment
-                #normalisation_constant_C = 0.0
-                @. pdf[:,1,1,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                    + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "no1st_vpa2"
-                normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                           (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-                normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
-                @. pdf[:,1,1,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
-            else
-                println("pdf_adjustment_option not recognised")
-                stop()
-            end
-
-            moments.electron.constraints_A_coefficient[1,ir] = normalisation_constant_A
-            moments.electron.constraints_B_coefficient[1,ir] = normalisation_constant_B
-            moments.electron.constraints_C_coefficient[1,ir] = normalisation_constant_C
+            pdf[sigma_ind:end,1,1,ir] .= reversed_pdf[sigma_ind:end]
 
-            # smooth the pdf at the boundary
-            #for ivpa ∈ 2:ivpa_max-1
-            #    pdf[ivpa,1,1,ir] = (pdf[ivpa-1,1,1,ir] + pdf[ivpa+1,1,1,ir]) / 2.0
-            #end
+            moments.electron.constraints_A_coefficient[1,ir] = A
+            moments.electron.constraints_B_coefficient[1,ir] = 0.0
+            moments.electron.constraints_C_coefficient[1,ir] = C
         end
     end
 
@@ -911,308 +846,162 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
     # the electrostatic potential at the boundary, which determines the critical speed, is unknown a priori;
     # use the constraint that the first moment of the normalised pdf be zero to choose the potential.
     
-    # io_pdf_stages = open("pdf_stages.txt", "w")
-    # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-    # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-    # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid^2
-    # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-    # @loop_vpa ivpa begin
-    #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,1], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-    #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 0)
-    # end
-    # println(io_pdf_stages,"")
-
     if z.irank == z.nrank - 1
+        if z.bc != "wall"
+            error("Options other than wall bc not implemented yet for electrons")
+        end
         @loop_r ir begin
-            # construct a grid of wpa = (vpa - upar)/vthe values corresponding to a vpa-symmetric grid
-            #@. wpa_values = vpa.grid # - upar[end,ir] / vthe[end,ir]
-            # Need to reverse vpa.grid because the grid passed as the second argument of
-            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+            # Impose sheath-edge boundary condition, while also imposing moment
+            # constraints and determining the cut-off velocity (and therefore the sheath
+            # potential).
+
+            # Delete the upar contribution here if ignoring the 'upar shift'
+            vpa_unnorm = @. vpa.scratch2 = vthe[end,ir] * vpa.grid + upar[end,ir]
+
+            u_over_vt = upar[end,ir] / vthe[end,ir]
+
+            # Initial guess for cut-off velocity is result from previous RK stage (which
+            # might be the previous timestep if this is the first stage). Recalculate this
+            # value from phi.
+            vmax = sqrt(phi[end,ir] / me_over_mi)
+
+            # -vmax is between vmax_ind and vmax_ind+1
+            vmax_ind = searchsortedlast(vpa_unnorm, vmax)
+
+            # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
+            # shift'
+            sigma = -u_over_vt
+
+            # sigma is between sigma_ind and sigma_ind+1
+            sigma_ind = searchsortedlast(vpa_unnorm, 0.0)
+
+            # sigma_fraction is the fraction of the distance between sigma_ind+1 and
+            # sigma_ind where sigma is.
+            sigma_fraction = (sigma - vpa_unnorm[sigma_ind+1]) / (vpa_unnorm[sigma_ind] - vpa_unnorm[sigma_ind+1])
+
+            # Per-grid-point contributions to moment integrals
+            # Note that we need to include the normalisation factor of 1/sqrt(pi) that
+            # would be factored in by integrate_over_vspace(). This will need to
+            # change/adapt when we support 2V as well as 1V.
+            density_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.wgts / sqrt(pi)
+            flow_integral_pieces = @views @. vpa.scratch4 = density_integral_pieces * vpa_unnorm / vthe[end,ir]
+            energy_integral_pieces = @views @. vpa.scratch5 = flow_integral_pieces * vpa_unnorm / vthe[end,ir]
+            cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[end,ir]
+            quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[end,ir]
+
+            function get_integrals_and_derivatives(vmax, vmax_ind)
+                # vmax_fraction is the fraction of the distance between vmax_ind and
+                # vmax_ind+1 where vmax is.
+                vmax_fraction = (vmax - vpa_unnorm[vmax_ind+1]) / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind+1])
+
+                function get_for_one_moment(integral_pieces, skip_part2=false)
+                    # Integral contribution from the cell containing vmax
+                    integral_vmax_cell = (0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[vmax_ind+1])
+
+                    part1 = sum(integral_pieces[vmax_ind+2:end])
+                    part1 += 0.5 * integral_pieces[vmax_ind+1] + vmax_fraction * integral_vmax_cell
+                    # part1prime is d(part1)/d(vmax)
+                    part1prime = integral_vmax_cell / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind+1])
+
+                    if skip_part2
+                        part2 = nothing
+                        part2prime = nothing
+                    else
+                        # Integral contribution from the cell containing sigma
+                        integral_sigma_cell = (0.5 * integral_pieces[sigma_ind] + 0.5 * integral_pieces[sigma_ind+1])
+
+                        part2 = sum(integral_pieces[sigma_ind+2:vmax_ind-1])
+                        part2 += (1.0 - vmax_fraction) * integral_vmax_cell + 0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[sigma_ind+1] + sigma_fraction * integral_sigma_cell
+                        # part2prime is d(part2)/d(vmax)
+                        part2prime = -part1prime
+                    end
 
+                    return part1, part1prime, part2, part2prime
+                end
+                a1, a1prime, a2, a2prime = get_for_one_moment(density_integral_pieces)
+                b1, b1prime, _, _ = get_for_one_moment(flow_integral_pieces, true)
+                c1, c1prime, c2, c2prime = get_for_one_moment(energy_integral_pieces)
+                d1, d1prime, _, _ = get_for_one_moment(cubic_integral_pieces, true)
+                e1, e1prime, e2, e2prime = get_for_one_moment(quartic_integral_pieces)
+
+                return get_residual_and_coefficients_for_bc(a1, a1prime, a2, a2prime, b1,
+                                                            b1prime, c1, c1prime, c2,
+                                                            c2prime, d1, d1prime, e1,
+                                                            e1prime, e2, e2prime,
+                                                            u_over_vt)
+            end
+
+            counter = 1
+            # Always do at least one update of vmax
+            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+            while true
+                # Newton iteration update. Note that primes denote derivatives with
+                # respect to vmax
+                delta_v = - epsilon / epsilonprime
+
+                # Prevent the step size from getting too big, to make Newton iteration
+                # more robust.
+                delta_v = min(delta_v, 0.1 * vthe[1,ir])
+                delta_v = max(delta_v, -0.1 * vthe[1,ir])
+
+                vmax = vmax + delta_v
+                vmax_ind = searchsortedlast(vpa_unnorm, vmax)
+
+                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+
+                if abs(epsilon) < newton_tol * abs(u_over_vt)
+                    break
+                end
+
+                if counter ≥ newton_max_its
+                    error("Newton iteration for electron upper-z boundary failed to "
+                          * "converge after $counter iterations")
+                end
+                counter += 1
+            end
+
+            # Adjust pdf so that after reflecting and cutting off tail, it will obey the
+            # constraints.
+            @. pdf[:,1,end,ir] *= A + C * vpa_unnorm^2 / vthe[end,ir]^2
+
+            # Want to construct the w-grid corresponding to -vpa.
+            #   wpa(vpa) = (vpa - upar)/vth
+            #   ⇒ vpa = vth*wpa(vpa) + upar
+            #   wpa(-vpa) = (-vpa - upar)/vth
+            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+            #             = (-vth*wpa - 2*upar)/vth
+            #             = -wpa - 2*upar/vth
             # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
             #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
             # Need to reverse vpa.grid because the grid passed as the second argument of
             # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 2.0 * upar[end,ir] / vthe[end,ir]
-            #reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid) .- 1.5 * upar[end,ir] / vthe[end,ir]
-            reversed_wpa_of_minus_vpa = vpa.scratch2 .= .-reverse(vpa.grid)
+            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
+            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
+            reverse!(reversed_wpa_of_minus_vpa)
+
             # interpolate the pdf onto this grid
-            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,end,ir], vpa, vpa_spectral)
+            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
             @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
             reverse!(reversed_pdf)
-            # fill in the vpa < 0 points of the pdf by mirroring the vpa > 0 points
-            #@. interpolated_pdf[ivpa_zero-1:-1:1] = interpolated_pdf[ivpa_zero+1:end]
-            # construct a grid of vpa/vthe = wpa + upar/vthe values corresponding to the wpa-symmetric grid
-            #@. wpa_values = vpa.grid #+ upar[end,ir] / vthe[end,ir]
-            # interpolate back onto the original wpa grid
-            #@views interpolate_to_grid_1d!(pdf[:,1,end,ir], wpa_values, interpolated_pdf, vpa, vpa_spectral)
-
-            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @. vpa.scratch3 *= vpa.grid
-            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @loop_vpa ivpa begin
-            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 1)
-            # end
-            # println(io_pdf_stages,"")
-
-            # construct wpa * pdf
-            #@. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid
-            # calculate the first moment of the normalised pdf
-            #first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # the initial min vpa index is the smallest one possible; this will be increased if the first moment is negative
-            ivpa = vpa.n+1
-            ivpa_min = 0
-            # adjust the critical (cutoff) speed until the first moment is as close to zero as possible
-            # if the first moment is negative, then the magnitude of the cutoff speed needs to be reduced
-            upar_over_vth = upar[end,ir] / vthe[end,ir]
-            #println("$first_vspace_moment, u/vth=$upar_over_vth")
-            vpa_unnorm = @. vpa.scratch3 = vthe[end,ir] * vpa.grid + upar[end,ir]
-            upar_integral = 0.0
-            #while first_vspace_moment < upar_over_vth # < 0.0
-            #    # zero out the pdf at the current cutoff velocity
-            #    pdf[ivpa_min,1,end,ir] = 0.0
-            #    # update wpa * pdf
-            #    vpa.scratch3[ivpa_min] = 0.0
-            #    # calculate the updated first moment of the normalised pdf
-            #    first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            #    if first_vspace_moment < upar_over_vth
-            #        ivpa_min += 1
-            #    end
-            #end
-            upar_end = upar[end,ir]
-            if upar_end <= 0.0
-                error("In upper-z boundary condition, upar_end=$upar_end has the wrong sign.")
-            elseif isnan(upar_end)
-                error("In upper-z boundary condition, upar_end=$upar_end.")
-            end
-            #println("before pdf ", pdf[:,1,end,ir])
-            while upar_integral < upar_end && ivpa > 1
-                ivpa -= 1
-                ivpa_min += 1
-                # zero out the reversed pdf at the current cutoff velocity
-                #reversed_pdf[ivpa_min] = 0.0
-                # calculate the updated first moment of the normalised pdf
-                upar_integral += vpa_unnorm[ivpa] * pdf[ivpa,1,end,ir] * vpa.wgts[ivpa]
-                #println("right ", ivpa, " ", upar_integral, " ", upar_end)
-            end
-            if ivpa ≤ 1
-                error("In upper-z boundary condition, upar_integral failed to reach upar_end")
-            elseif ivpa ≥ vpa.n
-                error("In upper-z boundary condition, ivpa=$ivpa ≥ vpa.n=$(vpa.n)")
-            end
-            integral_excess = upar_integral - upar_end
-            fraction_of_pdf = integral_excess / (vpa_unnorm[ivpa] * vpa.wgts[ivpa]) / pdf[ivpa,1,end,ir]
-            #println("B fraction_of_pdf=", fraction_of_pdf)
-
-            # Define so that when fraction_of_pdf=1 (when all of the contribution to the
-            # integral from the point at ivpa is required to make upar_integral>upar_end) the
-            # cut-off velocity is half way between ivpa-1 and ivpa, while when
-            # fraction_of_pdf=0 (none of the contribution to the integral from the point at
-            # ivpa is required to make upar_integral>upar_end) the cut-off is half way between
-            # ivpa and ivpa+1.
-            vmin = 0.5 * (vpa_unnorm[ivpa-1] + vpa_unnorm[ivpa]) +
-                   0.5 * fraction_of_pdf*(vpa_unnorm[ivpa+1] - vpa_unnorm[ivpa-1])
-
-            #println("vmin=$vmin, v-no-interp=", vpa_unnorm[ivpa])
-            #wmin = (-vmin - upar[end,ir]) / vthe[end,ir]
-            #@loop_vpa ivpa begin
-            #    reversed_pdf[ivpa] *= 0.5*(1.0 + tanh((vpa.grid[ivpa] - wmin) / cutoff_step_width))
-            #end
-            reversed_pdf[1:ivpa_min-1] .= 0
-            reversed_pdf[ivpa_min] *= fraction_of_pdf
-            #println("done second cutoff loop")
-
-            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @. vpa.scratch3 *= vpa.grid
-            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @loop_vpa ivpa begin
-            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 2)
-            # end
-            # println(io_pdf_stages,"")
+
+            ivpa_min = searchsortedfirst(vpa_unnorm, -vmax)
+            reversed_pdf[1:ivpa_min-1] .= 0.0
+            # vmax_fraction is the fraction of the distance between ivpa_min and
+            # ivpa_min-1 where -vmax is.
+            vmax_fraction = (-vmax - vpa_unnorm[ivpa_min]) / (vpa_unnorm[ivpa_min-1] - vpa_unnorm[ivpa_min])
+            reversed_pdf[ivpa_min] *= vmax_fraction
 
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-            #phi[end,ir] = me_over_mi * vthe[end,ir]^2 * vpa.grid[ivpa_min]^2
-            phi[end,ir] = me_over_mi * vmin^2
-            iv0 = findlast(x -> x<0.0, vpa_unnorm)
-            if iv0 === nothing
-                error("All unnormalised vpa values at upper-z sheath entrance are positive. "
-                      * "Cannot apply electron boundary condition.")
-            end
-            pdf[1:iv0,1,end,ir] .= reversed_pdf[1:iv0]
-            #println("after pdf ", pdf[:,1,end,ir])
-            # obtain the normalisation constants needed to ensure the zeroth, first and second moments
-            # of the modified pdf are 1, 0 and 1/2 respectively
-            # will need vpa / vthe = wpa + upar/vthe
-            @. vpa.scratch2 = vpa.grid + upar[end,ir] / vthe[end,ir]
-            # first need to calculate int dwpa pdf = zeroth_moment
-            zeroth_moment = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
-            # calculate int dwpa wpa^2 * pdf = wpa2_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.grid^2
-            wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * pdf = vpa2_moment
-            @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2
-            vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa vpa^2 * wpa * pdf = vpa2_wpa_moment
-            @. vpa.scratch3 = vpa.grid * vpa.scratch2^2 * pdf[:,1,end,ir]
-            vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # calculate int dwpa wpa^2 * vpa^2 * pdf = vpa2_wpa2_moment
-            @. vpa.scratch3 = vpa.grid^2 * vpa.scratch2^2 * pdf[:,1,end,ir]
-            vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            
-            normalisation_constant_A = 0.0
-            normalisation_constant_B = 0.0
-            normalisation_constant_C = 0.0
-            if pdf_adjustment_option == "absvpa"
-                # calculate int dwpa |vpa| * pdf = absvpa_moment
-                @. vpa.scratch3 = pdf[:,1,end,ir] * abs(vpa.scratch2)
-                absvpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa |vpa| * wpa * pdf = absvpa_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                absvpa_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa |vpa| * wpa^2 * pdf = absvpa_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                absvpa_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-                # calculate the 'B' normalisation constant
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) / 
-                    (absvpa_wpa2_moment  - absvpa_wpa_moment * vpa2_wpa2_moment / vpa2_wpa_moment
-                    + wpa2_moment / zeroth_moment * (absvpa_wpa_moment * vpa2_moment / vpa2_wpa_moment - absvpa_moment))
-                # calculate the 'A' normalisation constant
-                normalisation_constant_A = (1 + normalisation_constant_B * 
-                    (vpa2_moment * absvpa_wpa_moment / vpa2_wpa_moment - absvpa_moment)) / zeroth_moment
-                # calculate the 'C' normalisation constant
-                normalisation_constant_C = -normalisation_constant_B * absvpa_wpa_moment / vpa2_wpa_moment
-                # updated pdf is old pdf * (normalisation_constant_A + |vpa| * normalisation_constant_B + vpa^2 * normalisation_constant_C)
-                @. pdf[:,1,end,ir] *= (normalisation_constant_A + abs(vpa.scratch2) * normalisation_constant_B 
-                                    + vpa.scratch2^2 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "vpa4"
-                # calculate int dwpa vpa^4 * pdf = vpa4_moment
-                @. vpa.scratch3 = vpa.scratch2^4 * pdf[:,1,end,ir]
-                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + vpa^4 * normalisation_constant_C)
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-                normalisation_constant_A = (1 + normalisation_constant_B 
-                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-                @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_B 
-                                    + vpa.scratch2^4 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "vpa4_gaussian"
-                afac = 0.1
-                bfac = 0.2
-                # calculate int dwpa vpa^2 * exp(-vpa^2) * pdf = vpa2_moment
-                @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.scratch2^2 * exp(-afac * vpa.scratch2^2)
-                vpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^2 * exp(-vpa^2) * wpa * pdf = vpa2_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa2_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa wpa^2 * vpa^2 * exp(-vpa^2) * pdf = vpa2_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa2_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * exp(-vpa^2) * pdf = vpa4_moment
-                @. vpa.scratch3 = vpa.scratch2^4 * exp(-bfac * vpa.scratch2^2) * pdf[:,1,end,ir]
-                vpa4_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa * pf = vpa4_wpa_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # calculate int dwpa vpa^4 * wpa^2 * pdf = vpa4_wpa2_moment
-                @. vpa.scratch3 *= vpa.grid
-                vpa4_wpa2_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-                # assuming pdf_updated = pdf * (normalisation_constant_A + vpa^2 * normalisation_constant_B + exp(-vpa^2) * vpa^4 * normalisation_constant_C)
-                normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                     (vpa2_wpa2_moment - vpa2_wpa_moment * vpa4_wpa2_moment / vpa4_wpa_moment
-                     + wpa2_moment / zeroth_moment * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment))
-                normalisation_constant_A = (1 + normalisation_constant_B 
-                    * (vpa2_wpa_moment * vpa4_moment / vpa4_wpa_moment - vpa2_moment)) / zeroth_moment
-                normalisation_constant_C = -normalisation_constant_B * vpa2_wpa_moment / vpa4_wpa_moment
-                #normalisation_constant_B = (0.5 - wpa2_moment / zeroth_moment) /
-                #                           (vpa2_wpa2_moment
-                #                            - wpa2_moment / zeroth_moment * vpa2_moment)
-                #normalisation_constant_A = (1 - normalisation_constant_B 
-                #                                * vpa2_moment) / zeroth_moment
-                #normalisation_constant_C = 0.0
-                @. pdf[:,1,end,ir] *= (normalisation_constant_A + exp(-afac * vpa.scratch2^2) * vpa.scratch2^2 * normalisation_constant_B 
-                                    + exp(-bfac * vpa.scratch2^2) * vpa.scratch2^4 * normalisation_constant_C)
-            elseif pdf_adjustment_option == "no1st_vpa2"
-                normalisation_constant_C = (1.0 - 0.5*zeroth_moment/wpa2_moment) /
-                                           (vpa2_moment - zeroth_moment*vpa2_wpa2_moment / wpa2_moment)
-                normalisation_constant_A = (0.5 - normalisation_constant_C*vpa2_wpa2_moment) / wpa2_moment
-                @. pdf[:,1,end,ir] *= (normalisation_constant_A + vpa.scratch2^2 * normalisation_constant_C)
-            else
-                println("pdf_adjustment_option not recognised")
-                stop()
-            end
+            phi[end,ir] = me_over_mi * vmax^2
+            pdf[1:sigma_ind,1,end,ir] .= reversed_pdf[1:sigma_ind]
 
-            moments.electron.constraints_A_coefficient[end,ir] = normalisation_constant_A
-            moments.electron.constraints_B_coefficient[end,ir] = normalisation_constant_B
-            moments.electron.constraints_C_coefficient[end,ir] = normalisation_constant_C
-
-            # smooth the pdf at the boundary
-            #for ivpa ∈ ivpa_min+1:vpa.n-1
-            #    pdf[ivpa,1,end,ir] = (pdf[ivpa-1,1,end,ir] + pdf[ivpa+1,1,end,ir]) / 2.0
-            #end
-
-            # zeroth_vspace_moment = integrate_over_vspace(pdf[:,1,end,1], vpa.wgts)
-            # @. vpa.scratch3 = pdf[:,1,end,1] * vpa.grid
-            # first_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @. vpa.scratch3 *= vpa.grid
-            # second_vspace_moment = integrate_over_vspace(vpa.scratch3, vpa.wgts)
-            # @loop_vpa ivpa begin
-            #     println(io_pdf_stages, "vpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa,1,end,ir], " zeroth_vspace_moment: ", zeroth_vspace_moment, 
-            #         " first_vspace_moment: ", first_vspace_moment, " second_vspace_moment: ", second_vspace_moment, " stage: ", 3)
-            # end
-            # println(io_pdf_stages,"")
+            moments.electron.constraints_A_coefficient[end,ir] = A
+            moments.electron.constraints_B_coefficient[end,ir] = 0.0
+            moments.electron.constraints_C_coefficient[end,ir] = C
         end
     end
 
-    # # initialise the electron pdf for positive vpa to be the mirror reflection about vpa = 0
-    # ivpa_zero = (vpa.n+1)÷2
-    # ivpa_max = vpa.n
-    # @. pdf[ivpa_zero+1:end,:,1,:] = pdf[ivpa_zero-1:-1:1,:,1,:]
-    # # calculate the zeroth v-space moment of the normalised electron pdf:
-    # # if unity to within specified tolerance, then the boundary condition is satisfied;
-    # # otherwise, modify the cutoff velocity and repeat
-    # @loop_r ir begin
-    #     unity = 2.0
-    #     while unity > 1.0
-    #         unity = integrate_over_vspace(pdf[:,1,1,ir], vpa.wgts)
-    #         # if unity > 1.0, then the cutoff velocity is too high so reduce it
-    #         if unity > 1.0
-    #             pdf[ivpa_max,1,1,ir] = 0.0
-    #             ivpa_max -= 1
-    #         end
-    #     end
-    #     phi[1,ir] = vthe[1,ir]^2 * vpa.grid[ivpa_max]^2
-    # end
-    # # repeat the above procedure for the boundary at z_max
-    # @. pdf[ivpa_zero-1:-1:1,:,end,:] = pdf[ivpa_zero+1:end,:,end,:]
-    # ivpa_max = 1
-    # @loop_r ir begin
-    #     unity = 2.0
-    #     while unity > 1.0
-    #         unity = integrate_over_vspace(pdf[:,1,end,ir], vpa.wgts)
-    #         if unity > 1.0
-    #             pdf[ivpa_max,1,end,ir] = 0.0
-    #             ivpa_max += 1
-    #         end
-    #         phi[end,ir] = vthe[end,ir]^2 * vpa.grid[ivpa_max]^2
-    #     end
-    #     #println("unity: ", unity)
-    # end
+    return nothing
 end
 
 """
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index dee497b0b..4174c5954 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -332,7 +332,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
             [(speed=speed,)], moments, num_diss_params,
-            composition.me_over_mi, scratch_dummy)
+            composition.me_over_mi, scratch_dummy; newton_tol=0.1*t_params.electron.rtol)
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
@@ -1370,7 +1370,7 @@ this 'initital' value for the electron will just be the first guess in an iterat
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
         vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, moments, num_diss_params,
-        me_over_mi, scratch_dummy; restart_from_boltzmann=false)
+        me_over_mi, scratch_dummy; restart_from_boltzmann=false, newton_tol)
 
     if z.bc == "wall"
         begin_r_region()
@@ -1389,7 +1389,7 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
                                                     vperp_spectral, vpa_spectral,
                                                     vpa_advect, moments,
                                                     num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                    me_over_mi)
+                                                    me_over_mi, newton_tol)
 
         # Distribute the z-boundary pdf values to every process
         begin_serial_region()
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 7fbcb8112..7aa66846f 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -651,6 +651,14 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     # Now that `t_params` and `scratch` have been created, initialize electrons if
     # necessary
     if composition.electron_physics != restart_electron_physics
+        begin_serial_region()
+        @serial_region begin
+            # zero-initialise phi here, because the boundary points of phi are used as an
+            # effective 'cache' for the sheath-boundary cutoff speed for the electrons, so
+            # needs to be initialised to something, but phi cannot be calculated properly
+            # until after the electrons are initialised.
+            fields.phi .= 0.0
+        end
         initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions, gyroavs,

From 3d3b87a8e22a921140fb97b01f4c3d7f87fe2cdc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 31 May 2024 13:31:03 +0100
Subject: [PATCH 266/428] Add kinetic electron case to precompilation runs
 inputs

Also add a script to precompile specifically for kinetic electron runs.
---
 precompile-kinetic-electrons.jl          | 16 +++++++
 util/precompile_run.jl                   | 18 +++++--
 util/precompile_run_kinetic-electrons.jl | 61 ++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 3 deletions(-)
 create mode 100644 precompile-kinetic-electrons.jl
 create mode 100644 util/precompile_run_kinetic-electrons.jl

diff --git a/precompile-kinetic-electrons.jl b/precompile-kinetic-electrons.jl
new file mode 100644
index 000000000..9922963aa
--- /dev/null
+++ b/precompile-kinetic-electrons.jl
@@ -0,0 +1,16 @@
+using Pkg
+
+# Activate the moment_kinetics package
+Pkg.activate(".")
+
+using PackageCompiler
+
+# Create the sysimage 'moment_kinetics.so' in the base moment_kinetics source directory
+# with both moment_kinetics and the dependencies listed above precompiled.
+# Warning: editing the code will not affect what runs when using this .so, you
+# need to re-precompile if you change anything.
+create_sysimage(; sysimage_path="moment_kinetics.so",
+                precompile_execution_file="util/precompile_run_kinetic-electrons.jl",
+                include_transitive_dependencies=false, # This is needed to make MPI work, see https://github.com/JuliaParallel/MPI.jl/issues/518
+                sysimage_build_args=`-O3 --check-bounds=no`, # Assume if we are precompiling we want an optimized, production build
+               )
diff --git a/util/precompile_run.jl b/util/precompile_run.jl
index f5a7fb512..83c60f73c 100644
--- a/util/precompile_run.jl
+++ b/util/precompile_run.jl
@@ -27,7 +27,7 @@ base_input = Dict("run_name" => "precompilation",
                   "vpa_ngrid" => 7,
                   "vpa_nelement" => 3,
                   "vpa_bc" => "zero",
-                  "vpa_L" => 4.0,
+                  "vpa_L" => 8.0,
                   "vpa_discretization" => "finite_difference",
                   "vzeta_ngrid" => 5,
                   "vzeta_nelement" => 3,
@@ -42,7 +42,7 @@ base_input = Dict("run_name" => "precompilation",
                   "vz_ngrid" => 7,
                   "vz_nelement" => 3,
                   "vz_bc" => "zero",
-                  "vz_L" => 4.0,
+                  "vz_L" => 8.0,
                   "vz_discretization" => "finite_difference",
                   "timestepping" => Dict{String,Any}("nstep" => 1))
 cheb_input = merge(base_input, Dict("r_discretization" => "chebyshev_pseudospectral",
@@ -76,7 +76,19 @@ collisions_input = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
 geo_input1 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
                                             "geometry" => Dict{String,Any}("option" => "1D-mirror", "DeltaB" => 0.5, "pitch" => 0.5, "rhostar" => 1.0))) 
 
-push!(inputs_list, collisions_input, geo_input1)
+kinetic_electron_input = merge(cheb_input, Dict("evolve_moments_density" => true,
+                                                "evolve_moments_parallel_flow" => true,
+                                                "evolve_moments_parallel_pressure" => true,
+                                                "electron_physics" => "kinetic_electrons",
+                                                "electron_timestepping" => Dict{String,Any}("nstep" => 1,
+                                                                                            "dt" => 2.0e-11,
+                                                                                            "initialization_residual_value" => 1.0e10,
+                                                                                            "converged_residual_value" => 1.0e10,
+                                                                                            "rtol" => 1.0e10,
+                                                                                            "no_restart" => true),
+                                               ))
+
+push!(inputs_list, collisions_input, geo_input1, kinetic_electron_input)
 
 for input in inputs_list
     run_moment_kinetics(input)
diff --git a/util/precompile_run_kinetic-electrons.jl b/util/precompile_run_kinetic-electrons.jl
new file mode 100644
index 000000000..b605fcf16
--- /dev/null
+++ b/util/precompile_run_kinetic-electrons.jl
@@ -0,0 +1,61 @@
+# provide option of running from command line via 'julia moment_kinetics.jl'
+using Pkg
+Pkg.activate(".")
+
+using moment_kinetics
+
+# Create a temporary directory for test output
+test_output_directory = tempname()
+mkpath(test_output_directory)
+
+input = Dict("run_name" => "precompilation",
+             "base_directory" => test_output_directory,
+             "dt" => 0.0,
+             "evolve_moments_density" => true,
+             "evolve_moments_parallel_flow" => true,
+             "evolve_moments_parallel_pressure" => true,
+             "electron_physics" => "kinetic_electrons",
+             "r_ngrid" => 1,
+             "r_nelement" => 1,
+             "r_bc" => "periodic",
+             "r_discretization" => "chebyshev_pseudospectral",
+             "z_ngrid" => 5,
+             "z_nelement" => 4,
+             "z_bc" => "wall",
+             "z_discretization" => "chebyshev_pseudospectral",
+             "vperp_ngrid" => 1,
+             "vperp_nelement" => 1,
+             "vperp_bc" => "zero",
+             "vperp_L" => 4.0,
+             "vperp_discretization" => "chebyshev_pseudospectral",
+             "vpa_ngrid" => 7,
+             "vpa_nelement" => 8,
+             "vpa_bc" => "zero",
+             "vpa_L" => 8.0,
+             "vpa_discretization" => "chebyshev_pseudospectral",
+             "vzeta_ngrid" => 1,
+             "vzeta_nelement" => 1,
+             "vzeta_bc" => "zero",
+             "vzeta_L" => 4.0,
+             "vzeta_discretization" => "chebyshev_pseudospectral",
+             "vr_ngrid" => 1,
+             "vr_nelement" => 1,
+             "vr_bc" => "zero",
+             "vr_L" => 4.0,
+             "vr_discretization" => "chebyshev_pseudospectral",
+             "vz_ngrid" => 7,
+             "vz_nelement" => 8,
+             "vz_bc" => "zero",
+             "vz_L" => 8.0,
+             "vz_discretization" => "chebyshev_pseudospectral",
+             "timestepping" => Dict{String,Any}("nstep" => 1,
+                                                "dt" => 2.0e-11),
+             "electron_timestepping" => Dict{String,Any}("nstep" => 1,
+                                                         "dt" => 2.0e-11,
+                                                         "initialization_residual_value" => 1.0e10,
+                                                         "converged_residual_value" => 1.0e10,
+                                                         "rtol" => 1.0e10,
+                                                         "no_restart" => true))
+
+
+run_moment_kinetics(input)

From 3167e0d2e22c1986d5b12bfec3b873b2431ce63d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 31 May 2024 17:06:29 +0100
Subject: [PATCH 267/428] Always reset ion timestep when electron model changes

When restarting a run with changed electron model, it is probably best
to start with a small timestep, so use the value from the input file,
not the reloaded value from the previous run.
---
 moment_kinetics/src/time_advance.jl | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 7aa66846f..48abc61bc 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -426,14 +426,11 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     electron_mom_diss_coeff = num_diss_params.electron.moment_dissipation_coefficient
     neutral_mom_diss_coeff = num_diss_params.neutral.moment_dissipation_coefficient
 
-    if composition.electron_physics == braginskii_fluid &&
-            restart_electron_physics !== braginskii_fluid
+    if composition.electron_physics != restart_electron_physics
 
-        # When restarting braginskii_fluid from a different electron physics type, and
+        # When restarting from a different electron physics type, and
         # using an adaptive timestep, do not want to keep the `dt` from the previous
-        # simulation, as the diffusive behaviour of braginskii fluid means the CFL
-        # condition is likely to be very restrictive, so we probably need to reset to a
-        # small initial timestep.
+        # simulation, in case the new electron physics requires a smaller ion timestep.
         dt_reload = nothing
     end
 

From 5e91466dd8a762295298ac886a661908b3f40ac8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 12:31:13 +0100
Subject: [PATCH 268/428] Do not overwrite boundary values when calculating phi
 from Epar

z-boundary points in `phi` are now used as a cache for the cutoff
velocity used in the boundary condition for kinetic electrons.
---
 moment_kinetics/src/em_fields.jl | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 1d5aa8624..c6df3a405 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -117,7 +117,7 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments
             collisions.nu_ei, moments.electron.parallel_friction,
             composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
             fvec.density_neutral, fvec.uz_neutral, fvec.electron_upar)
-        calculate_phi_from_Epar!(fields.phi, fields.Ez, z)
+        calculate_phi_from_Epar!(fields.phi, fields.Ez, r, z)
     end
     ## can calculate phi at z = L and hence phi_wall(z=L) using jpar_i at z =L if needed
     _block_synchronize()
@@ -169,15 +169,26 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments
     return nothing
 end
 
-function calculate_phi_from_Epar!(phi, Epar, z)
-    # simple calculation of phi from Epar for now, with zero phi assumed at boundary
+function calculate_phi_from_Epar!(phi, Epar, r, z)
+    # simple calculation of phi from Epar for now. Assume phi is already set at the
+    # lower-ze boundary, e.g. by the kinetic electron boundary condition.
     begin_serial_region()
 
     dz = z.cell_width
     @serial_region begin
-        phi[1,:] .= 3.0
-        @loop_r_z ir iz begin
-            if iz > 1
+        # Need to broadcast the lower-z boundary value, because we only communicate
+        # delta_phi below, rather than passing the boundary values directly from block to
+        # block.
+        phi[1,:] .= MPI.bcast(@view(phi[1,:]), z.comm; root=0)
+
+        if z.irank == z.nrank - 1
+            # Don't want to change the upper-z boundary value, so save it here so we can
+            # restore it at the end
+            @views @. r.scratch = phi[end,:]
+        end
+
+        @loop_r ir begin
+            for iz ∈ 2:z.n
                 phi[iz,ir] = phi[iz-1,ir] - dz[iz-1]*Epar[iz,ir]
             end
         end
@@ -193,6 +204,11 @@ function calculate_phi_from_Epar!(phi, Epar, z)
                 end
             end
         end
+
+        if z.irank == z.nrank - 1
+            # Restore the upper-z boundary value
+            @views @. phi[end,:] = r.scratch
+        end
     end
 
     return nothing

From d91c3562e64b1b52e5c76fbc429afea9199fd2e5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 14:37:43 +0100
Subject: [PATCH 269/428] Don't set `next_output_time[]` for electrons when not
 doing electron I/O

When not doing electron I/O we do not need to use `next_output_time[]`
to step to exactly the requested output time. There is a bug where
occasionally the timestep will get set to a negative value, so this
update helps to avoid that bug.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index bb01ec011..d0a057bb6 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -517,6 +517,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             end
             output_counter += 1
             if t_params.adaptive && !t_params.write_after_fixed_step_count &&
+                    io_electron !== nothing &&
                     (output_counter - initial_output_counter ≤ length(dfns_output_times))
                 @serial_region begin
                     t_params.next_output_time[] =

From e6afaec90fdafc629c568c3e8203eddcf42baa92 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 16:40:06 +0100
Subject: [PATCH 270/428] Refactor adaptive timestep I/O times

Move more of the functionality into
`runge_kutta.adaptive_timestep_update_t_params!()` so that it is easier
to keep consistent.
---
 moment_kinetics/src/input_structs.jl |  7 ++-
 moment_kinetics/src/runge_kutta.jl   | 65 ++++++++++++++++++--------
 moment_kinetics/src/time_advance.jl  | 70 +++++++++++-----------------
 3 files changed, 79 insertions(+), 63 deletions(-)

diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 174affdd5..e8973de4d 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -48,8 +48,13 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     dt_before_output::MPISharedArray{mk_float,1}
     dt_before_last_fail::MPISharedArray{mk_float,1}
     CFL_prefactor::mk_float
-    step_to_output::MPISharedArray{Bool,1}
+    step_to_moments_output::MPISharedArray{Bool,1}
+    step_to_dfns_output::MPISharedArray{Bool,1}
+    write_moments_output::MPISharedArray{Bool,1}
+    write_dfns_output::MPISharedArray{Bool,1}
     step_counter::Ref{mk_int}
+    moments_output_counter::Ref{mk_int}
+    dfns_output_counter::Ref{mk_int}
     failure_counter::Ref{mk_int}
     failure_caused_by::Vector{mk_int}
     limit_caused_by::Vector{mk_int}
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 0bf317f00..803cc2715 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1056,8 +1056,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
         error("Unrecognized error_norm_method '$method'")
     end
 
-    just_completed_output_step = false
-
     if !success
         # Iteration failed in implicit part of timestep try decreasing timestep
 
@@ -1075,10 +1073,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
                 t_params.dt_before_last_fail[] = t_params.previous_dt[]
             end
 
-            # If we were trying to take a step to the output timestep, dt will be smaller on
-            # the re-try, so will not reach the output time.
-            t_params.step_to_output[] = false
-
             # Decrease timestep by 1/2 - this factor should probably be settable!
             # Note when nonlinear solve iteration fails, we do not enforce
             # minimum_dt, as the timesolver must error if we do not decrease dt.
@@ -1096,6 +1090,11 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             # Call the 'cause' of the timestep failure the variable that has the biggest
             # error norm here
             t_params.failure_caused_by[end] += 1
+
+            # If we were trying to take a step to the output timestep, dt will be smaller on
+            # the re-try, so will not reach the output time.
+            t_params.step_to_moments_output[] = false
+            t_params.step_to_dfns_output[] = false
         end
     elseif (error_norm > 1.0 || isnan(error_norm)) && current_dt > t_params.minimum_dt * (1.0 + 1.0e-13)
         # (1.0 + 1.0e-13) fudge factor accounts for possible rounding errors when
@@ -1120,10 +1119,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
                 t_params.dt_before_last_fail[] = t_params.previous_dt[]
             end
 
-            # If we were trying to take a step to the output timestep, dt will be smaller on
-            # the re-try, so will not reach the output time.
-            t_params.step_to_output[] = false
-
             # Get new timestep estimate using same formula as for a successful step, but
             # limit decrease to factor 1/2 - this factor should probably be settable!
             t_params.dt[] = max(t_params.dt[] / 2.0,
@@ -1137,6 +1132,11 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             # error norm here
             t_params.failure_caused_by[max_error_variable_index] += 1
 
+            # If we were trying to take a step to the output timestep, dt will be smaller on
+            # the re-try, so will not reach the output time.
+            t_params.step_to_moments_output[] = false
+            t_params.step_to_dfns_output[] = false
+
             #println("t=$t, timestep failed, error_norm=$error_norm, error_norms=$error_norms, decreasing timestep to ", t_params.dt[])
         end
     else
@@ -1145,17 +1145,23 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             # simulation time.
             t_params.previous_dt[] = t_params.dt[]
 
-            if t_params.step_to_output[]
+            if t_params.step_to_moments_output[] || t_params.step_to_dfns_output[]
                 # Completed an output step, reset dt to what it was before it was reduced to reach
                 # the output time
                 t_params.dt[] = t_params.dt_before_output[]
-                t_params.step_to_output[] = false
+
+                if t_params.step_to_moments_output[]
+                    t_params.step_to_moments_output[] = false
+                    t_params.write_moments_output[] = true
+                end
+                if t_params.step_to_dfns_output[]
+                    t_params.step_to_dfns_output[] = false
+                    t_params.write_dfns_output[] = true
+                end
 
                 if t_params.dt[] > CFL_limit
                     t_params.dt[] = CFL_limit
                 end
-
-                just_completed_output_step = true
             else
                 # Adjust timestep according to Fehlberg's suggestion
                 # (https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta%E2%80%93Fehlberg_method).
@@ -1245,12 +1251,33 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
         end
 
         current_time = t + t_params.previous_dt[]
-        if (!t_params.write_after_fixed_step_count && !just_completed_output_step
-            && (current_time + t_params.dt[] >= t_params.next_output_time[]))
+        # Store here to ensure dt_before_output is set correctly when both moments and
+        # dfns are written at the same time.
+        current_dt = t_params.dt[]
+        if (!t_params.write_after_fixed_step_count
+            && !t_params.write_moments_output[]
+            && length(t_params.moments_output_times) > 0
+            && (t_params.moments_output_counter[] ≤ length(t_params.moments_output_times) + 1)
+            && (current_time + t_params.dt[] >= t_params.moments_output_times[t_params.moments_output_counter[]-1]))
+
+            t_params.dt_before_output[] = current_dt
+            t_params.dt[] = t_params.moments_output_times[t_params.moments_output_counter[]-1] - current_time
+            t_params.step_to_moments_output[] = true
+
+            if t_params.dt[] < 0.0
+                error("When trying to step to next output time, made negative timestep "
+                      * "dt=$(t_params.dt[])")
+            end
+        end
+        if (!t_params.write_after_fixed_step_count
+            && !t_params.write_dfns_output[]
+            && length(t_params.dfns_output_times) > 0
+            && (t_params.dfns_output_counter[] ≤ length(t_params.dfns_output_times) + 1)
+            && (current_time + t_params.dt[] >= t_params.dfns_output_times[t_params.dfns_output_counter[]-1]))
 
-            t_params.dt_before_output[] = t_params.dt[]
-            t_params.dt[] = t_params.next_output_time[] - current_time
-            t_params.step_to_output[] = true
+            t_params.dt_before_output[] = current_dt
+            t_params.dt[] = t_params.dfns_output_times[t_params.dfns_output_counter[]-1] - current_time
+            t_params.step_to_dfns_output[] = true
 
             if t_params.dt[] < 0.0
                 error("When trying to step to next output time, made negative timestep "
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 48abc61bc..1650e04fc 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -313,14 +313,20 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
     next_output_time = allocate_shared_float(1)
     dt_before_output = allocate_shared_float(1)
     dt_before_last_fail = allocate_shared_float(1)
-    step_to_output = allocate_shared_bool(1)
+    step_to_moments_output = allocate_shared_bool(1)
+    step_to_dfns_output = allocate_shared_bool(1)
+    write_moments_output = allocate_shared_bool(1)
+    write_dfns_output = allocate_shared_bool(1)
     if block_rank[] == 0
         dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         previous_dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         next_output_time[] = 0.0
         dt_before_output[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         dt_before_last_fail[] = dt_before_last_fail_reload === nothing ? Inf : dt_before_last_fail_reload
-        step_to_output[] = false
+        step_to_moments_output[] = false
+        step_to_dfns_output[] = false
+        write_moments_output[] = false
+        write_dfns_output[] = false
     end
     _block_synchronize()
 
@@ -385,9 +391,11 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
     end
     return time_info(n_variables, t_input["nstep"], end_time, dt_shared, previous_dt_shared,
                      next_output_time, dt_before_output, dt_before_last_fail,
-                     CFL_prefactor, step_to_output, Ref(0), Ref(0), mk_int[], mk_int[],
-                     t_input["nwrite"], t_input["nwrite_dfns"], moments_output_times,
-                     dfns_output_times, t_input["type"], rk_coefs, rk_coefs_implicit,
+                     CFL_prefactor, step_to_moments_output, step_to_dfns_output,
+                     write_moments_output, write_dfns_output, Ref(0), Ref(2), Ref(2),
+                     Ref(0), mk_int[], mk_int[], t_input["nwrite"],
+                     t_input["nwrite_dfns"], moments_output_times, dfns_output_times,
+                     t_input["type"], rk_coefs, rk_coefs_implicit,
                      implicit_coefficient_is_zero, n_rk_stages, rk_order, adaptive,
                      low_storage, t_input["rtol"], t_input["atol"], t_input["atol_upar"],
                      t_input["step_update_prefactor"], t_input["max_increase_factor"],
@@ -1546,16 +1554,6 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
     start_time = now()
 
     epsilon = 1.e-11
-    moments_output_counter = 1
-    dfns_output_counter = 1
-    @serial_region begin
-        if t_params.adaptive && !t_params.write_after_fixed_step_count
-            t_params.next_output_time[] =
-                min(t_params.moments_output_times[moments_output_counter],
-                    t_params.dfns_output_times[dfns_output_counter])
-        end
-    end
-    _block_synchronize()
 
     # main time advance loop
     iwrite_moments = 2
@@ -1570,10 +1568,8 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
     while true
         
         if t_params.adaptive && !t_params.write_after_fixed_step_count
-            maybe_write_moments = (t + t_params.dt[] ≥ t_params.moments_output_times[moments_output_counter] - epsilon
-                                   || t + t_params.dt[] ≥ t_params.end_time - epsilon)
-            maybe_write_dfns = (t + t_params.dt[] ≥ t_params.dfns_output_times[dfns_output_counter] - epsilon
-                                || t + t_params.dt[] ≥ t_params.end_time - epsilon)
+            maybe_write_moments = t_params.step_to_moments_output[]
+            maybe_write_dfns = t_params.step_to_dfns_output[]
         else
             maybe_write_moments = (t_params.step_counter[] % t_params.nwrite_moments == 0
                                    || t_params.step_counter[] >= t_params.nstep)
@@ -1624,10 +1620,14 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
         end
 
         if t_params.adaptive && !t_params.write_after_fixed_step_count
-            write_moments = (t ≥ t_params.moments_output_times[moments_output_counter] - epsilon
-                             || t ≥ t_params.end_time - epsilon)
-            write_dfns = (t ≥ t_params.dfns_output_times[dfns_output_counter] - epsilon
-                          || t ≥ t_params.end_time - epsilon)
+            write_moments = t_params.write_moments_output[]
+            write_dfns = t_params.write_dfns_output[]
+
+            _block_synchronize()
+            @serial_region begin
+                t_params.write_moments_output[] = false
+                t_params.write_dfns_output[] = false
+            end
         else
             write_moments = (t_params.step_counter[] % t_params.nwrite_moments == 0
                              || t_params.step_counter[] >= t_params.nstep)
@@ -1635,26 +1635,10 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                           || t_params.step_counter[] >= t_params.nstep)
         end
         if write_moments
-            moments_output_counter += 1
-            if moments_output_counter ≤ length(t_params.moments_output_times)
-                @serial_region begin
-                    t_params.next_output_time[] =
-                        min(t_params.moments_output_times[moments_output_counter],
-                            t_params.dfns_output_times[dfns_output_counter])
-                end
-            end
-            write_moments = true
+            t_params.moments_output_counter[] += 1
         end
         if write_dfns
-            dfns_output_counter += 1
-            if dfns_output_counter ≤ length(t_params.dfns_output_times)
-                @serial_region begin
-                    t_params.next_output_time[] =
-                        min(t_params.moments_output_times[moments_output_counter],
-                            t_params.dfns_output_times[dfns_output_counter])
-                end
-            end
-            write_dfns = true
+            t_params.dfns_output_counter[] += 1
         end
 
         if write_moments || write_dfns || finish_now
@@ -1705,7 +1689,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             @serial_region begin
                 if global_rank[] == 0
                     print("writing moments output ",
-                          rpad(string(moments_output_counter - 1), 4), "  ",
+                          rpad(string(t_params.moments_output_counter[] - 1), 4), "  ",
                           "t = ", rpad(string(round(t, sigdigits=6)), 7), "  ",
                           "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ")
                     if t_params.adaptive
@@ -1795,7 +1779,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             @serial_region begin
                 if global_rank[] == 0
                     println("writing distribution functions output ",
-                            rpad(string(dfns_output_counter  - 1), 4), "  ",
+                            rpad(string(t_params.dfns_output_counter[] - 1), 4), "  ",
                             "t = ", rpad(string(round(t, sigdigits=6)), 7), "  ",
                             "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ",
                             Dates.format(now(), dateformat"H:MM:SS"))

From d19089e7d5a1f52510034d632d3d004474af550c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 16:48:30 +0100
Subject: [PATCH 271/428] Update electrons for refactored adaptive timestep I/O
 times

---
 .../src/electron_kinetic_equation.jl          | 45 +++++++------------
 moment_kinetics/src/initial_conditions.jl     | 12 ++---
 2 files changed, 21 insertions(+), 36 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index d0a057bb6..b3b000cbe 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -60,7 +60,7 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
-        initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
+        residual_tolerance=nothing, evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -73,7 +73,6 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
             vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
             external_source_settings, num_diss_params, max_electron_pdf_iterations;
             io_electron=io_electron, initial_time=initial_time,
-            initial_output_counter=initial_output_counter,
             residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar)
     elseif solution_method == "shooting_method"
         dens = moments.electron.dens
@@ -138,7 +137,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
-        initial_output_counter=0, residual_tolerance=nothing, evolve_ppar=false)
+        residual_tolerance=nothing, evolve_ppar=false)
 
     begin_r_z_region()
 
@@ -191,15 +190,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     else
         do_debug_io = false
     end
-    @serial_region begin
-        if t_params.adaptive && !t_params.write_after_fixed_step_count && io_electron !== nothing
-            t_params.next_output_time[] = dfns_output_times[1]
-        else
-            # Using fixed output step count, so we don't want to make the adaptive
-            # timestep adjust to output at specific times.
-            t_params.next_output_time[] = Inf
-        end
-    end
 
     #z_speedup_fac = 20.0
     #z_speedup_fac = 5.0
@@ -252,13 +242,13 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
     end
 
-    output_counter = initial_output_counter
     begin_serial_region()
-    output_counter += 1
+    t_params.moments_output_counter[] += 1
     @serial_region begin
         if io_electron !== nothing
             write_electron_state(scratch[1].pdf_electron, moments, t_params, time,
-                                 io_electron, output_counter, r, z, vperp, vpa)
+                                 io_electron, t_params.moments_output_counter[], r, z,
+                                 vperp, vpa)
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
@@ -487,14 +477,14 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 end
             end
         end
-        if ((t_params.adaptive && (time ≥ dfns_output_times[output_counter - initial_output_counter] - epsilon))
+        if ((t_params.adaptive && t_params.write_moments_output[])
             || (!t_params.adaptive && t_params.step_counter[] % t_params.nwrite_moments == 0)
             || (do_debug_io && (t_params.step_counter[] % debug_io_nwrite == 0)))
 
             begin_serial_region()
             @serial_region begin
                 if text_output
-                    if (mod(output_counter, 100) == 0)
+                    if (mod(t_params.moments_output_counter[], 100) == 0)
                         @loop_vpa ivpa begin
                             @loop_z iz begin
                                 println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
@@ -515,20 +505,14 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     println(io_vth,"")
                 end
             end
-            output_counter += 1
-            if t_params.adaptive && !t_params.write_after_fixed_step_count &&
-                    io_electron !== nothing &&
-                    (output_counter - initial_output_counter ≤ length(dfns_output_times))
-                @serial_region begin
-                    t_params.next_output_time[] =
-                        dfns_output_times[output_counter - initial_output_counter]
-                end
-            end
+            t_params.moments_output_counter[] += 1
             @serial_region begin
                 if io_electron !== nothing
+                    t_params.write_moments_output[] = false
                     write_electron_state(scratch[t_params.n_rk_stages+1].pdf_electron,
                                          moments, t_params, time, io_electron,
-                                         output_counter, r, z, vperp, vpa)
+                                         t_params.moments_output_counter[], r, z, vperp,
+                                         vpa)
                 end
             end
         end
@@ -566,9 +550,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
         if !electron_pdf_converged || do_debug_io
             if io_electron !== nothing && io_electron !== true
-                output_counter += 1
+                t_params.moments_output_counter[] += 1
                 write_electron_state(final_scratch_pdf, moments, t_params, time,
-                                     io_electron, output_counter, r, z, vperp, vpa)
+                                     io_electron, t_params.moments_output_counter[], r, z,
+                                     vperp, vpa)
                 finish_electron_io(io_electron)
             end
         end
@@ -578,7 +563,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         error("!!!max number of iterations for electron pdf update exceeded!!!\n"
               * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
     end
-    return time, output_counter
+    return time
 end
 
 function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=false)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 4174c5954..089714c74 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -639,7 +639,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                             t_params, t_params.debug_io[2], -1, nothing,
                                             "electron_debug")
         end
-        electron_pseudotime, n_debug_outputs =
+        electron_pseudotime =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
@@ -656,7 +656,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         if global_rank[] == 0
             println("Initializing electrons - evolving pdf_electron only to steady state")
         end
-        electron_pseudotime, n_debug_outputs =
+        electron_pseudotime =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
@@ -664,8 +664,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         external_source_settings, num_diss_params,
                                         max_electron_pdf_iterations;
                                         io_electron=io_initial_electron,
-                                        initial_time=electron_pseudotime,
-                                        initial_output_counter=n_debug_outputs)
+                                        initial_time=electron_pseudotime)
 
         begin_r_z_vperp_vpa_region()
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
@@ -675,9 +674,10 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
 
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
-        t_idx = n_debug_outputs + 1
+        t_params.moments_output_counter[] += 1
         write_electron_state(pdf.electron.norm, moments, t_params, electron_pseudotime,
-                             io_initial_electron, t_idx, r, z, vperp, vpa)
+                             io_initial_electron, t_params.moments_output_counter[], r, z,
+                             vperp, vpa)
         finish_electron_io(io_initial_electron)
 
     end

From d14be83cc5cf1cc92d75d94ba3849f945919a47b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 20:17:00 +0100
Subject: [PATCH 272/428] Do not adjust adaptive electron timestep for I/O
 times after init

This is not necessary, and can trigger a bug.
---
 moment_kinetics/src/initial_conditions.jl | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 089714c74..202f101ea 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -680,6 +680,13 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                              vperp, vpa)
         finish_electron_io(io_initial_electron)
 
+        # No need to do electron I/O (apart from possibly debug I/O) any more, so if
+        # adaptive timestep is used, it does not need to adjust to output times.
+        resize!(t_params.moments_output_times, 0)
+        resize!(t_params.dfns_output_times, 0)
+        t_params.moments_output_counter[] = 1
+        t_params.dfns_output_counter[] = 1
+
     end
     return nothing
 end

From fc1073579ba1b73cb49a325ebfbbf2a06ae59249 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 1 Jun 2024 22:33:24 +0100
Subject: [PATCH 273/428] Make plots for electron dudz, dpdz and dqdz

---
 moment_kinetics/src/load_data.jl | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index b67983f47..b80d1bb68 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -57,7 +57,8 @@ const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "collision_frequency_ei", "sound_speed", "mach_number")
 const electron_moment_variables = ("electron_density", "electron_parallel_flow",
                                    "electron_parallel_pressure", "electron_thermal_speed",
-                                   "electron_temperature", "electron_parallel_heat_flux")
+                                   "electron_temperature", "electron_parallel_heat_flux",
+                                   "electron_dudz", "electron_dpdz", "electron_dqdz")
 const neutral_moment_variables = ("density_neutral", "uz_neutral", "pz_neutral",
                                   "thermal_speed_neutral", "temperature_neutral",
                                   "qz_neutral")
@@ -4060,6 +4061,24 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
     elseif variable_name == "electron_temperature"
         vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
         variable = run_info.composition.me_over_mi .* vth.^2
+    elseif variable_name == "electron_dudz"
+        upar = postproc_load_variable(run_info, "electron_parallel_flow"; kwargs...)
+        variable = similar(upar)
+        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
+            @views derivative!(variable[:,ir,it], upar[:,ir,it], run_info.z, run_info.z_spectral)
+        end
+    elseif variable_name == "electron_dpdz"
+        ppar = postproc_load_variable(run_info, "electron_parallel_pressure"; kwargs...)
+        variable = similar(ppar)
+        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
+            @views derivative!(variable[:,ir,it], ppar[:,ir,it], run_info.z, run_info.z_spectral)
+        end
+    elseif variable_name == "electron_dqdz"
+        qpar = postproc_load_variable(run_info, "electron_parallel_heat_flux"; kwargs...)
+        variable = similar(qpar)
+        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
+            @views derivative!(variable[:,ir,it], qpar[:,ir,it], run_info.z, run_info.z_spectral)
+        end
     elseif variable_name == "temperature_neutral"
         vth = postproc_load_variable(run_info, "thermal_speed_neutral"; kwargs...)
         variable = vth.^2

From eec97d3c563638edef3d0977a5ee154ef3c64ab6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 2 Jun 2024 13:22:04 +0100
Subject: [PATCH 274/428] Bounds checks in kinetic electron boundary condition

Sometimes searches for index of sigma or vmax might not find a result in
vpa_unnorm - would get out-of-bounds errors if this happened, so check
explicitly.
---
 .../src/electron_kinetic_equation.jl          | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index b3b000cbe..23584cbc5 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -688,6 +688,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
             # -vmax is between vmax_ind-1 and vmax_ind
             vmax_ind = searchsortedfirst(vpa_unnorm, -vmax)
+            if vmax_ind < 2
+                error("In lower-z electron bc, failed to find vpa=-vmax point, vmax_ind=$vmax_ind")
+            end
+            if vmax_ind > vpa.n
+                error("In lower-z electron bc, failed to find vpa=-vmax point, vmax_ind=$vmax_ind")
+            end
 
             # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
             # shift'
@@ -695,6 +701,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
             # sigma is between sigma_ind-1 and sigma_ind
             sigma_ind = searchsortedfirst(vpa_unnorm, 0.0)
+            if sigma_ind < 2
+                error("In lower-z electron bc, failed to find vpa=0 point, sigma_ind=$sigma_ind")
+            end
+            if sigma_ind > vpa.n
+                error("In lower-z electron bc, failed to find vpa=0 point, sigma_ind=$sigma_ind")
+            end
 
             # sigma_fraction is the fraction of the distance between sigma_ind-1 and
             # sigma_ind where sigma is.
@@ -853,6 +865,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
             # -vmax is between vmax_ind and vmax_ind+1
             vmax_ind = searchsortedlast(vpa_unnorm, vmax)
+            if vmax_ind < 1
+                error("In upper-z electron bc, failed to find vpa=vmax point, vmax_ind=$vmax_ind")
+            end
+            if vmax_ind > vpa.n - 1
+                error("In upper-z electron bc, failed to find vpa=vmax point, vmax_ind=$vmax_ind")
+            end
 
             # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
             # shift'
@@ -860,6 +878,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
             # sigma is between sigma_ind and sigma_ind+1
             sigma_ind = searchsortedlast(vpa_unnorm, 0.0)
+            if sigma_ind < 1
+                error("In upper-z electron bc, failed to find vpa=0 point, sigma_ind=$sigma_ind")
+            end
+            if sigma_ind > vpa.n - 1
+                error("In upper-z electron bc, failed to find vpa=0 point, sigma_ind=$sigma_ind")
+            end
 
             # sigma_fraction is the fraction of the distance between sigma_ind+1 and
             # sigma_ind where sigma is.

From c5baa999f9f035718c79653566a4225828b56f32 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 2 Jun 2024 13:25:20 +0100
Subject: [PATCH 275/428] Fix parallelisation of electron I/O

---
 moment_kinetics/src/initial_conditions.jl | 33 ++++++++++++-----------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 202f101ea..5619ea764 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -557,6 +557,10 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         else
             restart_filename = get_default_restart_filename(io_input, "initial_electron";
                                                             error_if_no_file_found=false)
+            # Synchronize to ensure that some processes do not detect the restart file
+            # when they should not, because this function gets called after the other
+            # processes create the file.
+            MPI.Barrier(comm_world)
         end
         if restart_filename === nothing
             # No file to restart from
@@ -582,27 +586,24 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             # might have been loaded from a restart file).
             code_time = MPI.Bcast(code_time, 0, comm_block[])
         end
-        # Set to `true` rather than `nothing` so that processes that are not writing
-        # output (i.e. not rank-0 of their shared-memory block) know that 'initial
-        # electron output' is being written (so that they know not to activate 'debug
-        # I/O').
-        io_initial_electron = true
-        @serial_region begin
-            # Setup I/O for initial electron state
-            io_initial_electron = setup_electron_io(io_input, vpa, vperp, z, r,
-                                                    composition, collisions,
-                                                    moments.evolve_density,
-                                                    moments.evolve_upar,
-                                                    moments.evolve_ppar,
-                                                    external_source_settings, t_params,
-                                                    input_dict, restart_time_index,
-                                                    previous_runs_info,
-                                                    "initial_electron")
+        # Setup I/O for initial electron state
+        io_initial_electron = setup_electron_io(io_input, vpa, vperp, z, r,
+                                                composition, collisions,
+                                                moments.evolve_density,
+                                                moments.evolve_upar,
+                                                moments.evolve_ppar,
+                                                external_source_settings, t_params,
+                                                input_dict, restart_time_index,
+                                                previous_runs_info,
+                                                "initial_electron")
 
+        begin_serial_region()
+        @serial_region begin
             # update the electron pdf in the first scratch
             scratch[1].pdf_electron .= pdf.electron.norm
         end
 
+        begin_r_z_region()
         @loop_r_z ir iz begin
             # update the electron thermal speed using the updated electron parallel pressure
             moments.electron.vth[iz,ir] = sqrt(abs(2.0 * moments.electron.ppar[iz,ir] / (moments.electron.dens[iz,ir] * composition.me_over_mi)))

From af87dca147399ac285f15536b5045e54e3f678a1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 2 Jun 2024 16:01:53 +0100
Subject: [PATCH 276/428] Reload phi when restarting

z-boundary values of phi are used for boundary condition of kinetic
electrons, so useful to reload when restarting.
---
 moment_kinetics/src/load_data.jl       | 6 +++++-
 moment_kinetics/src/moment_kinetics.jl | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index b80d1bb68..1aff1b581 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -577,7 +577,7 @@ end
 """
 Reload pdf and moments from an existing output file.
 """
-function reload_evolving_fields!(pdf, moments, boundary_distributions,
+function reload_evolving_fields!(pdf, moments, fields, boundary_distributions,
                                  restart_prefix_iblock, time_index, composition, geometry,
                                  r, z, vpa, vperp, vzeta, vr, vz)
     code_time = 0.0
@@ -642,6 +642,10 @@ function reload_evolving_fields!(pdf, moments, boundary_distributions,
                 get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp,
                                   restart_vpa, restart_vzeta, restart_vr, restart_vz)
 
+            fields.phi .= reload_electron_moment("phi", dynamic, time_index, r, z,
+                                                 r_range, z_range, restart_r,
+                                                 restart_r_spectral, restart_z,
+                                                 restart_z_spectral, interpolation_needed)
             moments.ion.dens .= reload_moment("density", dynamic, time_index, r, z,
                                               r_range, z_range, restart_r,
                                               restart_r_spectral, restart_z,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index f76935e85..06d65f922 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -309,7 +309,7 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         # Reload pdf and moments from an existing output file
         code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
         previous_runs_info, restart_time_index, restart_electron_physics =
-            reload_evolving_fields!(pdf, moments, boundary_distributions,
+            reload_evolving_fields!(pdf, moments, fields, boundary_distributions,
                                     backup_prefix_iblock, restart_time_index,
                                     composition, geometry, r, z, vpa, vperp, vzeta, vr,
                                     vz)

From d806cca8079582927338e960c446bee8ec516f34 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 2 Jun 2024 17:30:40 +0100
Subject: [PATCH 277/428] Reduce initial temperature for kinetic electron
 simulation

---
 .../wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
index ceaf3a8c8..965f192b9 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
@@ -10,7 +10,7 @@ recycling_fraction = 0.5
 T_e = 0.2 # 1.0
 T_wall = 0.1
 initial_density1 = 1.0
-initial_temperature1 = 1.0
+initial_temperature1 = 0.1
 z_IC_option1 = "gaussian"
 z_IC_density_amplitude1 = 0.001
 z_IC_density_phase1 = 0.0

From ed8a3532c9b4c3e35356a01c514384cdbd88b247 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 2 Jun 2024 22:40:15 +0100
Subject: [PATCH 278/428] Decrease T_e in boltzmann run to restart kinetic
 electrons from

Maybe its better to have the Boltzmann, constant T_e lower than the
value that the kinetic electron simulation will have at the sheath
entrance.
---
 ...all-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
index 098618577..58998f2f0 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-vpadiss0.toml
@@ -7,12 +7,12 @@ evolve_moments_parallel_flow = true
 evolve_moments_parallel_pressure = true
 evolve_moments_conservation = true
 recycling_fraction = 0.5
-T_e = 0.2 # 1.0
+T_e = 0.05 # 1.0
 T_wall = 0.1
 initial_density1 = 1.0
 initial_temperature1 = 1.0
 z_IC_option1 = "gaussian"
-z_IC_density_amplitude1 = 0.001
+z_IC_density_amplitude1 = 1.0
 z_IC_density_phase1 = 0.0
 z_IC_upar_amplitude1 = 1.0
 z_IC_upar_phase1 = 0.0
@@ -66,7 +66,7 @@ vz_discretization = "chebyshev_pseudospectral"
 [timestepping]
 type = "Fekete4(3)"
 #nstep = 50000
-nstep = 1000000
+nstep = 3000000
 dt = 1.0e-6
 minimum_dt = 1.0e-6
 nwrite = 10000

From c486a28f614152761417e9dc087812ee8f2fda00 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 3 Jun 2024 17:30:17 +0100
Subject: [PATCH 279/428] Fix from merge of electrons and imex for timestep
 diagnostics

---
 .../src/makie_post_processing.jl                  | 15 ++++++++-------
 moment_kinetics/src/time_advance.jl               |  4 ++++
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index b24bfe497..9751b3770 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -7528,15 +7528,16 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                             label=prefix * "ion upar RK accuracy", ax=ax,
                             linestyle=:dash)
                 end
-                if electron || ri.evolve_ppar
+                if !electron && ri.evolve_ppar
+                    counter += 1
+                    plot_1d(time, @view limit_caused_by_per_output[counter,:];
+                            label=prefix * "ion ppar RK accuracy", ax=ax,
+                            linestyle=:dash)
+                end
+                if electron || ri.composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
                     counter += 1
-                    if electron
-                        label = prefix * "electron ppar RK accuracy"
-                    else
-                        label = prefix * "ion ppar RK accuracy"
-                    end
                     plot_1d(time, @view limit_caused_by_per_output[counter,:];
-                            label=label, ax=ax,
+                            label=prefix * "electron ppar RK accuracy", ax=ax,
                             linestyle=:dash)
                 end
                 if !electron && ri.n_neutral_species > 0
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 1650e04fc..e02363656 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -483,6 +483,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         # ion pressure
         n_variables += 1
     end
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        # electron pressure
+        n_variables += 1
+    end
     if composition.n_neutral_species > 0
         # neutral pdf
         n_variables += 1

From 0c454f114ab3f5af66ebaa2b024d1fca2379b60d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 6 Jun 2024 13:45:02 +0100
Subject: [PATCH 280/428] Electron time advance applies boundary conditions in
 timestep error

---
 .../src/electron_kinetic_equation.jl          | 156 ++++++++++++------
 1 file changed, 103 insertions(+), 53 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 23584cbc5..9522997af 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -16,7 +16,8 @@ using ..communication
 using ..interpolation: interpolate_to_grid_1d!
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float
-using ..electron_fluid_equations: calculate_electron_qpar_from_pdf!
+using ..electron_fluid_equations: update_electron_vth_temperature!,
+                                  calculate_electron_qpar_from_pdf!
 using ..electron_fluid_equations: electron_energy_equation!
 using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
 using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
@@ -24,7 +25,7 @@ using ..external_sources: external_electron_source!
 using ..file_io: get_electron_io_info, write_electron_state, finish_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
-using ..runge_kutta: rk_update_variable!, rk_error_variable!, local_error_norm,
+using ..runge_kutta: rk_update_variable!, rk_loworder_solution!, local_error_norm,
                      adaptive_timestep_update_t_params!
 using ..utils: get_minimum_CFL_z, get_minimum_CFL_vpa
 using ..velocity_moments: integrate_over_vspace
@@ -286,37 +287,29 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
             rk_update_variable!(scratch, nothing, :pdf_electron, t_params, istage)
 
-            latest_pdf = scratch[istage+1].pdf_electron
-            begin_r_z_vperp_vpa_region()
-            @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-                latest_pdf[ivpa,ivperp,iz,ir] = max(latest_pdf[ivpa,ivperp,iz,ir], 0.0)
-            end
-
-            # enforce the boundary condition(s) on the electron pdf
-            enforce_boundary_condition_on_electron_pdf!(scratch[istage+1].pdf_electron, phi,
-                                                        moments.electron.vth,
-                                                        moments.electron.upar, z, vperp,
-                                                        vpa, vperp_spectral, vpa_spectral,
-                                                        vpa_advect, moments,
-                                                        num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                        composition.me_over_mi, 0.1 * t_params.rtol)
+            if evolve_ppar
+                rk_update_variable!(scratch, nothing, :electron_ppar, t_params, istage)
 
-            begin_r_z_region()
-            A = moments.electron.constraints_A_coefficient
-            B = moments.electron.constraints_B_coefficient
-            C = moments.electron.constraints_C_coefficient
-            @loop_r_z ir iz begin
-                if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
-                    continue
+                begin_r_z_region()
+                moments_struct_ppar = moments.electron.ppar
+                scratch_ppar = scratch[istage+1].electron_ppar
+                @loop_r_z ir iz begin
+                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
                 end
-                (A[iz,ir], B[iz,ir], C[iz,ir]) =
-                    @views hard_force_moment_constraints!(latest_pdf[:,:,iz,ir],
-                                                          (evolve_density=true,
-                                                           evolve_upar=true,
-                                                           evolve_ppar=true), vpa)
+                _block_synchronize()
+
+                update_electron_vth_temperature!(moments, moments_struct_ppar,
+                                                 moments.electron.dens, composition)
             end
+
+            apply_electron_bc_and_constraints!(scratch[istage+1], phi, moments, z, vperp,
+                                               vpa, vperp_spectral, vpa_spectral,
+                                               vpa_advect, num_diss_params, composition,
+                                               t_params)
+
+            latest_pdf = scratch[istage+1].pdf_electron
             
-            function update_derived_moments_and_derivatives()
+            function update_derived_moments_and_derivatives(update_vth=false)
                 # update the electron heat flux
                 moments.electron.qpar_updated[] = false
                 calculate_electron_qpar_from_pdf!(moments.electron.qpar,
@@ -342,12 +335,16 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     this_dens = moments.electron.dens
                     this_vth = moments.electron.vth
                     this_dvth_dz = moments.electron.dvth_dz
+                    if update_vth
+                        @loop_r_z ir iz begin
+                            # update the electron thermal speed using the updated electron
+                            # parallel pressure
+                            this_vth[iz,ir] = sqrt(abs(2.0 * this_ppar[iz,ir] /
+                                                       (this_dens[iz,ir] *
+                                                        composition.me_over_mi)))
+                        end
+                    end
                     @loop_r_z ir iz begin
-                        # update the electron thermal speed using the updated electron
-                        # parallel pressure
-                        this_vth[iz,ir] = sqrt(abs(2.0 * this_ppar[iz,ir] /
-                                                   (this_dens[iz,ir] *
-                                                    composition.me_over_mi)))
                         # update the z-derivative of the electron thermal speed from the
                         # z-derivatives of the electron density and parallel pressure
                         this_dvth_dz[iz,ir] = 0.5 * this_vth[iz,ir] *
@@ -376,8 +373,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
             if t_params.adaptive && istage == t_params.n_rk_stages
                 electron_adaptive_timestep_update!(scratch, time, t_params, moments,
-                                                   z_advect, vpa_advect, r, z, vperp, vpa,
-                                                   external_source_settings;
+                                                   phi, z_advect, vpa_advect, composition,
+                                                   r, z, vperp, vpa, vperp_spectral,
+                                                   vpa_spectral, external_source_settings,
+                                                   num_diss_params;
                                                    evolve_ppar=evolve_ppar)
                 # Re-do this in case electron_adaptive_timestep_update!() re-arranged the
                 # `scratch` vector
@@ -388,17 +387,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     # Re-calculate moments and moment derivatives as the timstep needs to
                     # be re-done with a smaller dt, so scratch[t_params.n_rk_stages+1] has
                     # been reset to the values from the beginning of the timestep here.
-                    update_derived_moments_and_derivatives()
-                end
-            end
-            if evolve_ppar
-                rk_update_variable!(scratch, nothing, :electron_ppar, t_params, istage)
-                moments_struct_ppar = moments.electron.ppar
-                scratch_ppar = scratch[istage+1].electron_ppar
-                @loop_r_z ir iz begin
-                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
+                    update_derived_moments_and_derivatives(true)
                 end
-                _block_synchronize()
             end
         end
 
@@ -608,6 +598,40 @@ function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=fal
     return nothing
 end
 
+function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp, vpa,
+                                            vperp_spectral, vpa_spectral, vpa_advect,
+                                            num_diss_params, composition, t_params)
+    latest_pdf = this_scratch.pdf_electron
+
+    begin_r_z_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        latest_pdf[ivpa,ivperp,iz,ir] = max(latest_pdf[ivpa,ivperp,iz,ir], 0.0)
+    end
+
+    # enforce the boundary condition(s) on the electron pdf
+    enforce_boundary_condition_on_electron_pdf!(latest_pdf, phi, moments.electron.vth,
+                                                moments.electron.upar, z, vperp, vpa,
+                                                vperp_spectral, vpa_spectral, vpa_advect,
+                                                moments,
+                                                num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
+                                                composition.me_over_mi, 0.1 * t_params.rtol)
+
+    begin_r_z_region()
+    A = moments.electron.constraints_A_coefficient
+    B = moments.electron.constraints_B_coefficient
+    C = moments.electron.constraints_C_coefficient
+    @loop_r_z ir iz begin
+        if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
+            continue
+        end
+        (A[iz,ir], B[iz,ir], C[iz,ir]) =
+            @views hard_force_moment_constraints!(latest_pdf[:,:,iz,ir],
+                                                  (evolve_density=true,
+                                                   evolve_upar=true,
+                                                   evolve_ppar=true), vpa)
+    end
+end
+
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, moments, vpa_diffusion,
@@ -1015,14 +1039,20 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 end
 
 """
-    electron_adaptive_timestep_update!(scratch, t_params, rk_coefs, moments)
+    electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi, z_advect,
+                                       vpa_advect, composition, r, z, vperp, vpa,
+                                       vperp_spectral, vpa_spectral,
+                                       external_source_settings, num_diss_params;
+                                       evolve_ppar=false)
 
 Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
-function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_advect,
-                                            vpa_advect, r, z, vperp, vpa,
-                                            external_source_settings; evolve_ppar=false)
+function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi, z_advect,
+                                            vpa_advect, composition, r, z, vperp, vpa,
+                                            vperp_spectral, vpa_spectral,
+                                            external_source_settings, num_diss_params;
+                                            evolve_ppar=false)
     #error_norm_method = "Linf"
     error_norm_method = "L2"
 
@@ -1077,8 +1107,29 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     skip_z_lower = z.irank != 0
 
     # Calculate error ion distribution functions
-    # Note rk_error_variable!() stores the calculated error in `scratch[2]`.
-    rk_error_variable!(scratch, nothing, :pdf_electron, t_params)
+    # Note rk_loworder_solution!() stores the calculated error in `scratch[2]`.
+    rk_loworder_solution!(scratch, nothing, :pdf_electron, t_params)
+    if evolve_ppar
+        begin_r_z_region()
+        rk_loworder_solution!(scratch, nothing, :electron_ppar, t_params)
+
+        # Make vth consistent with `scratch[2]`, as it is needed for the electron pdf
+        # boundary condition.
+        update_electron_vth_temperature!(moments, scratch[2].electron_ppar,
+                                         moments.electron.dens, composition)
+    end
+    apply_electron_bc_and_constraints!(scratch[t_params.n_rk_stages+1], phi, moments, z,
+                                       vperp, vpa, vperp_spectral, vpa_spectral,
+                                       vpa_advect, num_diss_params, composition, t_params)
+    if evolve_ppar
+        # Reset vth in the `moments` struct to the result consistent with full-accuracy RK
+        # solution.
+        begin_r_z_region()
+        update_electron_vth_temperature!(moments,
+                                         scratch[t_params.n_rk_stages+1].electron_ppar,
+                                         moments.electron.dens, composition)
+    end
+
     pdf_error = local_error_norm(scratch[2].pdf_electron,
                                  scratch[t_params.n_rk_stages+1].pdf_electron,
                                  t_params.rtol, t_params.atol; method=error_norm_method,
@@ -1090,7 +1141,6 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, z_adv
     # Calculate error for moments, if necessary
     if evolve_ppar
         begin_r_z_region()
-        rk_error_variable!(scratch, nothing, :electron_ppar, t_params)
         p_err = local_error_norm(scratch[2].electron_ppar,
                                  scratch[t_params.n_rk_stages+1].electron_ppar,
                                  t_params.rtol, t_params.atol; method=error_norm_method,

From 6348968d204459fcb0c4631795d6b17aeb89a096 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 3 Jun 2024 22:45:02 +0100
Subject: [PATCH 281/428] Wall plots for electron pdf

---
 .../src/makie_post_processing.jl              | 745 +++++++++---------
 1 file changed, 392 insertions(+), 353 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 9751b3770..0dd67a2af 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -322,7 +322,10 @@ function makie_post_process(run_dir::Union{String,Tuple},
         end
     end
 
-    plot_ion_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
+    plot_charged_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
+    if has_electrons
+        plot_charged_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix, electron=true)
+    end
     if has_neutrals
         plot_neutral_pdf_2D_at_wall(run_info_dfns; plot_prefix=plot_prefix)
     end
@@ -721,6 +724,15 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         animation_ext=this_input_dict["animation_ext"],
        )
 
+    set_defaults_and_check_section!(
+        this_input_dict, "wall_pdf_electron";
+        plot=false,
+        animate=false,
+        advection_velocity=false,
+        colormap=this_input_dict["colormap"],
+        animation_ext=this_input_dict["animation_ext"],
+       )
+
     set_defaults_and_check_section!(
         this_input_dict, "wall_pdf_neutral";
         plot=false,
@@ -4453,7 +4465,7 @@ function animate_f_unnorm_vs_vpa_z(run_info; input=nothing, neutral=false, is=1,
 end
 
 """
-    plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
+    plot_charged_pdf_2D_at_wall(run_info; plot_prefix, electron=false)
 
 Make plots/animations of the ion distribution function at wall boundaries.
 
@@ -4468,79 +4480,61 @@ Settings are read from the `[wall_pdf]` section of the input.
 will be saved with the format `plot_prefix<some_identifying_string>.pdf`. When `run_info`
 is not a Tuple, `plot_prefix` is optional - plots/animations will be saved only if it is
 passed.
-"""
-function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
-    input = Dict_to_NamedTuple(input_dict_dfns["wall_pdf"])
-    if !(input.plot || input.animate || input.advection_velocity)
-        # nothing to do
-        return nothing
-    end
-    if !any(ri !== nothing for ri ∈ run_info)
-        println("Warning: no distribution function output, skipping wall_pdf plots")
-        return nothing
-    end
 
-    z_lower = 1
-    z_upper = run_info[1].z.n
-    if !all(ri.z.n == z_upper for ri ∈ run_info)
-        println("Cannot run plot_ion_pdf_2D_at_wall() for runs with different "
-                * "z-grid sizes. Got $(Tuple(ri.z.n for ri ∈ run_info))")
-        return nothing
-    end
-
-    println("Making plots of ion distribution function at walls")
-    flush(stdout)
+If `electron=true` is passed, plot electron distribution function instead of ion
+distribution function.
+"""
+function plot_charged_pdf_2D_at_wall(run_info; plot_prefix, electron=false)
+    try
+        if electron
+            electron_prefix = "electron_"
+            electron_suffix = "_electron"
+        else
+            electron_prefix = ""
+            electron_suffix = ""
+        end
+        input = Dict_to_NamedTuple(input_dict_dfns["wall_pdf$electron_suffix"])
+        if !(input.plot || input.animate || input.advection_velocity)
+            # nothing to do
+            return nothing
+        end
+        if !any(ri !== nothing for ri ∈ run_info)
+            println("Warning: no distribution function output, skipping wall_pdf plots")
+            return nothing
+        end
 
-    has_rdim = any(ri !== nothing && ri.r.n > 1 for ri ∈ run_info)
-    has_zdim = any(ri !== nothing && ri.z.n > 1 for ri ∈ run_info)
-    is_1V = all(ri !== nothing && ri.vperp.n == 1 for ri ∈ run_info)
-    moment_kinetic = any(ri !== nothing
-                         && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
-                         for ri ∈ run_info)
+        z_lower = 1
+        z_upper = run_info[1].z.n
+        if !all(ri.z.n == z_upper for ri ∈ run_info)
+            println("Cannot run plot_charged_pdf_2D_at_wall() for runs with different "
+                    * "z-grid sizes. Got $(Tuple(ri.z.n for ri ∈ run_info))")
+            return nothing
+        end
 
-    nt = minimum(ri.nt for ri ∈ run_info)
+        if electron
+            println("Making plots of electron distribution function at walls")
+        else
+            println("Making plots of ion distribution function at walls")
+        end
+        flush(stdout)
 
-    for (z, z_range, label) ∈ ((z_lower, z_lower:z_lower+4, "wall-"),
-                               (z_upper, z_upper-4:z_upper, "wall+"))
-        f_input = copy(input_dict_dfns["f"])
-        f_input["iz0"] = z
+        has_rdim = any(ri !== nothing && ri.r.n > 1 for ri ∈ run_info)
+        has_zdim = any(ri !== nothing && ri.z.n > 1 for ri ∈ run_info)
+        is_1V = all(ri !== nothing && ri.vperp.n == 1 for ri ∈ run_info)
+        moment_kinetic = !electron &&
+                         any(ri !== nothing
+                             && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                             for ri ∈ run_info)
 
-        if input.plot
-            fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    plot_vs_vpa(ri, "f"; is=1, iz=iz, input=f_input,
-                                label="$(run_label)iz=$iz", ax=ax)
-                end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "pdf_$(label)_vs_vpa.pdf"
-            save(outfile, fig)
+        nt = minimum(ri.nt for ri ∈ run_info)
 
-            fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    plot_vs_vpa(ri, "f"; is=1, iz=iz, input=f_input,
-                                label="$(run_label)iz=$iz", ax=ax, yscale=log10,
-                                transform=(x)->positive_or_nan(x; epsilon=1.e-20))
-                end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "logpdf_$(label)_vs_vpa.pdf"
-            save(outfile, fig)
+        for (z, z_range, label) ∈ ((z_lower, z_lower:z_lower+4, "wall-"),
+                                   (z_upper, z_upper-4:z_upper, "wall+"))
+            f_input = copy(input_dict_dfns["f"])
+            f_input["iz0"] = z
 
-            if moment_kinetic
-                fig, ax = get_1d_ax(; xlabel="vpa_unnorm", ylabel="f_unnorm")
+            if input.plot
+                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f$electron_suffix")
                 for iz ∈ z_range
                     for ri ∈ run_info
                         if length(run_info) > 1
@@ -4548,15 +4542,15 @@ function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        plot_f_unnorm_vs_vpa(ri; input=f_input, is=1, iz=iz,
-                                             label="$(run_label)iz=$iz", ax=ax)
+                        plot_vs_vpa(ri, "f$electron_suffix"; is=1, iz=iz, input=f_input,
+                                    label="$(run_label)iz=$iz", ax=ax)
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "pdf_unnorm_$(label)_vs_vpa.pdf"
+                outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa.pdf"
                 save(outfile, fig)
 
-                fig, ax = get_1d_ax(; xlabel="vpa_unnorm", ylabel="f_unnorm")
+                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
                 for iz ∈ z_range
                     for ri ∈ run_info
                         if length(run_info) > 1
@@ -4564,80 +4558,78 @@ function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        plot_f_unnorm_vs_vpa(ri; input=f_input, is=1, iz=iz,
-                                             label="$(run_label)iz=$iz", ax=ax, yscale=log10,
-                                             transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        plot_vs_vpa(ri, "f$electron_suffix"; is=1, iz=iz, input=f_input,
+                                    label="$(run_label)iz=$iz", ax=ax, yscale=log10,
+                                    transform=(x)->positive_or_nan(x; epsilon=1.e-20))
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "logpdf_unnorm_$(label)_vs_vpa.pdf"
+                outfile=plot_prefix * "logpdf$(electron_suffix)_$(label)_vs_vpa.pdf"
                 save(outfile, fig)
-            end
-
-            plot_f_unnorm_vs_vpa(run_info; f_over_vpa2=true, input=f_input, is=1,
-                                 outfile=plot_prefix * "pdf_unnorm_over_vpa2_$(label)_vs_vpa.pdf")
 
-            if !is_1V
-                plot_vs_vpa_vperp(run_info, "f"; is=1, input=f_input,
-                                  outfile=plot_prefix * "pdf_$(label)_vs_vpa_vperp.pdf")
-            end
+                if moment_kinetic
+                    fig, ax = get_1d_ax(; xlabel="vpa_unnorm", ylabel="f$(electron_suffix)_unnorm")
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            plot_f_unnorm_vs_vpa(ri; input=f_input, is=1, iz=iz,
+                                                 label="$(run_label)iz=$iz", ax=ax)
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "pdf_unnorm_$(label)_vs_vpa.pdf"
+                    save(outfile, fig)
+
+                    fig, ax = get_1d_ax(; xlabel="vpa_unnorm", ylabel="f_unnorm")
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            plot_f_unnorm_vs_vpa(ri; input=f_input, is=1, iz=iz,
+                                                 label="$(run_label)iz=$iz", ax=ax, yscale=log10,
+                                                 transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "logpdf_unnorm_$(label)_vs_vpa.pdf"
+                    save(outfile, fig)
+                end
 
-            if has_zdim
-                plot_vs_vpa_z(run_info, "f"; is=1, input=f_input, iz=z_range,
-                              outfile=plot_prefix * "pdf_$(label)_vs_vpa_z.pdf")
-            end
+                if !electron
+                    plot_f_unnorm_vs_vpa(run_info; f_over_vpa2=true, input=f_input, is=1,
+                                         outfile=plot_prefix * "pdf_unnorm_over_vpa2_$(label)_vs_vpa.pdf")
+                end
 
-            if has_rdim && has_zdim
-                plot_vs_z_r(run_info, "f"; is=1, input=f_input, iz=z_range,
-                            outfile=plot_prefix * "pdf_$(label)_vs_z_r.pdf")
-            end
+                if !is_1V
+                    plot_vs_vpa_vperp(run_info, "f$electron_suffix"; is=1, input=f_input,
+                                      outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_vperp.pdf")
+                end
 
-            if has_rdim
-                plot_vs_vpa_r(run_info, "f"; is=1, input=f_input,
-                              outfile=plot_prefix * "pdf_$(label)_vs_vpa_r.pdf")
-            end
-        end
+                if has_zdim
+                    plot_vs_vpa_z(run_info, "f$electron_suffix"; is=1, input=f_input, iz=z_range,
+                                  outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_z.pdf")
+                end
 
-        if input.animate
-            fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
-            frame_index = Observable(1)
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    animate_vs_vpa(ri, "f"; is=1, iz=iz, input=f_input,
-                                   label="$(run_label)iz=$iz", ax=ax,
-                                   frame_index=frame_index)
+                if has_rdim && has_zdim
+                    plot_vs_z_r(run_info, "f$electron_suffix"; is=1, input=f_input, iz=z_range,
+                                outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_z_r.pdf")
                 end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "pdf_$(label)_vs_vpa." * input.animation_ext
-            save_animation(fig, frame_index, nt, outfile)
 
-            fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f", yscale=log10)
-            frame_index = Observable(1)
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    animate_vs_vpa(ri, "f"; is=1, iz=iz, input=f_input,
-                                   label="$(run_label)iz=$iz", ax=ax,
-                                   frame_index=frame_index,
-                                   transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                if has_rdim
+                    plot_vs_vpa_r(run_info, "f$electron_suffix"; is=1, input=f_input,
+                                  outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_r.pdf")
                 end
             end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "logpdf_$(label)_vs_vpa." * input.animation_ext
-            save_animation(fig, frame_index, nt, outfile)
 
-            if moment_kinetic
-                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
+            if input.animate
+                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f$electron_suffix")
                 frame_index = Observable(1)
                 for iz ∈ z_range
                     for ri ∈ run_info
@@ -4646,16 +4638,16 @@ function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        animate_f_unnorm_vs_vpa(ri; is=1, iz=iz, input=f_input,
-                                                label="$(run_label)iz=$iz", ax=ax,
-                                                frame_index=frame_index)
+                        animate_vs_vpa(ri, "f$electron_suffix"; is=1, iz=iz, input=f_input,
+                                       label="$(run_label)iz=$iz", ax=ax,
+                                       frame_index=frame_index)
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "pdf_unnorm_$(label)_vs_vpa." * input.animation_ext
+                outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa." * input.animation_ext
                 save_animation(fig, frame_index, nt, outfile)
 
-                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
+                fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f$electron_suffix", yscale=log10)
                 frame_index = Observable(1)
                 for iz ∈ z_range
                     for ri ∈ run_info
@@ -4664,45 +4656,88 @@ function plot_ion_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        animate_f_unnorm_vs_vpa(ri; is=1, iz=iz, input=f_input,
-                                                label="$(run_label)iz=$iz", ax=ax,
-                                                frame_index=frame_index, yscale=log10,
-                                                transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        animate_vs_vpa(ri, "f$electron_suffix"; is=1, iz=iz, input=f_input,
+                                       label="$(run_label)iz=$iz", ax=ax,
+                                       frame_index=frame_index,
+                                       transform=(x)->positive_or_nan(x; epsilon=1.e-20))
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "logpdf_unnorm_$(label)_vs_vpa." * input.animation_ext
+                outfile=plot_prefix * "logpdf$(electron_suffix)_$(label)_vs_vpa." * input.animation_ext
                 save_animation(fig, frame_index, nt, outfile)
-            end
 
-            animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=true, input=f_input, is=1,
-                                    outfile=plot_prefix * "pdf_unnorm_over_vpa2_$(label)_vs_vpa." * input.animation_ext)
+                if moment_kinetic
+                    fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
+                    frame_index = Observable(1)
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            animate_f_unnorm_vs_vpa(ri; is=1, iz=iz, input=f_input,
+                                                    label="$(run_label)iz=$iz", ax=ax,
+                                                    frame_index=frame_index)
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "pdf_unnorm_$(label)_vs_vpa." * input.animation_ext
+                    save_animation(fig, frame_index, nt, outfile)
+
+                    fig, ax = get_1d_ax(; xlabel="vpa", ylabel="f")
+                    frame_index = Observable(1)
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            animate_f_unnorm_vs_vpa(ri; is=1, iz=iz, input=f_input,
+                                                    label="$(run_label)iz=$iz", ax=ax,
+                                                    frame_index=frame_index, yscale=log10,
+                                                    transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "logpdf_unnorm_$(label)_vs_vpa." * input.animation_ext
+                    save_animation(fig, frame_index, nt, outfile)
+                end
 
-            if !is_1V
-                animate_vs_vpa_vperp(run_info, "f"; is=1, input=f_input,
-                                     outfile=plot_prefix * "pdf_$(label)_vs_vpa_vperp." * input.animation_ext)
-            end
+                if !electron
+                    animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=true, input=f_input, is=1,
+                                            outfile=plot_prefix * "pdf_unnorm_over_vpa2_$(label)_vs_vpa." * input.animation_ext)
+                end
 
-            if has_zdim
-                animate_vs_vpa_z(run_info, "f"; is=1, input=f_input, iz=z_range,
-                                 outfile=plot_prefix * "pdf_$(label)_vs_vpa_z." * input.animation_ext)
-            end
+                if !is_1V
+                    animate_vs_vpa_vperp(run_info, "f$electron_suffix"; is=1, input=f_input,
+                                         outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_vperp." * input.animation_ext)
+                end
 
-            if has_rdim && has_zdim
-                animate_vs_z_r(run_info, "f"; is=1, input=f_input, iz=z_range,
-                               outfile=plot_prefix * "pdf_$(label)_vs_z_r." * input.animation_ext)
-            end
+                if has_zdim
+                    animate_vs_vpa_z(run_info, "f$electron_suffix"; is=1, input=f_input, iz=z_range,
+                                     outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_z." * input.animation_ext)
+                end
 
-            if has_rdim
-                animate_vs_vpa_r(run_info, "f"; is=1, input=f_input,
-                                 outfile=plot_prefix * "pdf_$(label)_vs_vpa_r." * input.animation_ext)
+                if has_rdim && has_zdim
+                    animate_vs_z_r(run_info, "f$electron_suffix"; is=1, input=f_input, iz=z_range,
+                                   outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_z_r." * input.animation_ext)
+                end
+
+                if has_rdim
+                    animate_vs_vpa_r(run_info, "f$electron_suffix"; is=1, input=f_input,
+                                     outfile=plot_prefix * "pdf$(electron_suffix)_$(label)_vs_vpa_r." * input.animation_ext)
+                end
             end
-        end
 
-        if input.advection_velocity
-            animate_vs_vpa(run_info, "vpa_advect_speed"; is=1, input=f_input,
-                           outfile=plot_prefix * "vpa_advect_speed_$(label)_vs_vpa." * input.animation_ext)
+            if input.advection_velocity
+                animate_vs_vpa(run_info, "$(electron_prefix)vpa_advect_speed"; is=1, input=f_input,
+                               outfile=plot_prefix * "$(electron_prefix)vpa_advect_speed_$(label)_vs_vpa." * input.animation_ext)
+            end
         end
+    catch e
+        println("Error in plot_charged_pdf_2D_at_wall(). Error was ", e)
     end
 
     return nothing
@@ -4726,76 +4761,43 @@ is not a Tuple, `plot_prefix` is optional - plots/animations will be saved only
 passed.
 """
 function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
-    input = Dict_to_NamedTuple(input_dict_dfns["wall_pdf_neutral"])
-    if !(input.plot || input.animate || input.advection_velocity)
-        # nothing to do
-        return nothing
-    end
-    if !any(ri !== nothing for ri ∈ run_info)
-        println("Warning: no distribution function output, skipping wall_pdf plots")
-        return nothing
-    end
-
-    z_lower = 1
-    z_upper = run_info[1].z.n
-    if !all(ri.z.n == z_upper for ri ∈ run_info)
-        println("Cannot run plot_neutral_pdf_2D_at_wall() for runs with different "
-                * "z-grid sizes. Got $(Tuple(ri.z.n for ri ∈ run_info))")
-        return nothing
-    end
-
-    println("Making plots of neutral distribution function at walls")
-    flush(stdout)
+    try
+        input = Dict_to_NamedTuple(input_dict_dfns["wall_pdf_neutral"])
+        if !(input.plot || input.animate || input.advection_velocity)
+            # nothing to do
+            return nothing
+        end
+        if !any(ri !== nothing for ri ∈ run_info)
+            println("Warning: no distribution function output, skipping wall_pdf plots")
+            return nothing
+        end
 
-    has_rdim = any(ri !== nothing && ri.r.n > 1 for ri ∈ run_info)
-    has_zdim = any(ri !== nothing && ri.z.n > 1 for ri ∈ run_info)
-    is_1V = all(ri !== nothing && ri.vzeta.n == 1 && ri.vr.n == 1 for ri ∈ run_info)
-    moment_kinetic = any(ri !== nothing
-                         && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
-                         for ri ∈ run_info)
-    nt = minimum(ri.nt for ri ∈ run_info)
+        z_lower = 1
+        z_upper = run_info[1].z.n
+        if !all(ri.z.n == z_upper for ri ∈ run_info)
+            println("Cannot run plot_neutral_pdf_2D_at_wall() for runs with different "
+                    * "z-grid sizes. Got $(Tuple(ri.z.n for ri ∈ run_info))")
+            return nothing
+        end
 
-    for (z, z_range, label) ∈ ((z_lower, z_lower:z_lower+4, "wall-"),
-                               (z_upper, z_upper-4:z_upper, "wall+"))
-        f_neutral_input = copy(input_dict_dfns["f_neutral"])
-        f_neutral_input["iz0"] = z
+        println("Making plots of neutral distribution function at walls")
+        flush(stdout)
 
-        if input.plot
-            fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    plot_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
-                               label="$(run_label)iz=$iz", ax=ax)
-                end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz.pdf"
-            save(outfile, fig)
+        has_rdim = any(ri !== nothing && ri.r.n > 1 for ri ∈ run_info)
+        has_zdim = any(ri !== nothing && ri.z.n > 1 for ri ∈ run_info)
+        is_1V = all(ri !== nothing && ri.vzeta.n == 1 && ri.vr.n == 1 for ri ∈ run_info)
+        moment_kinetic = any(ri !== nothing
+                             && (ri.evolve_density || ri.evolve_upar || ri.evolve_ppar)
+                             for ri ∈ run_info)
+        nt = minimum(ri.nt for ri ∈ run_info)
 
-            fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    plot_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
-                               label="$(run_label)iz=$iz", ax=ax, yscale=log10,
-                               transform=(x)->positive_or_nan(x; epsilon=1.e-20))
-                end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "logpdf_neutral_$(label)_vs_vpa.pdf"
-            save(outfile, fig)
+        for (z, z_range, label) ∈ ((z_lower, z_lower:z_lower+4, "wall-"),
+                                   (z_upper, z_upper-4:z_upper, "wall+"))
+            f_neutral_input = copy(input_dict_dfns["f_neutral"])
+            f_neutral_input["iz0"] = z
 
-            if moment_kinetic
-                fig, ax = get_1d_ax(; xlabel="vz_unnorm", ylabel="f_neutral_unnorm")
+            if input.plot
+                fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
                 for iz ∈ z_range
                     for ri ∈ run_info
                         if length(run_info) > 1
@@ -4803,16 +4805,15 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        plot_f_unnorm_vs_vpa(ri; neutral=true, input=f_neutral_input,
-                                             is=1, iz=iz, label="$(run_label)iz=$iz",
-                                             ax=ax)
+                        plot_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
+                                   label="$(run_label)iz=$iz", ax=ax)
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "pdf_neutral_unnorm_$(label)_vs_vpa.pdf"
+                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz.pdf"
                 save(outfile, fig)
 
-                fig, ax = get_1d_ax(; xlabel="vz_unnorm", ylabel="f_neutral_unnorm")
+                fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
                 for iz ∈ z_range
                     for ri ∈ run_info
                         if length(run_info) > 1
@@ -4820,95 +4821,92 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        plot_f_unnorm_vs_vpa(ri; neutral=true, input=f_neutral_input,
-                                             is=1, iz=iz, label="$(run_label)iz=$iz",
-                                             ax=ax, yscale=log10,
-                                             transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        plot_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
+                                   label="$(run_label)iz=$iz", ax=ax, yscale=log10,
+                                   transform=(x)->positive_or_nan(x; epsilon=1.e-20))
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "logpdf_neutral_unnorm_$(label)_vs_vpa.pdf"
+                outfile=plot_prefix * "logpdf_neutral_$(label)_vs_vpa.pdf"
                 save(outfile, fig)
-            end
-
-            if !is_1V
-                plot_vs_vzeta_vr(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_vzeta.pdf")
-                plot_vs_vzeta_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vzeta.pdf")
-                plot_vs_vr_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                              outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vr.pdf")
-            end
 
-            if has_zdim
-                plot_vs_vz_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                             outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_z.pdf")
-            end
-
-            if has_zdim && !is_1V
-                plot_vs_vzeta_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_z.pdf")
-                plot_vs_vr_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                             outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_z.pdf")
-            end
-
-            if has_rdim && has_zdim
-                plot_vs_z_r(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                            outfile=plot_prefix * "pdf_neutral_$(label)_vs_z_r.pdf")
-            end
+                if moment_kinetic
+                    fig, ax = get_1d_ax(; xlabel="vz_unnorm", ylabel="f_neutral_unnorm")
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            plot_f_unnorm_vs_vpa(ri; neutral=true, input=f_neutral_input,
+                                                 is=1, iz=iz, label="$(run_label)iz=$iz",
+                                                 ax=ax)
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "pdf_neutral_unnorm_$(label)_vs_vpa.pdf"
+                    save(outfile, fig)
+
+                    fig, ax = get_1d_ax(; xlabel="vz_unnorm", ylabel="f_neutral_unnorm")
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            plot_f_unnorm_vs_vpa(ri; neutral=true, input=f_neutral_input,
+                                                 is=1, iz=iz, label="$(run_label)iz=$iz",
+                                                 ax=ax, yscale=log10,
+                                                 transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "logpdf_neutral_unnorm_$(label)_vs_vpa.pdf"
+                    save(outfile, fig)
+                end
 
-            if has_rdim
-                plot_vs_vz_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                             outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_r.pdf")
                 if !is_1V
-                    plot_vs_vzeta_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_r.pdf")
+                    plot_vs_vzeta_vr(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                     outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_vzeta.pdf")
+                    plot_vs_vzeta_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                     outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vzeta.pdf")
+                    plot_vs_vr_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                  outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vr.pdf")
+                end
 
-                    plot_vs_vr_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_r.pdf")
+                if has_zdim
+                    plot_vs_vz_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_z.pdf")
                 end
-            end
-        end
 
-        if input.animate
-            fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
-            frame_index = Observable(1)
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
-                    end
-                    animate_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
-                                  label="$(run_label)iz=$iz", ax=ax,
-                                  frame_index=frame_index)
+                if has_zdim && !is_1V
+                    plot_vs_vzeta_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_z.pdf")
+                    plot_vs_vr_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_z.pdf")
                 end
-            end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz." * input.animation_ext
-            save_animation(fig, frame_index, nt, outfile)
 
-            fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral", yscale=log10)
-            frame_index = Observable(1)
-            for iz ∈ z_range
-                for ri ∈ run_info
-                    if length(run_info) > 1
-                        run_label = ri.run_name * " "
-                    else
-                        run_label = ""
+                if has_rdim && has_zdim
+                    plot_vs_z_r(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                outfile=plot_prefix * "pdf_neutral_$(label)_vs_z_r.pdf")
+                end
+
+                if has_rdim
+                    plot_vs_vz_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_r.pdf")
+                    if !is_1V
+                        plot_vs_vzeta_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                        outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_r.pdf")
+
+                        plot_vs_vr_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                     outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_r.pdf")
                     end
-                    animate_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
-                                  label="$(run_label)iz=$iz", ax=ax,
-                                  frame_index=frame_index,
-                                  transform=(x)->positive_or_nan(x; epsilon=1.e-20))
                 end
             end
-            put_legend_right(fig, ax)
-            outfile=plot_prefix * "logpdf_neutral_$(label)_vs_vz." * input.animation_ext
-            save_animation(fig, frame_index, nt, outfile)
 
-            if moment_kinetic
+            if input.animate
                 fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
                 frame_index = Observable(1)
                 for iz ∈ z_range
@@ -4918,17 +4916,16 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        animate_f_unnorm_vs_vpa(ri; neutral=true, is=1, iz=iz,
-                                                input=f_neutral_input,
-                                                label="$(run_label)iz=$iz", ax=ax,
-                                                frame_index=frame_index)
+                        animate_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
+                                      label="$(run_label)iz=$iz", ax=ax,
+                                      frame_index=frame_index)
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "pdf_neutral_unnorm_$(label)_vs_vz." * input.animation_ext
+                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz." * input.animation_ext
                 save_animation(fig, frame_index, nt, outfile)
 
-                fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
+                fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral", yscale=log10)
                 frame_index = Observable(1)
                 for iz ∈ z_range
                     for ri ∈ run_info
@@ -4937,59 +4934,101 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
                         else
                             run_label = ""
                         end
-                        animate_f_unnorm_vs_vpa(ri; neutral=true, is=1, iz=iz,
-                                                input=f_neutral_input, label="$(run_label)iz=$iz",
-                                                ax=ax, frame_index=frame_index, yscale=log10,
-                                                transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        animate_vs_vz(ri, "f_neutral"; is=1, iz=iz, input=f_neutral_input,
+                                      label="$(run_label)iz=$iz", ax=ax,
+                                      frame_index=frame_index,
+                                      transform=(x)->positive_or_nan(x; epsilon=1.e-20))
                     end
                 end
                 put_legend_right(fig, ax)
-                outfile=plot_prefix * "logpdf_neutral_unnorm_$(label)_vs_vz." * input.animation_ext
+                outfile=plot_prefix * "logpdf_neutral_$(label)_vs_vz." * input.animation_ext
                 save_animation(fig, frame_index, nt, outfile)
-            end
 
-            if !is_1V
-                animate_vs_vzeta_vr(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_vzeta." * input.animation_ext)
-                animate_vs_vzeta_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vzeta." * input.animation_ext)
-                animate_vs_vr_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                 outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vr." * input.animation_ext)
-            end
+                if moment_kinetic
+                    fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
+                    frame_index = Observable(1)
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            animate_f_unnorm_vs_vpa(ri; neutral=true, is=1, iz=iz,
+                                                    input=f_neutral_input,
+                                                    label="$(run_label)iz=$iz", ax=ax,
+                                                    frame_index=frame_index)
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "pdf_neutral_unnorm_$(label)_vs_vz." * input.animation_ext
+                    save_animation(fig, frame_index, nt, outfile)
+
+                    fig, ax = get_1d_ax(; xlabel="vz", ylabel="f_neutral")
+                    frame_index = Observable(1)
+                    for iz ∈ z_range
+                        for ri ∈ run_info
+                            if length(run_info) > 1
+                                run_label = ri.run_name * " "
+                            else
+                                run_label = ""
+                            end
+                            animate_f_unnorm_vs_vpa(ri; neutral=true, is=1, iz=iz,
+                                                    input=f_neutral_input, label="$(run_label)iz=$iz",
+                                                    ax=ax, frame_index=frame_index, yscale=log10,
+                                                    transform=(x)->positive_or_nan(x; epsilon=1.e-20))
+                        end
+                    end
+                    put_legend_right(fig, ax)
+                    outfile=plot_prefix * "logpdf_neutral_unnorm_$(label)_vs_vz." * input.animation_ext
+                    save_animation(fig, frame_index, nt, outfile)
+                end
 
-            if has_zdim
-                animate_vs_vz_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_z." * input.animation_ext)
-            end
+                if !is_1V
+                    animate_vs_vzeta_vr(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                        outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_vzeta." * input.animation_ext)
+                    animate_vs_vzeta_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                        outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vzeta." * input.animation_ext)
+                    animate_vs_vr_vz(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                     outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_vr." * input.animation_ext)
+                end
 
-            if has_zdim && !is_1V
-                animate_vs_vzeta_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                                   outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_z." * input.animation_ext)
-                animate_vs_vr_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_z." * input.animation_ext)
-            end
+                if has_zdim
+                    animate_vs_vz_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_z." * input.animation_ext)
+                end
 
-            if has_rdim && has_zdim
-                animate_vs_z_r(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
-                               outfile=plot_prefix * "pdf_neutral_$(label)_vs_z_r." * input.animation_ext)
-            end
+                if has_zdim && !is_1V
+                    animate_vs_vzeta_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                       outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_z." * input.animation_ext)
+                    animate_vs_vr_z(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_z." * input.animation_ext)
+                end
 
-            if has_rdim
-                animate_vs_vz_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_r." * input.animation_ext)
-                if !is_1V
-                    animate_vs_vzeta_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                       outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_r." * input.animation_ext)
-                    animate_vs_vr_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
-                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_r." * input.animation_ext)
+                if has_rdim && has_zdim
+                    animate_vs_z_r(run_info, "f_neutral"; is=1, input=f_neutral_input, iz=z_range,
+                                   outfile=plot_prefix * "pdf_neutral_$(label)_vs_z_r." * input.animation_ext)
+                end
+
+                if has_rdim
+                    animate_vs_vz_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                    outfile=plot_prefix * "pdf_neutral_$(label)_vs_vz_r." * input.animation_ext)
+                    if !is_1V
+                        animate_vs_vzeta_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                           outfile=plot_prefix * "pdf_neutral_$(label)_vs_vzeta_r." * input.animation_ext)
+                        animate_vs_vr_r(run_info, "f_neutral"; is=1, input=f_neutral_input,
+                                        outfile=plot_prefix * "pdf_neutral_$(label)_vs_vr_r." * input.animation_ext)
+                    end
                 end
             end
-        end
 
-        if input.advection_velocity
-            animate_vs_vz(run_info, "neutral_vz_advect_speed"; is=1, input=f_neutral_input,
-                          outfile=plot_prefix * "neutral_vz_advect_speed_$(label)_vs_vz." * input.animation_ext)
+            if input.advection_velocity
+                animate_vs_vz(run_info, "neutral_vz_advect_speed"; is=1, input=f_neutral_input,
+                              outfile=plot_prefix * "neutral_vz_advect_speed_$(label)_vs_vz." * input.animation_ext)
+            end
         end
+    catch e
+        println("Error in plot_neutral_pdf_2D_at_wall(). Error was ", e)
     end
 
     return nothing

From bfd86c141fd759081ca03b5c4729e9fbafc26a80 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 3 Jun 2024 22:45:24 +0100
Subject: [PATCH 282/428] Fix restarting of electron initialisation

Need to remove output times that have already been passed, as having
times less than or equal to the initial time would mess up the I/O
logic. Slight hack to workaround the fact that when `t_params.electron`
is set up, we do not know if we are restarting from an existing
`*.initial_electron.h5` file or not.
---
 moment_kinetics/src/initial_conditions.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 5619ea764..d8338c028 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -640,6 +640,15 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                             t_params, t_params.debug_io[2], -1, nothing,
                                             "electron_debug")
         end
+        if code_time > 0.0
+            tind = searchsortedfirst(t_params.moments_output_times, code_time)
+            n_truncated = length(t_params.moments_output_times) - tind
+            truncated_times = t_params.moments_output_times[tind+1:end]
+            resize!(t_params.moments_output_times, n_truncated)
+            t_params.moments_output_times .= truncated_times
+            resize!(t_params.dfns_output_times, n_truncated)
+            t_params.dfns_output_times .= truncated_times
+        end
         electron_pseudotime =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,

From 05648ae2bb9db23ab1f445e060ae7157cd5474ef Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 4 Jun 2024 12:26:47 +0100
Subject: [PATCH 283/428] Fix factors of me_over_mi in electron external source
 term

---
 moment_kinetics/src/electron_kinetic_equation.jl |  3 ++-
 moment_kinetics/src/external_sources.jl          | 11 +++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 9522997af..2190c49a1 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1408,7 +1408,8 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
     if external_source_settings.electron.active
         external_electron_source!(fvec_out.pdf_electron, fvec_in.pdf_electron,
                                   moments.electron.dens, moments.electron.upar, moments,
-                                  external_source_settings.electron, vperp, vpa, dt)
+                                  composition, external_source_settings.electron, vperp,
+                                  vpa, dt)
     end
 
     if evolve_ppar
diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index f83a6a88b..066de0691 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -745,15 +745,18 @@ end
 Add external source term to the electron kinetic equation.
 """
 function external_electron_source!(pdf_out, pdf_in, electron_density, electron_upar,
-                                   moments, electron_source_settings, vperp, vpa, dt)
+                                   moments, composition, electron_source_settings, vperp,
+                                   vpa, dt)
     begin_r_z_vperp_region()
 
+    me_over_mi = composition.me_over_mi
+
     source_amplitude = moments.electron.external_source_amplitude
     source_T = electron_source_settings.source_T
     if vperp.n == 1
-        vth_factor = 1.0 / sqrt(source_T)
+        vth_factor = 1.0 / sqrt(source_T / me_over_mi)
     else
-        vth_factor = 1.0 / source_T^1.5
+        vth_factor = 1.0 / (source_T / me_over_mi)^1.5
     end
     vpa_grid = vpa.grid
     vperp_grid = vperp.grid
@@ -771,7 +774,7 @@ function external_electron_source!(pdf_out, pdf_in, electron_density, electron_u
             vpa_unnorm = vpa_grid[ivpa] * this_vth + this_upar
             pdf_out[ivpa,ivperp,iz,ir] +=
                 this_prefactor *
-                exp(-(vperp_unnorm^2 + vpa_unnorm^2) / source_T)
+                exp(-(vperp_unnorm^2 + vpa_unnorm^2) * me_over_mi / source_T)
         end
     end
 

From 229660e99535cf410d727394df318f3d141fa079 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 4 Jun 2024 18:06:19 +0100
Subject: [PATCH 284/428] No ion implicit options in electron timestepping
 input

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

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index e02363656..2c8c2cba8 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -401,8 +401,9 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      t_input["step_update_prefactor"], t_input["max_increase_factor"],
                      t_input["max_increase_factor_near_last_fail"],
                      t_input["last_fail_proximity_factor"], t_input["minimum_dt"],
-                     t_input["maximum_dt"], t_input["implicit_ion_advance"],
-                     t_input["implicit_vpa_advection"],
+                     t_input["maximum_dt"],
+                     electron !== nothing && t_input["implicit_ion_advance"],
+                     electron !== nothing && t_input["implicit_vpa_advection"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
                      t_input["split_operators"], t_input["steady_state_residual"],
                      t_input["converged_residual_value"],

From 3b9a54ad14200072189b42f0bb938934f8119516 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 6 May 2024 22:35:16 +0100
Subject: [PATCH 285/428] Ensure parallel_friction is always initialised

---
 moment_kinetics/src/time_advance.jl | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 2c8c2cba8..06504b02d 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -692,6 +692,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     # computing the electrostatic potential (and components of the electric field)
     calculate_electron_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
                                            electron_mom_diss_coeff, composition.electron_physics)
+    # calculate the electron-ion parallel friction force
+    calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+        moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+        composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
     # initialize the electrostatic potential
     begin_serial_region()
     update_phi!(fields, scratch[1], vperp, z, r, composition, collisions, moments, z_spectral, r_spectral, scratch_dummy, gyroavs)
@@ -917,10 +921,6 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         if composition.electron_physics == braginskii_fluid
             electron_fluid_qpar_boundary_condition!(moments.electron, z)
         end
-        # calculate the electron-ion parallel friction force
-        calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
-            moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
-            composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
         # update the electron moment entries in the scratch array
         begin_r_z_region()
         @loop_r_z ir iz begin
@@ -938,6 +938,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
             scratch[1].pz_neutral[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
         end
     end
+    # calculate the electron-ion parallel friction force
+    calculate_electron_parallel_friction_force!(moments.electron.parallel_friction, moments.electron.dens,
+        moments.electron.upar, moments.ion.upar, moments.electron.dT_dz,
+        composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
 
     calculate_ion_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
                                       ion_mom_diss_coeff)

From 5116852d2e728e247f10e2a88fc4c73e8daaa40e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 7 May 2024 17:11:28 +0100
Subject: [PATCH 286/428] Use correct RK-stage-updated variables for electron
 qpar bc

---
 moment_kinetics/src/electron_fluid_equations.jl | 17 ++++++++++-------
 moment_kinetics/src/initial_conditions.jl       |  8 ++++++--
 moment_kinetics/src/time_advance.jl             |  4 +++-
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 108145a02..cf6f4fb90 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -135,7 +135,10 @@ function calculate_electron_moments!(scratch, pdf, moments, composition, collisi
                              scratch.electron_upar, scratch.upar, collisions.nu_ei,
                              composition.me_over_mi, composition.electron_physics, vpa)
     if composition.electron_physics == braginskii_fluid
-        electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        electron_fluid_qpar_boundary_condition!(scratch.electron_ppar,
+                                                scratch.electron_upar,
+                                                scratch.electron_density,
+                                                moments.electron, z)
     end
     return nothing
 end
@@ -421,7 +424,7 @@ end
 Impose fluid approximation to electron sheath boundary condition on the parallel heat
 flux. See Stangeby textbook, equations (2.89) and (2.90).
 """
-function electron_fluid_qpar_boundary_condition!(electron_moments, z)
+function electron_fluid_qpar_boundary_condition!(ppar, upar, dens, electron_moments, z)
     begin_r_region()
 
     if z.irank == 0 && (z.irank == z.nrank - 1)
@@ -436,10 +439,10 @@ function electron_fluid_qpar_boundary_condition!(electron_moments, z)
 
     @loop_r ir begin
         for iz ∈ z_indices
-            ppar = electron_moments.ppar[iz,ir]
-            upar = electron_moments.upar[iz,ir]
-            dens = electron_moments.dens[iz,ir]
-            particle_flux = dens * upar
+            this_ppar = ppar[iz,ir]
+            this_upar = electron_moments.upar[iz,ir]
+            this_dens = electron_moments.dens[iz,ir]
+            particle_flux = this_dens * this_upar
             T_e = electron_moments.temp[iz,ir]
 
             # Stangeby (2.90)
@@ -450,7 +453,7 @@ function electron_fluid_qpar_boundary_condition!(electron_moments, z)
 
             # E.g. Helander&Sigmar (2.14), neglecting electron viscosity and kinetic
             # energy fluxes due to small mass ratio
-            conductive_heat_flux = total_heat_flux - 2.5 * ppar * upar
+            conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar
 
             electron_moments.qpar[iz,ir] = conductive_heat_flux
         end
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index d8338c028..d64a64e62 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -337,7 +337,9 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
     if composition.electron_physics == braginskii_fluid
-        electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        electron_fluid_qpar_boundary_condition!(
+            moments.electron.ppar, moments.electron.upar, moments.electron.dens,
+            moments.electron, z)
         if restart_electron_physics ∉ (nothing, braginskii_fluid, kinetic_electrons)
             # Restarting from Boltzmann. If we use an exactly constant T_e profile,
             # qpar for the electrons will be non-zero only at the boundary points,
@@ -413,7 +415,9 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         moments.electron.upar, moments.ion.upar, collisions.nu_ei, composition.me_over_mi,
         composition.electron_physics, vpa)
     if composition.electron_physics == braginskii_fluid
-        electron_fluid_qpar_boundary_condition!(moments.electron, z)
+        electron_fluid_qpar_boundary_condition!(
+            moments.electron.ppar, moments.electron.upar, moments.electron.dens,
+            moments.electron, z)
     end
     # calculate the zed derivative of the initial electron parallel heat flux
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 06504b02d..77e40d2d1 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -919,7 +919,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
             moments.electron.upar, moments.ion.upar, collisions.nu_ei,
             composition.me_over_mi, composition.electron_physics, vpa)
         if composition.electron_physics == braginskii_fluid
-            electron_fluid_qpar_boundary_condition!(moments.electron, z)
+            electron_fluid_qpar_boundary_condition!(
+                moments.electron.ppar, moments.electron.upar, moments.electron.dens,
+                moments.electron, z)
         end
         # update the electron moment entries in the scratch array
         begin_r_z_region()

From 16a570a2143cfe5dbe25ce15442f454c2202ca0f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 5 Jun 2024 09:11:22 +0100
Subject: [PATCH 287/428] Fix Braginskii electron restart from Boltzmann

`derivative_z!()` function must not be called inside a `@serial_region`.
---
 moment_kinetics/src/initial_conditions.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index d64a64e62..fb1945302 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -402,12 +402,12 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                 @. moments.electron.vth = sqrt(moments.electron.temp /
                                                composition.me_over_mi)
                 @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
-                @views derivative_z!(moments.electron.dT_dz, moments.electron.temp,
-                                     scratch_dummy.buffer_rs_1[:,1],
-                                     scratch_dummy.buffer_rs_2[:,1],
-                                     scratch_dummy.buffer_rs_3[:,1],
-                                     scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
             end
+            @views derivative_z!(moments.electron.dT_dz, moments.electron.temp,
+                                 scratch_dummy.buffer_rs_1[:,1],
+                                 scratch_dummy.buffer_rs_2[:,1],
+                                 scratch_dummy.buffer_rs_3[:,1],
+                                 scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
         end
     end
     moments.electron.qpar_updated[] = false

From 57c241164542ec8e3e9a743e67f26f8f3e5ae0a0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 5 May 2024 18:57:52 +0100
Subject: [PATCH 288/428] Implicit solve for Braginskii electron conduction,
 when using an IMEX solver

---
 .../src/electron_fluid_equations.jl           | 125 +++++++++++++++++-
 moment_kinetics/src/input_structs.jl          |   2 +
 moment_kinetics/src/moment_kinetics_input.jl  |   1 +
 moment_kinetics/src/time_advance.jl           |  68 +++++++++-
 4 files changed, 186 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index cf6f4fb90..b181e7c84 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -10,10 +10,13 @@ export calculate_electron_qpar_from_pdf!
 export update_electron_vth_temperature!
 
 using ..communication
+using ..derivatives: derivative_z!
 using ..looping
 using ..input_structs: boltzmann_electron_response_with_simple_sheath
 using ..input_structs: braginskii_fluid, kinetic_electrons
-using ..moment_kinetics_structs: electron_pdf_substruct
+using ..moment_kinetics_structs: electron_pdf_substruct, moments_electron_substruct
+using ..nonlinear_solvers
+using ..type_definitions: mk_float
 using ..velocity_moments: integrate_over_vspace
 
 using MPI
@@ -152,7 +155,8 @@ an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron_upar,
                                    ion_upar, ion_ppar, density_neutral, uz_neutral,
                                    pz_neutral, moments, collisions, dt, composition,
-                                   electron_source_settings, num_diss_params, z)
+                                   electron_source_settings, num_diss_params, z;
+                                   conduction=true)
     begin_r_z_region()
     # define some abbreviated variables for convenient use in rest of function
     me_over_mi = composition.me_over_mi
@@ -161,9 +165,13 @@ function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron
     # arising from derivatives of ppar, qpar and upar
     @loop_r_z ir iz begin
         ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
-                               + moments.dqpar_dz[iz,ir]
                                + 3*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
     end
+    if conduction
+        @loop_r_z ir iz begin
+            ppar_out[iz,ir] -= dt*moments.dqpar_dz[iz,ir]
+        end
+    end
     # @loop_r_z ir iz begin
     #     ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
     #                 + (2/3)*moments.dqpar_dz[iz,ir]
@@ -224,6 +232,106 @@ function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron
     return nothing
 end
 
+"""
+Add just the braginskii conduction contribution to the electron pressure, and assume that
+we have to calculate qpar and dqpar_dz from ppar within this function (they are not
+pre-calculated).
+"""
+function electron_braginskii_conduction!(ppar_out::AbstractVector{mk_float},
+                                         ppar_in::AbstractVector{mk_float},
+                                         dens::AbstractVector{mk_float},
+                                         upar_e::AbstractVector{mk_float},
+                                         upar_i::AbstractVector{mk_float},
+                                         electron_moments, collisions, composition, z,
+                                         z_spectral, scratch_dummy, dt, ir)
+
+    buffer_r_1 = @view scratch_dummy.buffer_rs_1[ir,1]
+    buffer_r_2 = @view scratch_dummy.buffer_rs_2[ir,1]
+    buffer_r_3 = @view scratch_dummy.buffer_rs_3[ir,1]
+    buffer_r_4 = @view scratch_dummy.buffer_rs_4[ir,1]
+
+    temp = @view electron_moments.temp[:,ir]
+    dT_dz = @view electron_moments.dT_dz[:,ir]
+    qpar = @view electron_moments.qpar[:,ir]
+    dqpar_dz = @view electron_moments.dqpar_dz[:,ir]
+
+    update_electron_temperature!(temp, ppar_in, dens, composition)
+    derivative_z!(dT_dz, temp, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4, z_spectral,
+                  z)
+    electron_moments.qpar_updated[] = false
+    calculate_electron_qpar!(electron_moments, nothing, ppar_in, upar_e, upar_i,
+                             collisions.nu_ei, composition.me_over_mi,
+                             composition.electron_physics, nothing)
+    electron_fluid_qpar_boundary_condition!(ppar_in, upar_e, dens, electron_moments, z)
+    derivative_z!(dqpar_dz, qpar, buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                  z_spectral, z)
+
+    @loop_r_z ir iz begin
+        ppar_out[iz,ir] -= dt*electron_moments.dqpar_dz[iz,ir]
+    end
+
+    return nothing
+end
+
+function implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt, z_spectral,
+                                         composition, collisions, scratch_dummy,
+                                         nl_solver_params)
+    begin_z_region()
+
+    for ir ∈ 1:r.n
+        ppar_out = @view fvec_out.electron_ppar[:,ir]
+        ppar_in = @view fvec_in.electron_ppar[:,ir]
+        dens = @view fvec_in.electron_density[:,ir]
+        upar_e = @view fvec_in.electron_upar[:,ir]
+        upar_i = @view fvec_in.upar[:,ir]
+
+        # Explicit timestep to give initial guess for implicit solve
+        electron_braginskii_conduction!(ppar_out, ppar_in, dens, upar_e, upar_i,
+                                        moments.electron, collisions, composition, z,
+                                        z_spectral, scratch_dummy, dt, ir)
+
+        # Define a function whose input is `electron_ppar`, so that when it's output
+        # `residual` is zero, electron_ppar is the result of a backward-Euler timestep:
+        #   (f_new - f_old) / dt = RHS(f_new)
+        # ⇒ (f_new - f_old)/dt - RHS(f_new) = 0
+        function residual_func!(residual, electron_ppar)
+            begin_z_region()
+            @loop_z iz begin
+                residual[iz] = ppar_in[iz]
+            end
+            electron_braginskii_conduction!(residual, electron_ppar, dens, upar_e,
+                                            upar_i, moments.electron, collisions,
+                                            composition, z, z_spectral, scratch_dummy,
+                                            dt, ir)
+            # Now
+            #   residual = f_old + dt*RHS(f_new)
+            # so update to desired residual
+            begin_z_region()
+            @loop_z iz begin
+                residual[iz] = (electron_ppar[iz] - residual[iz])
+            end
+        end
+
+        # Shared-memory buffers
+        residual = @view scratch_dummy.buffer_zs_1[:,1]
+        delta_x = @view scratch_dummy.buffer_zs_2[:,1]
+        rhs_delta = @view scratch_dummy.buffer_zs_3[:,1]
+        v = @view scratch_dummy.buffer_zs_4[:,1]
+        w = @view scratch_dummy.buffer_zrs_1[:,1,1]
+
+        success = newton_solve!(ppar_out, residual_func!, residual, delta_x, rhs_delta, v,
+                                w, nl_solver_params; left_preconditioner=nothing,
+                                right_preconditioner=nothing, coords=(z=z,))
+        if !success
+            return success
+        end
+    end
+
+    nl_solver_params.stage_counter[] += 1
+
+    return true
+end
+
 """
 solve the electron force balance (parallel momentum) equation for the
 parallel electric field, Epar:
@@ -418,6 +526,17 @@ function update_electron_vth_temperature!(moments, ppar, dens, composition)
     return nothing
 end
 
+function update_electron_temperature!(temp, ppar, dens, composition)
+    begin_z_region()
+
+    @loop_z iz begin
+        p = max(ppar[iz], 0.0)
+        temp[iz] = 2 * p / dens[iz]
+    end
+
+    return nothing
+end
+
 """
     electron_fluid_qpar_boundary_condition!(electron_moments, z)
 
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index e8973de4d..acf7d0b26 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -79,6 +79,7 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     last_fail_proximity_factor::mk_float
     minimum_dt::mk_float
     maximum_dt::mk_float
+    implicit_braginskii_conduction::Bool
     implicit_ion_advance::Bool
     implicit_vpa_advection::Bool
     write_after_fixed_step_count::Bool
@@ -121,6 +122,7 @@ mutable struct advance_info
     force_balance::Bool
     energy::Bool
     electron_energy::Bool
+    electron_conduction::Bool
     neutral_external_source::Bool
     neutral_source_terms::Bool
     neutral_continuity::Bool
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 2ce15e5c1..d8bcd18eb 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -221,6 +221,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         last_fail_proximity_factor=1.05,
         minimum_dt=0.0,
         maximum_dt=Inf,
+        implicit_braginskii_conduction=true,
         implicit_ion_advance=true,
         implicit_vpa_advection=false,
         write_after_fixed_step_count=false,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 77e40d2d1..c58c278fb 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -81,7 +81,9 @@ using ..electron_fluid_equations: calculate_electron_density!
 using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservation!
 using ..electron_fluid_equations: calculate_electron_qpar!, electron_fluid_qpar_boundary_condition!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
-using ..electron_fluid_equations: electron_energy_equation!, update_electron_vth_temperature!
+using ..electron_fluid_equations: electron_energy_equation!, update_electron_vth_temperature!,
+                                  electron_braginskii_conduction!,
+                                  implicit_braginskii_conduction!
 using ..input_structs: braginskii_fluid
 using ..derivatives: derivative_z!
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
@@ -359,6 +361,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
 
     if rk_coefs_implicit === nothing
         # Not an IMEX scheme, so cannot have any implicit terms
+        t_input["implicit_braginskii_conduction"] = false
         t_input["implicit_ion_advance"] = false
         t_input["implicit_vpa_advection"] = false
     end
@@ -402,6 +405,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      t_input["max_increase_factor_near_last_fail"],
                      t_input["last_fail_proximity_factor"], t_input["minimum_dt"],
                      t_input["maximum_dt"],
+                     electron !== nothing && t_input["implicit_braginskii_conduction"],
                      electron !== nothing && t_input["implicit_ion_advance"],
                      electron !== nothing && t_input["implicit_vpa_advection"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
@@ -597,6 +601,14 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
 
     # Set up parameters for Jacobian-free Newton-Krylov solver used for implicit part of
     # timesteps.
+    if t_params.implicit_braginskii_conduction
+        # Should really have options to set solver tolerance, etc.
+        electron_conduction_nl_solve_parameters = setup_nonlinear_solve(input_dict, (z=z,);
+                                                                        default_rtol=t_params.rtol / 10.0,
+                                                                        default_atol=t_params.atol / 10.0)
+    else
+       electron_conduction_nl_solve_parameters = nothing
+    end
     if t_params.implicit_ion_advance
         # Implicit solve for vpa_advection term should be done in serial, as it will be
         # called within a parallelised s_r_z_vperp loop.
@@ -628,7 +640,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         error("Cannot use implicit_ion_advance and implicit_vpa_advection at the same "
               * "time")
     end
-    nl_solver_params = (ion_advance=nl_solver_ion_advance_params,
+    nl_solver_params = (electron_conduction=electron_conduction_nl_solve_parameters,
+                        ion_advance=nl_solver_ion_advance_params,
                         vpa_advection=nl_solver_vpa_advection_params,)
 
     begin_serial_region()
@@ -996,6 +1009,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
     advance_force_balance = false
     advance_energy = false
     advance_electron_energy = false
+    advance_electron_conduction = false
     advance_neutral_z_advection = false
     advance_neutral_r_advection = false
     advance_neutral_vz_advection = false
@@ -1115,8 +1129,21 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         end
         # if treating the electrons as a fluid with Braginskii closure, or
         # moment-kinetically then advance the electron energy equation
-        if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        if composition.electron_physics == kinetic_electrons
             advance_electron_energy = true
+            advance_electron_conduction = true
+        elseif composition.electron_physics == braginskii_fluid
+            if t_params.implicit_braginskii_conduction
+                # if treating the electrons as a fluid with Braginskii closure, and using
+                # an IMEX scheme, advance the conduction part of the electron energy
+                # equation implicitly.
+                advance_electron_energy = true
+                advance_electron_conduction = false
+            else
+                # If not using an IMEX scheme, treat the conduction explicitly.
+                advance_electron_energy = true
+                advance_electron_conduction = true
+            end
         end
 
         # *_diffusion flags are set regardless of whether diffusion is included in explicit or
@@ -1146,7 +1173,8 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                         advance_external_source, advance_ion_numerical_dissipation,
                         advance_neutral_numerical_dissipation, advance_sources,
                         advance_continuity, advance_force_balance, advance_energy,
-                        advance_electron_energy, advance_neutral_external_source,
+                        advance_electron_energy, advance_electron_conduction,
+                        advance_neutral_external_source,
                         advance_neutral_sources, advance_neutral_continuity,
                         advance_neutral_force_balance, advance_neutral_energy,
                         manufactured_solns_test, r_diffusion, vpa_diffusion,
@@ -1185,6 +1213,7 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     advance_force_balance = false
     advance_energy = false
     advance_electron_energy = false
+    advance_electron_conduction = false
     advance_neutral_z_advection = false
     advance_neutral_r_advection = false
     advance_neutral_vz_advection = false
@@ -1255,6 +1284,13 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
     vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0)
 
+    if t_params.implicit_braginskii_conduction
+        # if treating the electrons as a fluid with Braginskii closure, and using an IMEX
+        # scheme, advance the conduction part of the electron energy equation implicitly.
+        advance_electron_energy = false
+        advance_electron_conduction = true
+    end
+
     manufactured_solns_test = manufactured_solns_input.use_for_advance
 
     return advance_info(advance_vpa_advection, advance_vperp_advection, advance_z_advection, advance_r_advection,
@@ -1267,8 +1303,8 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
                         explicit_weakform_fp_collisions,
                         advance_external_source, advance_ion_numerical_dissipation,
                         advance_neutral_numerical_dissipation, advance_sources,
-                        advance_continuity, advance_electron_energy,
-                        advance_force_balance, advance_energy,
+                        advance_continuity, advance_force_balance, advance_energy,
+                        advance_electron_energy, advance_electron_conduction,
                         advance_neutral_external_source, advance_neutral_sources,
                         advance_neutral_continuity, advance_neutral_force_balance,
                         advance_neutral_energy, manufactured_solns_test, r_diffusion,
@@ -3064,7 +3100,18 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                   fvec_in.ppar, fvec_in.density_neutral,
                                   fvec_in.uz_neutral, fvec_in.pz_neutral,
                                   moments.electron, collisions, dt, composition,
-                                  external_source_settings.electron, num_diss_params, z)
+                                  external_source_settings.electron, num_diss_params, z;
+                                  conduction=advance.electron_conduction)
+    elseif advance.electron_conduction
+        # Explicit version of the implicit part of the IMEX timestep, need to evaluate
+        # only the conduction term.
+        for ir ∈ 1:r.n
+            @views electron_braginskii_conduction!(
+                fvec_out.electron_ppar[:,ir], fvec_in.electron_ppar[:,ir],
+                fvec_in.electron_density[:,ir], fvec_in.electron_upar[:,ir],
+                fvec_in.upar[:,ir], moments.electron, collisions, composition, z,
+                z_spectral, scratch_dummy, dt, ir)
+        end
     end
     # reset "xx.updated" flags to false since ff has been updated
     # and the corresponding moments have not
@@ -3084,6 +3131,13 @@ function backward_euler!(fvec_out, fvec_in, pdf, fields, moments, advect_objects
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
+    if advance.electron_conduction
+        success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,
+                                                  z_spectral, composition, collisions,
+                                                  scratch_dummy,
+                                                  nl_solver_params.electron_conduction)
+    end
+
     if nl_solver_params.ion_advance !== nothing
         success = implicit_ion_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                         advect_objects, vz, vr, vzeta, vpa, vperp,

From 3ff61115d93d2915361547ce9dd181ff9a9c0d7f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 7 May 2024 18:50:18 +0100
Subject: [PATCH 289/428] Example input file using IMEX solver

---
 ...on0.5_split3_braginskii-vpadiss0-IMEX.toml | 111 ++++++++++++++++++
 1 file changed, 111 insertions(+)
 create mode 100644 examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_braginskii-vpadiss0-IMEX.toml

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_braginskii-vpadiss0-IMEX.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_braginskii-vpadiss0-IMEX.toml
new file mode 100644
index 000000000..c9230a33a
--- /dev/null
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_braginskii-vpadiss0-IMEX.toml
@@ -0,0 +1,111 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "braginskii_fluid"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 0.2 # 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "gaussian"
+z_IC_density_amplitude1 = 0.001
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "gaussian"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = -1.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.5
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 5
+z_nelement = 32
+#z_nelement_local = 16
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+z_element_spacing_option = "sqrt"
+vpa_ngrid = 6
+vpa_nelement = 63
+vpa_L = 36.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 63
+vz_L = 36.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "KennedyCarpenterARK324"
+#type = "KennedyCarpenterARK324-explicit"
+#type = "KennedyCarpenterARK437"
+#type = "fake_imex"
+#nstep = 50000
+implicit_ion_advance = false
+implicit_vpa_advection = false
+implicit_braginskii_conduction = true
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-9
+rtol = 1.0e-7
+#rtol = 1.0e-9
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+#write_after_fixed_step_count = true
+
+[ion_source]
+active = true
+z_profile = "gaussian"
+z_width = 0.125
+source_strength = 2.0
+source_T = 2.0
+
+[ion_numerical_dissipation]
+#vpa_dissipation_coefficient = 1.0e-1
+#vpa_dissipation_coefficient = 1.0e-2
+#vpa_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient = 2.0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+#vz_dissipation_coefficient = 1.0e-1
+#vz_dissipation_coefficient = 1.0e-2
+#vz_dissipation_coefficient = 1.0e-3
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true

From 282c3aea8b4e561faff3b3ac3c00468faa6b852a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 6 Jun 2024 18:05:23 +0100
Subject: [PATCH 290/428] Allow periodic boundary conditions for Braginskii
 electrons

---
 moment_kinetics/src/electron_fluid_equations.jl | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index b181e7c84..8efe49368 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -544,6 +544,12 @@ Impose fluid approximation to electron sheath boundary condition on the parallel
 flux. See Stangeby textbook, equations (2.89) and (2.90).
 """
 function electron_fluid_qpar_boundary_condition!(ppar, upar, dens, electron_moments, z)
+    if z.bc == "periodic"
+        # Nothing to do as z-derivative used to calculate qpar already imposed
+        # periodicity.
+        return nothing
+    end
+
     begin_r_region()
 
     if z.irank == 0 && (z.irank == z.nrank - 1)

From ea239d86b420d376836152079dac005a669a0b63 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 6 Jun 2024 21:53:05 +0100
Subject: [PATCH 291/428] Example with Braginskii electrons in periodic domain

---
 .../periodic_split3_braginskii-IMEX.toml      | 110 ++++++++++++++++++
 1 file changed, 110 insertions(+)
 create mode 100644 examples/kinetic-electrons/periodic_split3_braginskii-IMEX.toml

diff --git a/examples/kinetic-electrons/periodic_split3_braginskii-IMEX.toml b/examples/kinetic-electrons/periodic_split3_braginskii-IMEX.toml
new file mode 100644
index 000000000..4dfe44807
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_braginskii-IMEX.toml
@@ -0,0 +1,110 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "braginskii_fluid"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "KennedyCarpenterARK324"
+implicit_ion_advance = false
+implicit_vpa_advection = false
+implicit_braginskii_conduction = true
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-7
+rtol = 1.0e-7
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+#[timestepping]
+#type = "KennedyCarpenterARK324"
+#implicit_ion_advance = false
+#implicit_vpa_advection = false
+#implicit_braginskii_conduction = true
+#nstep = 1000000
+#dt = 1.0e-6
+#minimum_dt = 1.0e-7
+#rtol = 1.0e-7
+#nwrite = 100
+#nwrite_dfns = 100
+#steady_state_residual = true
+#converged_residual_value = 1.0e-3
+#write_after_fixed_step_count = true
+#write_error_diagnostics = true
+#write_steady_state_diagnostics = true
+
+[nonlinear_solver]
+nonlinear_max_iterations = 100
+#rtol = 1.0e-9
+#atol = 1.0e-12
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true

From 12bd861ea873a5a8edcd3d42db9047b5407d2f21 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 4 Jun 2024 18:33:39 +0100
Subject: [PATCH 292/428] Function to load electron moments data

---
 moment_kinetics/src/load_data.jl | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 1aff1b581..9051eee8c 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -5,6 +5,7 @@ module load_data
 export open_readonly_output_file
 export load_fields_data
 export load_ion_moments_data
+export load_electron_moments_data
 export load_neutral_particle_moments_data
 export load_pdf_data
 export load_neutral_pdf_data
@@ -508,6 +509,31 @@ function load_ion_moments_data(fid; printout=false, extended_moments = false)
     end
 end
 
+"""
+"""
+function load_electron_moments_data(fid; printout=false)
+    if printout
+        print("Loading electron velocity moments data...")
+    end
+
+    group = get_group(fid, "dynamic_data")
+
+    # Read electron parallel pressure
+    parallel_pressure = load_variable(group, "electron_parallel_pressure")
+
+    # Read electron parallel heat flux
+    parallel_heat_flux = load_variable(group, "electron_parallel_heat_flux")
+
+    # Read electron thermal speed
+    thermal_speed = load_variable(group, "electron_thermal_speed")
+
+    if printout
+        println("done.")
+    end
+
+    return parallel_pressure, parallel_heat_flux, thermal_speed
+end
+
 function load_neutral_particle_moments_data(fid; printout=false)
     if printout
         print("Loading neutral particle velocity moments data...")

From ee3921727c3210f4fa0bfc9717a0ecc1c1ed0b04 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 8 May 2024 12:16:38 +0100
Subject: [PATCH 293/428] Test case with Braginskii electrons

---
 .../test/braginskii_electrons_imex_tests.jl   | 312 ++++++++++++++++++
 moment_kinetics/test/runtests.jl              |   1 +
 2 files changed, 313 insertions(+)
 create mode 100644 moment_kinetics/test/braginskii_electrons_imex_tests.jl

diff --git a/moment_kinetics/test/braginskii_electrons_imex_tests.jl b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
new file mode 100644
index 000000000..97f0a4f33
--- /dev/null
+++ b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
@@ -0,0 +1,312 @@
+module BraginskiiElectronsIMEX
+
+# Regression test using wall boundary conditions, with recycling fraction less than 1 and
+# a plasma source. Runs for a while and then checks phi profile against saved reference
+# output.
+
+include("setup.jl")
+
+using Base.Filesystem: tempname
+using MPI
+
+using moment_kinetics.coordinates: define_coordinate
+using moment_kinetics.input_structs: grid_input, advection_input
+using moment_kinetics.interpolation: interpolate_to_grid_z
+using moment_kinetics.load_data: open_readonly_output_file
+using moment_kinetics.load_data: load_electron_moments_data
+
+# default inputs for tests
+test_input = Dict("n_ion_species" => 1,
+                  "n_neutral_species" => 1,
+                  "electron_physics" => "braginskii_fluid",
+                  "run_name" => "braginskii-electrons-imex",
+                  "evolve_moments_density" => true,
+                  "evolve_moments_parallel_flow" => true,
+                  "evolve_moments_parallel_pressure" => true,
+                  "evolve_moments_conservation" => true,
+                  "T_e" => 0.2,
+                  "nu_ei" => 1.0e3,
+                  "initial_density1" => 1.0,
+                  "initial_temperature1" => 1.0,
+                  "z_IC_option1" => "sinusoid",
+                  "z_IC_density_amplitude1" => 0.1,
+                  "z_IC_density_phase1" => 0.0,
+                  "z_IC_upar_amplitude1" => 1.0,
+                  "z_IC_upar_phase1" => 0.0,
+                  "z_IC_temperature_amplitude1" => 0.1,
+                  "z_IC_temperature_phase1" => 1.0,
+                  "vpa_IC_option1" => "gaussian",
+                  "vpa_IC_density_amplitude1" => 1.0,
+                  "vpa_IC_density_phase1" => 0.0,
+                  "vpa_IC_upar_amplitude1" => 0.0,
+                  "vpa_IC_upar_phase1" => 0.0,
+                  "vpa_IC_temperature_amplitude1" => 0.0,
+                  "vpa_IC_temperature_phase1" => 0.0,
+                  "initial_density2" => 1.0,
+                  "initial_temperature2" => 1.0,
+                  "z_IC_option2" => "sinusoid",
+                  "z_IC_density_amplitude2" => 0.001,
+                  "z_IC_density_phase2" => 0.0,
+                  "z_IC_upar_amplitude2" => 0.0,
+                  "z_IC_upar_phase2" => 0.0,
+                  "z_IC_temperature_amplitude2" => 0.0,
+                  "z_IC_temperature_phase2" => 0.0,
+                  "vpa_IC_option2" => "gaussian",
+                  "vpa_IC_density_amplitude2" => 1.0,
+                  "vpa_IC_density_phase2" => 0.0,
+                  "vpa_IC_upar_amplitude2" => 0.0,
+                  "vpa_IC_upar_phase2" => 0.0,
+                  "vpa_IC_temperature_amplitude2" => 0.0,
+                  "vpa_IC_temperature_phase2" => 0.0,
+                  "charge_exchange_frequency" => 0.75,
+                  "ionization_frequency" => 0.5,
+                  "constant_ionization_rate" => false,
+                  "timestepping" => Dict{String,Any}("type" => "KennedyCarpenterARK324",
+                                                     "implicit_ion_advance" => false,
+                                                     "implicit_vpa_advection" => false,
+                                                     "nstep" => 10000,
+                                                     "dt" => 1.0e-6,
+                                                     "minimum_dt" => 1.e-7,
+                                                     "rtol" => 1.0e-7,
+                                                     "nwrite" => 10000,
+                                                     "high_precision_error_sum" => true),
+                  "nonlinear_solver" => Dict{String,Any}("nonlinear_max_iterations" => 100),
+                  "r_ngrid" => 1,
+                  "r_nelement" => 1,
+                  "z_ngrid" => 17,
+                  "z_nelement" => 16,
+                  "z_bc" => "periodic",
+                  "z_discretization" => "chebyshev_pseudospectral",
+                  "vpa_ngrid" => 6,
+                  "vpa_nelement" => 31,
+                  "vpa_L" => 12.0,
+                  "vpa_bc" => "zero",
+                  "vpa_discretization" => "chebyshev_pseudospectral",
+                  "vz_ngrid" => 6,
+                  "vz_nelement" => 31,
+                  "vz_L" => 12.0,
+                  "vz_bc" => "zero",
+                  "vz_discretization" => "chebyshev_pseudospectral",
+                  "ion_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0,
+                                                                  "vpa_dissipation_coefficient" => 1e0),
+                  "neutral_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0,
+                                                                      "vz_dissipation_coefficient" => 1e-1))
+
+if global_size[] > 2 && global_size[] % 2 == 0
+    # Test using distributed-memory
+    test_input["z_nelement_local"] = test_input["z_nelement"] ÷ 2
+end
+
+"""
+Run a test for a single set of parameters
+"""
+function run_test(test_input, expected_p, expected_q, expected_vt; rtol=4.e-14,
+                  atol=1.e-15, args...)
+    # by passing keyword arguments to run_test, args becomes a Tuple of Pairs which can be
+    # used to update the default inputs
+
+    # Make a copy to make sure nothing modifies the input Dicts defined in this test
+    # script.
+    test_input = deepcopy(test_input)
+
+    # Convert keyword arguments to a unique name
+    name = test_input["run_name"]
+    if length(args) > 0
+        name = string(name, "_", (string(k, "-", v, "_") for (k, v) in args)...)
+
+        # Remove trailing "_"
+        name = chop(name)
+    end
+
+    # Provide some progress info
+    println("    - testing ", name)
+
+    # Convert dict from symbol keys to String keys
+    modified_inputs = Dict(String(k) => v for (k, v) in args)
+
+    # Update default inputs with values to be changed
+    input = merge(test_input, modified_inputs)
+
+    input["run_name"] = name
+
+    # Suppress console output while running
+    p = undef
+    q = undef
+    vt = undef
+    quietoutput() do
+        # run simulation
+        run_moment_kinetics(input)
+    end
+
+    if global_rank[] == 0
+        quietoutput() do
+            # Load and analyse output
+            #########################
+
+            path = joinpath(realpath(input["base_directory"]), name, name)
+
+            # open the netcdf file and give it the handle 'fid'
+            fid = open_readonly_output_file(path,"moments")
+
+            # load fields data
+            parallel_pressure_zrt, parallel_heat_flux_zrt, thermal_speed_zrt =
+            load_electron_moments_data(fid)
+
+            close(fid)
+            
+            p = parallel_pressure_zrt[:,1,:]
+            q = parallel_heat_flux_zrt[:,1,:]
+            vt = thermal_speed_zrt[:,1,:]
+        end
+
+        # Regression test
+        actual_p = p[begin:3:end, end]
+        actual_q = q[begin:3:end, end]
+        actual_vt = vt[begin:3:end, end]
+        if expected_p == nothing
+            # Error: no expected input provided
+            println("data tested would be: ", actual_p)
+            @test false
+        else
+            @test isapprox(actual_p, expected_p, rtol=rtol, atol=atol)
+        end
+        if expected_q == nothing
+            # Error: no expected input provided
+            println("data tested would be: ", actual_q)
+            @test false
+        else
+            @test isapprox(actual_q, expected_q, rtol=rtol, atol=atol)
+        end
+        if expected_vt == nothing
+            # Error: no expected input provided
+            println("data tested would be: ", actual_vt)
+            @test false
+        else
+            @test isapprox(actual_vt, expected_vt, rtol=rtol, atol=atol)
+        end
+    end
+end
+
+function runtests()
+    # Create a temporary directory for test output
+    test_output_directory = get_MPI_tempdir()
+
+    expected_p = [0.44951398349003646, 0.4481071794001275, 0.4446923360697697,
+                  0.44101458684278777, 0.4384371320097695, 0.43732241798810445,
+                  0.4369770330480599, 0.4358502579878266, 0.434769684040917,
+                  0.43442263207284637, 0.43458925629798073, 0.4346831326574436,
+                  0.4350671485602161, 0.4362345316295716, 0.438235619583819,
+                  0.4402734411528, 0.4411447249385275, 0.4420727185406772,
+                  0.44481447378612293, 0.44890303238070045, 0.4530257687705389,
+                  0.4554124231936834, 0.4562726972649091, 0.4597319740656588,
+                  0.46533878727932276, 0.4713773067908409, 0.47559731298259184,
+                  0.4767158034722211, 0.47998294113890483, 0.4863021470825949,
+                  0.4937119852803594, 0.49963221191774776, 0.5018851701590787,
+                  0.5041515792976728, 0.5102381811386025, 0.5181350099985299,
+                  0.525127401588624, 0.52883991125208, 0.530126542118299,
+                  0.5350544857912003, 0.5423139848496283, 0.5492817298572252,
+                  0.5536930071759895, 0.5548035254843448, 0.5579127055178238,
+                  0.5633811620967428, 0.5689229399038384, 0.5726992489776701,
+                  0.5739864539763129, 0.5751981737538447, 0.5780380624660489,
+                  0.5808091879153607, 0.582371244425747, 0.5828442626520689,
+                  0.5829483019641976, 0.5830505324419235, 0.582289662647816,
+                  0.5804174302674363, 0.5785673270305791, 0.5780104294159015,
+                  0.5762363804591635, 0.5722372963746002, 0.5667182906200348,
+                  0.5617560490141181, 0.5597527656925925, 0.5576781862353368,
+                  0.5518325952127137, 0.5437117958155238, 0.5360738744744328,
+                  0.5318632735767702, 0.5303800815638172, 0.52458857332585,
+                  0.5157429166779243, 0.5068932260184091, 0.5010911813300937,
+                  0.4996034253232407, 0.49537353092296554, 0.4876627194305651,
+                  0.47936652714814293, 0.4732825296867743, 0.47108892532041974,
+                  0.46894816235702474, 0.4635185649133237, 0.4571465360817948,
+                  0.452118238600268, 0.4496786524823512]
+    expected_q = [0.6831704519267778, 0.6763596731108003, 0.6550248548395122,
+                  0.621137479330641, 0.5855826766086369, 0.564624348243424,
+                  0.5570221429154738, 0.5262612511440813, 0.4759763560524193,
+                  0.4215209814677685, 0.3833879518808367, 0.3732765291862611,
+                  0.34373639264749195, 0.2865988207461233, 0.21960124384883264,
+                  0.16604253330270563, 0.14564326706615302, 0.1251073951350269,
+                  0.06985203286011953, -0.002189422724091073, -0.06650025444931523,
+                  -0.10092291876877173, -0.11290749491826756, -0.15911660726647084,
+                  -0.22829407549560776, -0.296381042779731, -0.34065325931034834,
+                  -0.35197568475314056, -0.38412394718970916, -0.44264467058363494,
+                  -0.5056139329952417, -0.5518262734055381, -0.5684878570815051,
+                  -0.5847404886752257, -0.6258728785916116, -0.6737205371368551,
+                  -0.7106893697823674, -0.728139347093567, -0.733817023138346,
+                  -0.7537144673017869, -0.7772295129285097, -0.7923028306406594,
+                  -0.7973216863277829, -0.7979485411906018, -0.7981733058556373,
+                  -0.7921566532586084, -0.7749750759560777, -0.7538532886753724,
+                  -0.7441527309138952, -0.7334870232950073, -0.7004061657361658,
+                  -0.6479952272743823, -0.5926982319505321, -0.5599217222341915,
+                  -0.5480028121415452, -0.49964097922347867, -0.4203001696008753,
+                  -0.3344059204499378, -0.2745563922655548, -0.258755893042823,
+                  -0.21281014297392614, -0.1251037616154198, -0.02486081101616467,
+                  0.052711670524221245, 0.08156346618184351, 0.11018860125597137,
+                  0.184956394584215, 0.2769990401549583, 0.3533862259663847,
+                  0.39185061791993997, 0.4048279100244915, 0.4527903930933007,
+                  0.5181604219781556, 0.5745067111528425, 0.6066373424047252,
+                  0.6142607159164151, 0.6345436396519946, 0.6660807860559502,
+                  0.6916551627488451, 0.7043061657863463, 0.7074613359074287,
+                  0.7097605448059333, 0.7117980175789959, 0.7059983443925062,
+                  0.693373867220681, 0.6839039812030164]
+    expected_vt = [60.50323643979975, 60.46152282311776, 60.35195588066853,
+                   60.21533395043499, 60.09982168515988, 60.04066964102238,
+                   60.02053283714272, 59.94518544878992, 59.83955672767372,
+                   59.74465499035528, 59.688176306690224, 59.674434375999695,
+                   59.63705080297302, 59.575593022878486, 59.5200362867479,
+                   59.487302095351524, 59.477403317059846, 59.46882152619191,
+                   59.45242049432979, 59.44514092412334, 59.45163053025197,
+                   59.460024506000224, 59.463744883587445, 59.481935299569656,
+                   59.52060886409094, 59.572294471701106, 59.61340646988501,
+                   59.624904994314015, 59.6598130677176, 59.73238906214771,
+                   59.82488626148378, 59.9037999413929, 59.934901010623335,
+                   59.966754837133514, 60.05500712816102, 60.17514068478408,
+                   60.28669865440711, 60.34792638918577, 60.36947690125584,
+                   60.45364714009672, 60.58266887898263, 60.71293986868975,
+                   60.79930625245654, 60.82160071089043, 60.8853594969313,
+                   61.00319313188073, 61.132712888454186, 61.229832513831234,
+                   61.265349742432086, 61.30028735647212, 61.39021480819434,
+                   61.49843415833545, 61.586182165745534, 61.62960105016664,
+                   61.644121615044114, 61.69715263940207, 61.76740693027959,
+                   61.825107969209625, 61.8560293732873, 61.86304841080026,
+                   61.880900287439, 61.90495043959934, 61.91707450705104,
+                   61.91528849056989, 61.912066851775705, 61.907449920498564,
+                   61.88836178502154, 61.84936547495241, 61.801519377780714,
+                   61.77096806875504, 61.75953059854668, 61.711577658100396,
+                   61.62836639460284, 61.53297073619124, 61.46358406675648,
+                   61.44488197626709, 61.38959988569385, 61.280332330036074,
+                   61.14919211755058, 61.04273591939215, 61.00193140312018,
+                   60.96075236762326, 60.84965248313868, 60.70479794465818,
+                   60.57625670314503, 60.50800809804206]
+
+    @testset "Braginskii electron IMEX timestepping" verbose=use_verbose begin
+        println("Braginskii electron IMEX timestepping tests")
+
+        @testset "Split 3" begin
+            test_input["base_directory"] = test_output_directory
+            run_test(test_input, expected_p, expected_q, expected_vt)
+        end
+        @long @testset "Check other timestep - $type" for
+                type ∈ ("KennedyCarpenterARK437",)
+
+            timestep_check_input = deepcopy(test_input)
+            timestep_check_input["base_directory"] = test_output_directory
+            timestep_check_input["run_name"] = type
+            timestep_check_input["timestepping"]["type"] = type
+            run_test(timestep_check_input, expected_p, expected_q, expected_vt,
+                     rtol=2.e-4, atol=1.e-10)
+        end
+    end
+
+    if global_rank[] == 0
+        # Delete output directory to avoid using too much disk space
+        rm(realpath(test_output_directory); recursive=true)
+    end
+end
+
+end # BraginskiiElectronsIMEX
+
+
+using .BraginskiiElectronsIMEX
+
+BraginskiiElectronsIMEX.runtests()
diff --git a/moment_kinetics/test/runtests.jl b/moment_kinetics/test/runtests.jl
index fa0f5d64f..73d688f06 100644
--- a/moment_kinetics/test/runtests.jl
+++ b/moment_kinetics/test/runtests.jl
@@ -16,6 +16,7 @@ function runtests()
         include(joinpath(@__DIR__, "harrisonthompson.jl"))
         include(joinpath(@__DIR__, "wall_bc_tests.jl"))
         include(joinpath(@__DIR__, "recycling_fraction_tests.jl"))
+        include(joinpath(@__DIR__, "braginskii_electrons_imex_tests.jl"))
         include(joinpath(@__DIR__, "fokker_planck_tests.jl"))
         include(joinpath(@__DIR__, "fokker_planck_time_evolution_tests.jl"))
         include(joinpath(@__DIR__, "gyroaverage_tests.jl"))

From 7c8b2e08dc3bfc099ca0a7037d865931897d22a5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 09:23:09 +0100
Subject: [PATCH 294/428] Do not do kinetic electron solve for lower-order RK
 solution

When using adaptive timestepping, we use the RK scheme to calculate a
lower-order approximation to the solution, which is used to estimate the
error (which is then used to update the timestep). To apply boundary
conditions and constraints to the lower-order approximations of the ion
and neutral distribution functions, we use the
`apply_all_bcs_constraints_update_moments!()`
function. This function calls the (expensive!) `update_electron_pdf!()`
function that runs the kinetic electron update, but this kinetic
electron update is not needed for the lower-order approximation, so it
is better to skip it for that step.
---
 moment_kinetics/src/time_advance.jl | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index c58c278fb..d8b500027 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2134,7 +2134,8 @@ function apply_all_bcs_constraints_update_moments!(
         this_scratch, pdf, moments, fields, boundary_distributions, scratch_electron, vz,
         vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
         collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-        t_params, advance, scratch_dummy, diagnostic_moments; pdf_bc_constraints=true)
+        t_params, advance, scratch_dummy, diagnostic_moments; pdf_bc_constraints=true,
+        update_electrons=true)
 
     begin_s_r_z_region()
 
@@ -2215,7 +2216,7 @@ function apply_all_bcs_constraints_update_moments!(
         # (because this function is being called after a failed timestep, to reset to the
         # state at the beginning of the step), we also do not need to update the
         # electrons.
-        if pdf_bc_constraints
+        if update_electrons
             update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, fields.phi,
                                  r, z, vperp, vpa, z_spectral, vperp_spectral,
                                  vpa_spectral, z_advect, vpa_advect, scratch_dummy,
@@ -2423,7 +2424,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         scratch[2], pdf, moments, fields, boundary_distributions, scratch_electron, vz,
         vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
         collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-        t_params, advance, scratch_dummy, false)
+        t_params, advance, scratch_dummy, false; update_electrons=false)
 
     # Re-calculate moment derivatives in the `moments` struct, in case they were changed
     # by the previous call
@@ -2432,7 +2433,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
         advect_objects, composition, collisions, geometry, gyroavs,
         external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
-        false; pdf_bc_constraints=false)
+        false; pdf_bc_constraints=false, update_electrons=false)
 
     # Calculate the timstep error estimates
     ion_pdf_error = local_error_norm(scratch[2].pdf, scratch[t_params.n_rk_stages+1].pdf,
@@ -2587,7 +2588,8 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
             scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
             vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
             collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false)
+            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false,
+            update_electrons=false)
     end
 
     if composition.electron_physics == kinetic_electrons
@@ -2751,7 +2753,8 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
             scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
             advect_objects, composition, collisions, geometry, gyroavs,
             external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
-            diagnostic_moments; pdf_bc_constraints=apply_bc_constraints)
+            diagnostic_moments; pdf_bc_constraints=apply_bc_constraints,
+            update_electrons=apply_bc_constraints)
     end
 
     if t_params.adaptive

From 8359bb74c7b37dd2c02f37d445ef28052692bc8d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 09:39:51 +0100
Subject: [PATCH 295/428] If electron solve takes too many iterations, trigger
 timestep failure

Decreasing the ion/neutral timestep makes it easier for the kinetic
electron solve to converge (because the input quantities change less
from the previous converged solution). If the kinetic electron solve
takes too many iterations, we can trigger an ion/neutral timestep
failure which will decrease the ion/neutral timestep, which may allow
the kinetic electron solve to converge.
---
 .../src/electron_kinetic_equation.jl          |  9 ++---
 moment_kinetics/src/initial_conditions.jl     | 13 +++++-
 moment_kinetics/src/time_advance.jl           | 40 ++++++++++++++-----
 3 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 2190c49a1..055eb0df3 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1,6 +1,5 @@
 module electron_kinetic_equation
 
-using Dates
 using LinearAlgebra
 using MPI
 
@@ -549,11 +548,11 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
     end
     if !electron_pdf_converged
-        # need to exit or handle this appropriately
-        error("!!!max number of iterations for electron pdf update exceeded!!!\n"
-              * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+        success = false
+    else
+        success = true
     end
-    return time
+    return time, success
 end
 
 function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=false)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index fb1945302..3f732d3fa 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -11,6 +11,7 @@ export create_boundary_distributions
 export create_pdf
 
 # package
+using Dates
 using SpecialFunctions: erfc
 # modules
 using ..type_definitions: mk_float, mk_int
@@ -653,7 +654,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             resize!(t_params.dfns_output_times, n_truncated)
             t_params.dfns_output_times .= truncated_times
         end
-        electron_pseudotime =
+        electron_pseudotime, success =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
@@ -664,13 +665,17 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         initial_time=code_time,
                                         residual_tolerance=t_input["initialization_residual_value"],
                                         evolve_ppar=true)
+        if !success
+            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+        end
 
         # Now run without evolve_ppar=true to get pdf_electron fully to steady state,
         # ready for the start of the ion time advance.
         if global_rank[] == 0
             println("Initializing electrons - evolving pdf_electron only to steady state")
         end
-        electron_pseudotime =
+        electron_pseudotime, success =
             @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
                                         vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
@@ -679,6 +684,10 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         max_electron_pdf_iterations;
                                         io_electron=io_initial_electron,
                                         initial_time=electron_pseudotime)
+        if !success
+            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+        end
 
         begin_r_z_vperp_vpa_region()
         @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index d8b500027..464585bd5 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2144,6 +2144,8 @@ function apply_all_bcs_constraints_update_moments!(
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
+    success = true
+
     if pdf_bc_constraints
         # Ensure there are no negative values in the pdf before applying boundary
         # conditions, so that negative deviations do not mess up the integral-constraint
@@ -2216,13 +2218,17 @@ function apply_all_bcs_constraints_update_moments!(
         # (because this function is being called after a failed timestep, to reset to the
         # state at the beginning of the step), we also do not need to update the
         # electrons.
-        if update_electrons
-            update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, fields.phi,
-                                 r, z, vperp, vpa, z_spectral, vperp_spectral,
-                                 vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                 t_params.electron, collisions, composition,
-                                 external_source_settings, num_diss_params,
-                                 max_electron_pdf_iterations)
+        # Note that if some solve for the implicit timestep already failed, we will reset
+        # to the beginning of the ion/neutral timestep, so the electron solution
+        # calculated here would be discarded - we might as well skip calculating it in
+        # that case.
+        if update_electrons && success
+            _, kinetic_electron_success = update_electron_pdf!(
+               scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
+               z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect,
+               scratch_dummy, t_params.electron, collisions, composition,
+               external_source_settings, num_diss_params, max_electron_pdf_iterations)
+            success = success && kinetic_electron_success
         end
     end
     # update the electron parallel friction force
@@ -2292,6 +2298,8 @@ function apply_all_bcs_constraints_update_moments!(
                                               z_spectral,
                                               num_diss_params.neutral.moment_dissipation_coefficient)
     end
+
+    return success
 end
 
 """
@@ -2712,12 +2720,19 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # The result of the implicit solve gives the state vector at 'istage'
                 # which is used as input to the explicit part of the IMEX time step.
                 old_scratch = scratch_implicit[istage]
-                apply_all_bcs_constraints_update_moments!(
+                kinetic_electron_success = apply_all_bcs_constraints_update_moments!(
                     scratch_implicit[istage], pdf, moments, fields,
                     boundary_distributions, scratch_electron, vz, vr, vzeta, vpa, vperp,
                     z, r, spectral_objects, advect_objects, composition, collisions,
                     geometry, gyroavs, external_source_settings, num_diss_params,
                     t_params, advance, scratch_dummy, false)
+                success = success && kinetic_electron_success
+                if !success
+                    # Break out of the istage loop, as passing `success = false` to the
+                    # adaptive timestep update function will signal a failed timestep, so
+                    # that we restart this timestep with a smaller `dt`.
+                    break
+                end
             end
         else
             # Fully explicit method starts the forward-Euler step with the result from the
@@ -2748,13 +2763,20 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                                 || istage == n_rk_stages
                                 || t_params.implicit_coefficient_is_zero[istage+1])
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
-        apply_all_bcs_constraints_update_moments!(
+        kinetic_electron_success = apply_all_bcs_constraints_update_moments!(
             scratch[istage+1], pdf, moments, fields, boundary_distributions,
             scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
             advect_objects, composition, collisions, geometry, gyroavs,
             external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
             diagnostic_moments; pdf_bc_constraints=apply_bc_constraints,
             update_electrons=apply_bc_constraints)
+        success = success && kinetic_electron_success
+        if !success
+            # Break out of the istage loop, as passing `success = false` to the
+            # adaptive timestep update function will signal a failed timestep, so
+            # that we restart this timestep with a smaller `dt`.
+            break
+        end
     end
 
     if t_params.adaptive

From ebd7a4b0bb13d0972fabb3bad714f1dc70175fac Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 10:12:10 +0100
Subject: [PATCH 296/428] Fix local error norm for electron distribution
 function

---
 moment_kinetics/src/runge_kutta.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 803cc2715..07fc72649 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -914,7 +914,7 @@ function local_error_norm(f_loworder::MPISharedArray{mk_float,4},
         end
         # Will sum results from different processes in shared memory block after returning
         # from this function.
-        nvpa, nvperp, nz, nr, nspecies = size(f_loworder)
+        nvpa, nvperp, nz, nr = size(f_loworder)
         if skip_r_inner
             nr -= 1
         end

From a13e215a2f4c29af4f7dc102e787b41ad43176fc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 12:35:30 +0100
Subject: [PATCH 297/428] Comment out check for qpar_updated[] in
 calculate_electron_qpar!()

Do not seem to be resetting to qpar_update[]=false correctly. Not often
enough? Not reliable enough? Not in the right place?
---
 moment_kinetics/src/electron_fluid_equations.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 8efe49368..74816d1f2 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -419,7 +419,7 @@ function calculate_electron_qpar!(electron_moments, pdf, ppar_e, upar_e, upar_i,
                                   me_over_mi, electron_model, vpa)
     # only calculate qpar_e if needs updating
     qpar_updated = electron_moments.qpar_updated
-    if !qpar_updated[]
+    #if !qpar_updated[]
         qpar_e = electron_moments.qpar
         vth_e = electron_moments.vth
         dTe_dz = electron_moments.dT_dz
@@ -453,7 +453,7 @@ function calculate_electron_qpar!(electron_moments, pdf, ppar_e, upar_e, upar_i,
                 qpar_e[iz,ir] = 0.0
             end
         end
-    end
+    #end
     # qpar has been updated
     qpar_updated[] = true
     return nothing

From 67c459eeac234b51c072760566cf5a5189bb6e63 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 12:37:30 +0100
Subject: [PATCH 298/428] Update moments.electron.dens and
 moments.electron.upar

These need to be copied out of `scratch`.
---
 moment_kinetics/src/time_advance.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 464585bd5..cf868a0f2 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2211,6 +2211,8 @@ function apply_all_bcs_constraints_update_moments!(
             moments.neutral.pz[iz,ir,isn] = this_scratch.pz_neutral[iz,ir,isn]
         end
         @loop_r_z ir iz begin
+            moments.electron.dens[iz,ir] = this_scratch.electron_density[iz,ir]
+            moments.electron.upar[iz,ir] = this_scratch.electron_upar[iz,ir]
             moments.electron.ppar[iz,ir] = this_scratch.electron_ppar[iz,ir]
         end
 

From 3cfc8d5fb6c48cd2b9ec575fe07dcf31f0eb435a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 12:38:24 +0100
Subject: [PATCH 299/428] Reset electron pdf before resetting moments when
 timestep fails

When using kinetic electrons, the electron parallel heat flux needs to
be calculated from the electron distribution function, so when resetting
variables because a timestep failed, we need to reset the electron
distribution function before resetting the moments.
---
 moment_kinetics/src/time_advance.jl | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index cf868a0f2..a663c91c8 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2590,18 +2590,6 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
                                        total_points, current_dt, error_norm_method,
                                        success, nl_max_its_fraction)
 
-    if t_params.previous_dt[] == 0.0
-        # Re-update remaining velocity moments that are calculable from the evolved
-        # pdf These need to be re-calculated because `scratch[istage+1]` is now the
-        # state at the beginning of the timestep, because the timestep failed
-        apply_all_bcs_constraints_update_moments!(
-            scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
-            vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
-            collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false,
-            update_electrons=false)
-    end
-
     if composition.electron_physics == kinetic_electrons
         if t_params.previous_dt[] == 0.0
             # Reset electron pdf to its value at the beginning of this step.
@@ -2623,6 +2611,18 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         end
     end
 
+    if t_params.previous_dt[] == 0.0
+        # Re-update remaining velocity moments that are calculable from the evolved
+        # pdf These need to be re-calculated because `scratch[istage+1]` is now the
+        # state at the beginning of the timestep, because the timestep failed
+        apply_all_bcs_constraints_update_moments!(
+            scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
+            vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+            collisions, geometry, gyroavs, external_source_settings, num_diss_params,
+            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false,
+            update_electrons=false)
+    end
+
     return nothing
 end
 

From df01251bb683c2faf75e381c18a5a925a416bcd0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 13:52:12 +0100
Subject: [PATCH 300/428] Decrease max_electron_pdf_iterations to 1000

Now that the ion timestep can be decreased if the kinetic electron
iteration takes too long, it makes sense to decrease the maximum
iteration count (for the solves during the ion time-loop) so that we
don't waste time when it would be more efficient to slightly decrease
the ion timestep.
---
 moment_kinetics/src/time_advance.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index a663c91c8..4726d50a5 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2195,7 +2195,7 @@ function apply_all_bcs_constraints_update_moments!(
                                            num_diss_params.electron.moment_dissipation_coefficient, 
                                            composition.electron_physics)
     if composition.electron_physics == kinetic_electrons
-        max_electron_pdf_iterations = 100000
+        max_electron_pdf_iterations = 1000
 
         # Copy ion and electron moments from `scratch` into `moments` to be used in
         # electron kinetic equation update

From 97fdfd4a9237e2f8279e07c93d92465ed5a1eb97 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 14:58:43 +0100
Subject: [PATCH 301/428] Use max_increase_factor_near_last_fail in kinetic
 electron example input

Now that the timestep fails (decreasing the step size) if the kinetic
electron solve hits its maximum iterations limit, it helps to use
max_increase_factor_near_last_fail to keep the ion timestep near to the
last successful value for a reasonably long time, to avoid too-frequent
timestep failures.
---
 .../wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml   | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
index 965f192b9..7828c61fa 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
@@ -69,8 +69,10 @@ type = "Fekete4(3)"
 nstep = 1000000
 dt = 1.0e-6
 minimum_dt = 1.0e-6
+max_increase_factor_near_last_fail = 1.001
+last_fail_proximity_factor = 1.1
 nwrite = 10000
-nwrite_dfns = 100000
+nwrite_dfns = 10000
 steady_state_residual = true
 converged_residual_value = 1.0e-3
 

From 864eb2e9c3221f7bb86a828b314ec7f905d8797d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 7 Jun 2024 22:48:27 +0100
Subject: [PATCH 302/428] Fix post-processing for failure caused by kinetic
 electrons diagnostic

---
 .../src/makie_post_processing.jl              |  7 +++
 .../src/electron_kinetic_equation.jl          |  6 +-
 moment_kinetics/src/initial_conditions.jl     |  4 +-
 moment_kinetics/src/runge_kutta.jl            | 18 +++++-
 moment_kinetics/src/time_advance.jl           | 57 ++++++++++---------
 5 files changed, 57 insertions(+), 35 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 0dd67a2af..53ed2585d 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -7413,6 +7413,13 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
                                 linestyle=:dot,
                                 label=prefix * "nonlinear iteration convergence failure", ax=ax_failures)
                     end
+                    if ri.composition.electron_physics == kinetic_electrons
+                        # Kinetic electron iteration failed to converge
+                        counter += 1
+                        plot_1d(time, @view failure_caused_by_per_output[counter,:];
+                                linestyle=:dot,
+                                label=prefix * "nonlinear iteration convergence failure", ax=ax_failures)
+                    end
                 end
 
                 if counter > size(failure_caused_by_per_output, 1)
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 055eb0df3..c09c969c1 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -548,9 +548,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
     end
     if !electron_pdf_converged
-        success = false
+        success = "kinetic-electrons"
     else
-        success = true
+        success = ""
     end
     return time, success
 end
@@ -1150,7 +1150,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
     end
 
     adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
-                                       total_points, current_dt, error_norm_method, true,
+                                       total_points, current_dt, error_norm_method, "",
                                        0.0, electron=true)
 
     return nothing
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 3f732d3fa..e7b1a7d2e 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -665,7 +665,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         initial_time=code_time,
                                         residual_tolerance=t_input["initialization_residual_value"],
                                         evolve_ppar=true)
-        if !success
+        if success != ""
             error("!!!max number of iterations for electron pdf update exceeded!!!\n"
                   * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
         end
@@ -684,7 +684,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
                                         max_electron_pdf_iterations;
                                         io_electron=io_initial_electron,
                                         initial_time=electron_pseudotime)
-        if !success
+        if success != ""
             error("!!!max number of iterations for electron pdf update exceeded!!!\n"
                   * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
         end
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 07fc72649..f6912d379 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1056,7 +1056,7 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
         error("Unrecognized error_norm_method '$method'")
     end
 
-    if !success
+    if success != ""
         # Iteration failed in implicit part of timestep try decreasing timestep
 
         # Set scratch[end] equal to scratch[1] to start the timestep over
@@ -1088,8 +1088,20 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
             t_params.previous_dt[] = 0.0
 
             # Call the 'cause' of the timestep failure the variable that has the biggest
-            # error norm here
-            t_params.failure_caused_by[end] += 1
+            # error norm here.
+            # Could do with a better way to sort the different possible types of
+            # convergence failure...
+            if t_params.rk_coefs_implicit !== nothing && composition.electron_physics == kinetic_electrons
+                if success == "nonlinear-solver"
+                    t_params.failure_caused_by[end-1] += 1
+                elseif success == "kinetic-electrons"
+                    t_params.failure_caused_by[end] += 1
+                else
+                    error("Unrecognised cause of convergence failure: \"$success\"")
+                end
+            else
+                t_params.failure_caused_by[end] += 1
+            end
 
             # If we were trying to take a step to the output timestep, dt will be smaller on
             # the re-try, so will not reach the output time.
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 4726d50a5..375b20cd0 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -549,7 +549,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
         # electron pressure
         push!(t_params.limit_caused_by, 0) # RK accuracy
-        push!(t_params.failure_caused_by, 0)
+        push!(t_params.failure_caused_by, 0) # RK accuracy for electron_ppar
+        if composition.electron_physics == kinetic_electrons
+            push!(t_params.failure_caused_by, 0) # Convergence failure for kinetic electron solve
+        end
     end
     if composition.n_neutral_species > 0
         # neutral pdf
@@ -2144,7 +2147,7 @@ function apply_all_bcs_constraints_update_moments!(
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
-    success = true
+    success = ""
 
     if pdf_bc_constraints
         # Ensure there are no negative values in the pdf before applying boundary
@@ -2224,13 +2227,13 @@ function apply_all_bcs_constraints_update_moments!(
         # to the beginning of the ion/neutral timestep, so the electron solution
         # calculated here would be discarded - we might as well skip calculating it in
         # that case.
-        if update_electrons && success
+        if update_electrons && success == ""
             _, kinetic_electron_success = update_electron_pdf!(
                scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
                z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect,
                scratch_dummy, t_params.electron, collisions, composition,
                external_source_settings, num_diss_params, max_electron_pdf_iterations)
-            success = success && kinetic_electron_success
+            success = kinetic_electron_success
         end
     end
     # update the electron parallel friction force
@@ -2677,7 +2680,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
     end
 
     # success is set to false if an iteration failed to converge in an implicit solve
-    success = true
+    success = ""
     for istage ∈ 1:n_rk_stages
         if t_params.rk_coefs_implicit !== nothing
             update_solution_vector!(scratch_implicit[istage], scratch[istage], moments,
@@ -2702,19 +2705,21 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # Note the timestep for this solve is rk_coefs_implict[istage,istage]*dt.
                 # The diagonal elements are equal to the Butcher 'a' coefficients
                 # rk_coefs_implicit[istage,istage]=a[istage,istage].
-                success = backward_euler!(scratch_implicit[istage], scratch[istage], pdf,
-                                          fields, moments, advect_objects, vz, vr, vzeta,
-                                          vpa, vperp, gyrophase, z, r, t, t_params.dt[] *
-                                          t_params.rk_coefs_implicit[istage,istage],
-                                          spectral_objects, composition, collisions,
-                                          geometry, scratch_dummy,
-                                          manufactured_source_list,
-                                          external_source_settings, num_diss_params,
-                                          gyroavs, nl_solver_params, advance_implicit,
-                                          fp_arrays, istage)
-                success = MPI.Allreduce(success, &, comm_world)
-                if !success
-                    # Break out of the istage loop, as passing `success = false` to the
+                nl_success = backward_euler!(scratch_implicit[istage], scratch[istage],
+                                             pdf, fields, moments, advect_objects, vz, vr,
+                                             vzeta, vpa, vperp, gyrophase, z, r, t,
+                                             t_params.dt[] *
+                                             t_params.rk_coefs_implicit[istage,istage],
+                                             spectral_objects, composition, collisions,
+                                             geometry, scratch_dummy,
+                                             manufactured_source_list,
+                                             external_source_settings, num_diss_params,
+                                             gyroavs, nl_solver_params, advance_implicit,
+                                             fp_arrays, istage)
+                nl_success = MPI.Allreduce(nl_success, &, comm_world)
+                if !nl_success
+                    success = "nonlinear-solver"
+                    # Break out of the istage loop, as passing `success != ""` to the
                     # adaptive timestep update function will signal a failed timestep, so
                     # that we restart this timestep with a smaller `dt`.
                     break
@@ -2722,15 +2727,14 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # The result of the implicit solve gives the state vector at 'istage'
                 # which is used as input to the explicit part of the IMEX time step.
                 old_scratch = scratch_implicit[istage]
-                kinetic_electron_success = apply_all_bcs_constraints_update_moments!(
+                success = apply_all_bcs_constraints_update_moments!(
                     scratch_implicit[istage], pdf, moments, fields,
                     boundary_distributions, scratch_electron, vz, vr, vzeta, vpa, vperp,
                     z, r, spectral_objects, advect_objects, composition, collisions,
                     geometry, gyroavs, external_source_settings, num_diss_params,
                     t_params, advance, scratch_dummy, false)
-                success = success && kinetic_electron_success
-                if !success
-                    # Break out of the istage loop, as passing `success = false` to the
+                if success != ""
+                    # Break out of the istage loop, as passing `success != ""` to the
                     # adaptive timestep update function will signal a failed timestep, so
                     # that we restart this timestep with a smaller `dt`.
                     break
@@ -2765,16 +2769,15 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                                 || istage == n_rk_stages
                                 || t_params.implicit_coefficient_is_zero[istage+1])
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
-        kinetic_electron_success = apply_all_bcs_constraints_update_moments!(
+        success = apply_all_bcs_constraints_update_moments!(
             scratch[istage+1], pdf, moments, fields, boundary_distributions,
             scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
             advect_objects, composition, collisions, geometry, gyroavs,
             external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
             diagnostic_moments; pdf_bc_constraints=apply_bc_constraints,
             update_electrons=apply_bc_constraints)
-        success = success && kinetic_electron_success
-        if !success
-            # Break out of the istage loop, as passing `success = false` to the
+        if success != ""
+            # Break out of the istage loop, as passing `success != ""` to the
             # adaptive timestep update function will signal a failed timestep, so
             # that we restart this timestep with a smaller `dt`.
             break
@@ -2797,7 +2800,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                                   advect_objects, gyroavs, num_diss_params,
                                   advance, scratch_dummy, r, z, vperp, vpa,
                                   vzeta, vr, vz, success, nl_max_its_fraction)
-    elseif !success
+    elseif success != ""
         error("Implicit part of timestep failed")
     end
 

From b22629988ea5154845a58c500e16448830b987c0 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 13:10:16 +0100
Subject: [PATCH 303/428] Save whole `io_input` into I/O structs and
 info-Tuples

Makes it easier to implement more I/O options.
---
 moment_kinetics/ext/file_io_netcdf.jl |   6 +-
 moment_kinetics/src/file_io.jl        | 135 +++++++++++++-------------
 moment_kinetics/src/file_io_hdf5.jl   |   6 +-
 3 files changed, 75 insertions(+), 72 deletions(-)

diff --git a/moment_kinetics/ext/file_io_netcdf.jl b/moment_kinetics/ext/file_io_netcdf.jl
index d1e52c737..9ca463674 100644
--- a/moment_kinetics/ext/file_io_netcdf.jl
+++ b/moment_kinetics/ext/file_io_netcdf.jl
@@ -21,9 +21,9 @@ function io_has_parallel(::Val{netcdf})
     return false
 end
 
-function open_output_file_implementation(::Val{netcdf}, prefix, parallel_io, io_comm,
+function open_output_file_implementation(::Val{netcdf}, prefix, io_input, io_comm,
                                          mode="c")
-    parallel_io && error("NetCDF interface does not support parallel I/O")
+    io_input.parallel_io && error("NetCDF interface does not support parallel I/O")
 
     # the netcdf file will be given by output_dir/run_name with .cdf appended
     filename = string(prefix, ".cdf")
@@ -32,7 +32,7 @@ function open_output_file_implementation(::Val{netcdf}, prefix, parallel_io, io_
     # create the new NetCDF file
     fid = NCDataset(filename, mode)
 
-    return fid, (filename, parallel_io, io_comm)
+    return fid, (filename, io_input, io_comm)
 end
 
 function create_io_group(parent::NCDataset, name; description=nothing)
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index d9d070f38..84f68ede5 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -57,7 +57,8 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
                        Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
                        Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Texte1, Texte2,
                        Texte3, Texte4, Tconstri, Tconstrn, Tconstre, Tint, Tfailcause,
-                       Telectrontime, Telectronint, Telectronfailcause, Tnldiagnostics}
+                       Telectrontime, Telectronint, Telectronfailcause, Tnldiagnostics,
+                       Tinput}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the time variable
@@ -164,15 +165,15 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     # write diagnostics generically for as many nonlinear solvers as are created.
     nl_solver_diagnostics::Tnldiagnostics
 
-    # Use parallel I/O?
-    parallel_io::Bool
+    # Settings for I/O
+    io_input::Tinput
  end
 
 """
 structure containing the data/metadata needed for binary file i/o
 distribution function data only
 """
-struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments}
+struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments, Tinput}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the ion species distribution function variable
@@ -182,8 +183,8 @@ struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments}
     # handle for the neutral species distribution function variable
     f_neutral::Tfn
 
-    # Use parallel I/O?
-    parallel_io::Bool
+    # Settings for I/O
+    io_input::Tinput
 
     # Handles for moment variables
     io_moments::Tmoments
@@ -194,7 +195,8 @@ structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
 struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4,
-                                Tconstr, Telectrontime, Telectronint, Telectronfailcause}
+                                Tconstr, Telectrontime, Telectronint, Telectronfailcause,
+                                Tinput}
     # file identifier for the binary file to which data is written
     fid::Tfile
     # handle for the pseudotime variable
@@ -234,8 +236,8 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3,
     # by adaptve timestepping algorithm
     electron_dt_before_last_fail::Telectrontime
 
-    # Use parallel I/O?
-    parallel_io::Bool
+    # Settings for I/O
+    io_input::Tinput
 end
 
 """
@@ -305,20 +307,18 @@ function setup_file_io(io_input, boundary_distributions, vz, vr, vzeta, vpa, vpe
 
         run_id = io_input.run_id
 
-        io_moments = setup_moments_io(out_prefix, io_input.binary_format, vz, vr, vzeta,
-                                      vpa, vperp, r, z, composition, collisions,
-                                      evolve_density, evolve_upar, evolve_ppar,
-                                      external_source_settings, input_dict,
-                                      io_input.parallel_io, comm_inter_block[], run_id,
+        io_moments = setup_moments_io(out_prefix, io_input, vz, vr, vzeta, vpa, vperp, r,
+                                      z, composition, collisions, evolve_density,
+                                      evolve_upar, evolve_ppar, external_source_settings,
+                                      input_dict, comm_inter_block[], run_id,
                                       restart_time_index, previous_runs_info,
                                       time_for_setup, t_params, nl_solver_params)
-        io_dfns = setup_dfns_io(out_prefix, io_input.binary_format,
-                                boundary_distributions, r, z, vperp, vpa, vzeta, vr, vz,
-                                composition, collisions, evolve_density, evolve_upar,
-                                evolve_ppar, external_source_settings, input_dict,
-                                io_input.parallel_io, comm_inter_block[], run_id,
-                                restart_time_index, previous_runs_info, time_for_setup,
-                                t_params, nl_solver_params)
+        io_dfns = setup_dfns_io(out_prefix, io_input, boundary_distributions, r, z, vperp,
+                                vpa, vzeta, vr, vz, composition, collisions,
+                                evolve_density, evolve_upar, evolve_ppar,
+                                external_source_settings, input_dict, comm_inter_block[],
+                                run_id, restart_time_index, previous_runs_info,
+                                time_for_setup, t_params, nl_solver_params)
 
         return ascii, io_moments, io_dfns
     end
@@ -350,8 +350,7 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
         if !parallel_io
             electrons_prefix *= ".$(iblock_index[])"
         end
-        fid, file_info = open_output_file(electrons_prefix, io_input.binary_format,
-                                          parallel_io, io_comm)
+        fid, file_info = open_output_file(electrons_prefix, io_input, io_comm)
 
         # write a header to the output file
         add_attribute!(fid, "file_info",
@@ -417,7 +416,7 @@ function get_electron_io_info(io_input, prefix_label)
         error("Unrecognized binary_format=$(io_input.binary_format)")
     end
 
-    return (filename, io_input.parallel_io, comm_inter_block[])
+    return (filename, io_input, comm_inter_block[])
 end
 
 """
@@ -425,8 +424,8 @@ Reopen an existing initial electron output file to append more data
 """
 function reopen_initial_electron_io(file_info)
     @serial_region begin
-        filename, parallel_io, io_comm = file_info
-        fid = reopen_output_file(filename, parallel_io, io_comm)
+        filename, io_input, io_comm = file_info
+        fid = reopen_output_file(filename, io_input, io_comm)
         dyn = get_group(fid, "dynamic_data")
 
         variable_list = get_variable_keys(dyn)
@@ -456,7 +455,7 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_failure_caused_by"),
                                         getvar("electron_limit_caused_by"),
                                         getvar("electron_dt_before_last_fail"),
-                                        parallel_io)
+                                        io_input)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -870,11 +869,12 @@ function create_dynamic_variable! end
 define dynamic (time-evolving) moment variables for writing to the hdf5 file
 """
 function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
-                                          r::coordinate, z::coordinate, parallel_io,
+                                          r::coordinate, z::coordinate, io_input,
                                           external_source_settings, evolve_density,
                                           evolve_upar, evolve_ppar, electron_physics,
                                           t_params, nl_solver_params)
     @serial_region begin
+        parallel_io = io_input.parallel_io
         dynamic = create_io_group(fid, "dynamic_data", description="time evolving variables")
 
         io_time = create_dynamic_variable!(dynamic, "time", mk_float; parallel_io=parallel_io,
@@ -1006,7 +1006,7 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                io_electron_failure_counter, io_electron_failure_caused_by,
                                io_electron_limit_caused_by,
                                io_electron_dt_before_last_fail, io_nl_solver_diagnostics,
-                               parallel_io)
+                               io_input)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1467,14 +1467,15 @@ define dynamic (time-evolving) distribution function variables for writing to th
 file
 """
 function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, composition,
-                                       parallel_io, external_source_settings,
+                                       io_input, external_source_settings,
                                        evolve_density, evolve_upar, evolve_ppar, t_params,
                                        nl_solver_params)
 
     @serial_region begin
+        parallel_io = io_input.parallel_io
         io_moments = define_dynamic_moment_variables!(fid, composition.n_ion_species,
                                                       composition.n_neutral_species, r, z,
-                                                      parallel_io,
+                                                      io_input,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
                                                       evolve_ppar,
@@ -1528,25 +1529,25 @@ function open_output_file_implementation end
 """
 Open an output file, selecting the backend based on io_option
 """
-function open_output_file(prefix, binary_format, parallel_io, io_comm)
-    check_io_implementation(binary_format)
+function open_output_file(prefix, io_input, io_comm)
+    check_io_implementation(io_input.binary_format)
 
-    return open_output_file_implementation(Val(binary_format), prefix, parallel_io,
+    return open_output_file_implementation(Val(io_input.binary_format), prefix, io_input,
                                            io_comm)
 end
 
 """
 Re-open an existing output file, selecting the backend based on io_option
 """
-function reopen_output_file(filename, parallel_io, io_comm)
+function reopen_output_file(filename, io_input, io_comm)
     prefix, format_string = splitext(filename)
     if format_string == ".h5"
         check_io_implementation(hdf5)
-        return open_output_file_implementation(Val(hdf5), prefix, parallel_io, io_comm,
+        return open_output_file_implementation(Val(hdf5), prefix, io_input, io_comm,
                                                "r+")[1]
     elseif format_string == ".cdf"
         check_io_implementation(netcdf)
-        return open_output_file_implementation(Val(netcdf), prefix, parallel_io, io_comm,
+        return open_output_file_implementation(Val(netcdf), prefix, io_input, io_comm,
                                                "a")[1]
     else
         error("Unsupported I/O format $binary_format")
@@ -1556,17 +1557,18 @@ end
 """
 setup file i/o for moment variables
 """
-function setup_moments_io(prefix, binary_format, vz, vr, vzeta, vpa, vperp, r, z,
+function setup_moments_io(prefix, io_input, vz, vr, vzeta, vpa, vperp, r, z,
                           composition, collisions, evolve_density, evolve_upar,
-                          evolve_ppar, external_source_settings, input_dict, parallel_io,
+                          evolve_ppar, external_source_settings, input_dict,
                           io_comm, run_id, restart_time_index, previous_runs_info,
                           time_for_setup, t_params, nl_solver_params)
     @serial_region begin
         moments_prefix = string(prefix, ".moments")
+        parallel_io = io_input.parallel_io
         if !parallel_io
             moments_prefix *= ".$(iblock_index[])"
         end
-        fid, file_info = open_output_file(moments_prefix, binary_format, parallel_io, io_comm)
+        fid, file_info = open_output_file(moments_prefix, io_input, io_comm)
 
         # write a header to the output file
         add_attribute!(fid, "file_info", "Output moments data from the moment_kinetics code")
@@ -1589,7 +1591,7 @@ function setup_moments_io(prefix, binary_format, vz, vr, vzeta, vpa, vperp, r, z
         ### in a struct for later access ###
         io_moments = define_dynamic_moment_variables!(
             fid, composition.n_ion_species, composition.n_neutral_species, r, z,
-            parallel_io, external_source_settings, evolve_density, evolve_upar,
+            io_input, external_source_settings, evolve_density, evolve_upar,
             evolve_ppar, composition.electron_physics, t_params, nl_solver_params)
 
         close(fid)
@@ -1606,8 +1608,8 @@ Reopen an existing moments output file to append more data
 """
 function reopen_moments_io(file_info)
     @serial_region begin
-        filename, parallel_io, io_comm = file_info
-        fid = reopen_output_file(filename, parallel_io, io_comm)
+        filename, io_input, io_comm = file_info
+        fid = reopen_output_file(filename, io_input, io_comm)
         dyn = get_group(fid, "dynamic_data")
 
         variable_list = get_variable_keys(dyn)
@@ -1673,7 +1675,7 @@ function reopen_moments_io(file_info)
                                getvar("electron_failure_caused_by"),
                                getvar("electron_limit_caused_by"),
                                getvar("electron_dt_before_last_fail"),
-                               getvar("nl_solver_diagnostics"), parallel_io)
+                               getvar("nl_solver_diagnostics"), io_input)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1683,18 +1685,19 @@ end
 """
 setup file i/o for distribution function variables
 """
-function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vperp, vpa,
-                       vzeta, vr, vz, composition, collisions, evolve_density,
-                       evolve_upar, evolve_ppar, external_source_settings, input_dict,
-                       parallel_io, io_comm, run_id, restart_time_index,
-                       previous_runs_info, time_for_setup, t_params, nl_solver_params)
+function setup_dfns_io(prefix, io_input, boundary_distributions, r, z, vperp, vpa, vzeta,
+                       vr, vz, composition, collisions, evolve_density, evolve_upar,
+                       evolve_ppar, external_source_settings, input_dict, io_comm, run_id,
+                       restart_time_index, previous_runs_info, time_for_setup, t_params,
+                       nl_solver_params)
 
     @serial_region begin
         dfns_prefix = string(prefix, ".dfns")
+        parallel_io = io_input.parallel_io
         if !parallel_io
             dfns_prefix *= ".$(iblock_index[])"
         end
-        fid, file_info = open_output_file(dfns_prefix, binary_format, parallel_io, io_comm)
+        fid, file_info = open_output_file(dfns_prefix, io_input, io_comm)
 
         # write a header to the output file
         add_attribute!(fid, "file_info",
@@ -1722,7 +1725,7 @@ function setup_dfns_io(prefix, binary_format, boundary_distributions, r, z, vper
         ### create variables for time-dependent quantities and store them ###
         ### in a struct for later access ###
         io_dfns = define_dynamic_dfn_variables!(
-            fid, r, z, vperp, vpa, vzeta, vr, vz, composition, parallel_io,
+            fid, r, z, vperp, vpa, vzeta, vr, vz, composition, io_input,
             external_source_settings, evolve_density, evolve_upar, evolve_ppar, t_params,
             nl_solver_params)
 
@@ -1740,8 +1743,8 @@ Reopen an existing distribution-functions output file to append more data
 """
 function reopen_dfns_io(file_info)
     @serial_region begin
-        filename, parallel_io, io_comm = file_info
-        fid = reopen_output_file(filename, parallel_io, io_comm)
+        filename, io_input, io_comm = file_info
+        fid = reopen_output_file(filename, io_input, io_comm)
         dyn = get_group(fid, "dynamic_data")
 
         variable_list = get_variable_keys(dyn)
@@ -1810,10 +1813,10 @@ function reopen_dfns_io(file_info)
                                      getvar("electron_failure_caused_by"),
                                      getvar("electron_limit_caused_by"),
                                      getvar("electron_dt_before_last_fail"),
-                                     getvar("nl_solver_diagnostics"), parallel_io)
+                                     getvar("nl_solver_diagnostics"), io_input)
 
         return io_dfns_info(fid, getvar("f"), getvar("f_electron"), getvar("f_neutral"),
-                            parallel_io, io_moments)
+                            io_input, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1861,7 +1864,7 @@ function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
             closefile = true
         end
 
-        parallel_io = io_moments.parallel_io
+        parallel_io = io_moments.io_input.parallel_io
 
         # add the time for this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.time, t, t_idx, parallel_io)
@@ -1915,7 +1918,7 @@ function write_em_fields_data_to_binary(fields, io_moments::io_moments_info, t_i
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_moments.parallel_io
+        parallel_io = io_moments.io_input.parallel_io
 
         # add the electrostatic potential and electric field components at this time slice to the hdf5 file
         append_to_dynamic_var(io_moments.phi, fields.phi, t_idx, parallel_io, z, r)
@@ -1936,7 +1939,7 @@ function write_ion_moments_data_to_binary(moments, n_ion_species,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_moments.parallel_io
+        parallel_io = io_moments.io_input.parallel_io
 
         # add the density data at this time slice to the output file
         append_to_dynamic_var(io_moments.density, moments.ion.dens, t_idx,
@@ -2029,7 +2032,7 @@ function write_electron_moments_data_to_binary(moments, t_params,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_moments.parallel_io
+        parallel_io = io_moments.io_input.parallel_io
 
         append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx,
                               parallel_io, z, r)
@@ -2098,7 +2101,7 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_moments.parallel_io
+        parallel_io = io_moments.io_input.parallel_io
 
         append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
                               parallel_io, z, r, n_neutral_species)
@@ -2207,7 +2210,7 @@ function write_ion_dfns_data_to_binary(ff, n_ion_species, io_dfns::io_dfns_info,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_dfns.parallel_io
+        parallel_io = io_dfns.io_input.parallel_io
 
         append_to_dynamic_var(io_dfns.f, ff, t_idx, parallel_io, vpa, vperp, z, r,
                               n_ion_species)
@@ -2226,7 +2229,7 @@ function write_electron_dfns_data_to_binary(ff_electron,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_dfns.parallel_io
+        parallel_io = io_dfns.io_input.parallel_io
 
         if io_dfns.f_electron !== nothing
             append_to_dynamic_var(io_dfns.f_electron, ff_electron, t_idx, parallel_io,
@@ -2247,7 +2250,7 @@ function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
     @serial_region begin
         # Only read/write from first process in each 'block'
 
-        parallel_io = io_dfns.parallel_io
+        parallel_io = io_dfns.io_input.parallel_io
 
         if n_neutral_species > 0
             append_to_dynamic_var(io_dfns.f_neutral, ff_neutral, t_idx, parallel_io, vz,
@@ -2277,7 +2280,7 @@ function write_electron_state(pdf, moments, t_params, t, io_or_file_info_initial
             closefile = true
         end
 
-        parallel_io = io_initial_electron.parallel_io
+        parallel_io = io_initial_electron.io_input.parallel_io
 
         # add the pseudo-time for this time slice to the hdf5 file
         append_to_dynamic_var(io_initial_electron.pseudotime, t, t_idx, parallel_io)
@@ -2577,7 +2580,7 @@ function debug_dump(vz::coordinate, vr::coordinate, vzeta::coordinate, vpa::coor
             ### in a struct for later access ###
             io_moments = define_dynamic_moment_variables!(fid, composition.n_ion_species,
                                                           composition.n_neutral_species,
-                                                          r, z, false,
+                                                          r, z, nothing,
                                                           external_source_settings,
                                                           evolve_density, evolve_upar,
                                                           evolve_ppar,
@@ -2585,7 +2588,7 @@ function debug_dump(vz::coordinate, vr::coordinate, vzeta::coordinate, vpa::coor
                                                           nl_solver_params)
             io_dfns = define_dynamic_dfn_variables!(
                 fid, r, z, vperp, vpa, vzeta, vr, vz, composition.n_ion_species,
-                composition.n_neutral_species, false, external_source_settings,
+                composition.n_neutral_species, nothing, external_source_settings,
                 evolve_density, evolve_upar, evolve_ppar, t_params, nl_solver_params)
 
             # create the "istage" variable, used to identify the rk stage where
diff --git a/moment_kinetics/src/file_io_hdf5.jl b/moment_kinetics/src/file_io_hdf5.jl
index c40dea9d5..e39066f4c 100644
--- a/moment_kinetics/src/file_io_hdf5.jl
+++ b/moment_kinetics/src/file_io_hdf5.jl
@@ -7,7 +7,7 @@ function io_has_parallel(::Val{hdf5})
     return HDF5.has_parallel()
 end
 
-function open_output_file_implementation(::Val{hdf5}, prefix, parallel_io, io_comm, mode="cw")
+function open_output_file_implementation(::Val{hdf5}, prefix, io_input, io_comm, mode="cw")
     # the hdf5 file will be given by output_dir/run_name with .h5 appended
     filename = string(prefix, ".h5")
 
@@ -16,7 +16,7 @@ function open_output_file_implementation(::Val{hdf5}, prefix, parallel_io, io_co
               * "characters), which will cause an error in HDF5.")
     end
     # create the new HDF5 file
-    if parallel_io
+    if io_input.parallel_io
         # if a file with the requested name already exists, remove it
         if mode == "cw" && MPI.Comm_rank(io_comm) == 0 && isfile(filename)
             rm(filename)
@@ -32,7 +32,7 @@ function open_output_file_implementation(::Val{hdf5}, prefix, parallel_io, io_co
         fid = h5open(filename, mode)
     end
 
-    return fid, (filename, parallel_io, io_comm)
+    return fid, (filename, io_input, io_comm)
 end
 
 # HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group

From 251c9e696b698eb1b6053bc9c4dc66880dc648d6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 14:37:10 +0100
Subject: [PATCH 304/428] Write output from scratch_pdf, instead of pdf_struct
 and moments_struct

This gives more flexibility to, in future, save for example error
estimates, etc. using the information in the `scratch_pdf` struct.
---
 .../src/electron_kinetic_equation.jl          |  17 +-
 moment_kinetics/src/file_io.jl                | 190 +++++++++++-------
 moment_kinetics/src/initial_conditions.jl     |  25 +--
 moment_kinetics/src/load_data.jl              |  11 -
 moment_kinetics/src/moment_kinetics.jl        |   4 +-
 moment_kinetics/src/time_advance.jl           |  28 ++-
 6 files changed, 159 insertions(+), 116 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index c09c969c1..c145a902c 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -246,9 +246,12 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     t_params.moments_output_counter[] += 1
     @serial_region begin
         if io_electron !== nothing
-            write_electron_state(scratch[1].pdf_electron, moments, t_params, time,
-                                 io_electron, t_params.moments_output_counter[], r, z,
-                                 vperp, vpa)
+            # Copy scratch arrays in order to save initial state - not optimal, but should
+            # be done at most once or twice while initialising a simulation
+            scratch[t_params.n_rk_stages+1].pdf_electron .= scratch[1].pdf_electron
+            scratch[t_params.n_rk_stages+1].electron_ppar .= scratch[1].electron_ppar
+            write_electron_state(scratch, moments, t_params, time, io_electron,
+                                 t_params.moments_output_counter[], r, z, vperp, vpa)
         end
     end
     # evolve (artificially) in time until the residual is less than the tolerance
@@ -498,8 +501,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             @serial_region begin
                 if io_electron !== nothing
                     t_params.write_moments_output[] = false
-                    write_electron_state(scratch[t_params.n_rk_stages+1].pdf_electron,
-                                         moments, t_params, time, io_electron,
+                    write_electron_state(scratch, moments, t_params, time, io_electron,
                                          t_params.moments_output_counter[], r, z, vperp,
                                          vpa)
                 end
@@ -540,9 +542,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         if !electron_pdf_converged || do_debug_io
             if io_electron !== nothing && io_electron !== true
                 t_params.moments_output_counter[] += 1
-                write_electron_state(final_scratch_pdf, moments, t_params, time,
-                                     io_electron, t_params.moments_output_counter[], r, z,
-                                     vperp, vpa)
+                write_electron_state(scratch, moments, t_params, time, io_electron,
+                                     t_params.moments_output_counter[], r, z, vperp, vpa)
                 finish_electron_io(io_electron)
             end
         end
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 84f68ede5..3b766ec98 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -204,9 +204,9 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3,
     # handle for the electron distribution function variable
     f_electron::Tfe
     # handle for the electron density variable
-    electron_density::Tmom
+    electron_density::Union{Tmom,Nothing}
     # handle for the electron parallel flow variable
-    electron_parallel_flow::Tmom
+    electron_parallel_flow::Union{Tmom,Nothing}
     # handle for the electron parallel pressure variable
     electron_parallel_pressure::Tmom
     # handle for the electron parallel heat flux variable
@@ -392,7 +392,9 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
-                                                      evolve_ppar, kinetic_electrons, t_params)
+                                                      evolve_ppar, kinetic_electrons,
+                                                      t_params;
+                                                      electron_only_io=true)
 
         close(fid)
 
@@ -1209,21 +1211,26 @@ define dynamic (time-evolving) electron moment variables for writing to the hdf5
 """
 function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordinate,
         parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar,
-        electron_physics, t_params)
+        electron_physics, t_params; electron_only_io=false)
 
     dynamic = get_group(fid, "dynamic_data")
 
-    # io_density is the handle for the ion particle density
-    io_electron_density = create_dynamic_variable!(dynamic, "electron_density", mk_float, z, r;
-                                          parallel_io=parallel_io,
-                                          description="electron species density",
-                                          units="n_ref")
-
-    # io_electron_upar is the handle for the electron parallel flow density
-    io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
-                                       parallel_io=parallel_io,
-                                       description="electron species parallel flow",
-                                       units="c_ref = sqrt(2*T_ref/mi)")
+    if !electron_only_io
+        # io_density is the handle for the ion particle density
+        io_electron_density = create_dynamic_variable!(dynamic, "electron_density", mk_float, z, r;
+                                              parallel_io=parallel_io,
+                                              description="electron species density",
+                                              units="n_ref")
+
+        # io_electron_upar is the handle for the electron parallel flow density
+        io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
+                                           parallel_io=parallel_io,
+                                           description="electron species parallel flow",
+                                           units="c_ref = sqrt(2*T_ref/mi)")
+    else
+        io_electron_density = nothing
+        io_electron_upar = nothing
+    end
 
     # io_electron_ppar is the handle for the electron parallel pressure
     io_electron_ppar = create_dynamic_variable!(dynamic, "electron_parallel_pressure", mk_float, z, r;
@@ -1848,7 +1855,7 @@ end
 write time-dependent moments data for ions, electrons and neutrals to the binary output
 file
 """
-function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
+function write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_species,
                                           n_neutral_species, io_or_file_info_moments,
                                           t_idx, time_for_run, t_params, nl_solver_params,
                                           r, z)
@@ -1871,13 +1878,14 @@ function write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
 
         write_em_fields_data_to_binary(fields, io_moments, t_idx, r, z)
 
-        write_ion_moments_data_to_binary(moments, n_ion_species, io_moments, t_idx, r, z)
+        write_ion_moments_data_to_binary(scratch, moments, n_ion_species, t_params,
+                                         io_moments, t_idx, r, z)
 
-        write_electron_moments_data_to_binary(moments, t_params.electron, io_moments,
-                                              t_idx, r, z)
+        write_electron_moments_data_to_binary(scratch, moments, t_params,
+                                              t_params.electron, io_moments, t_idx, r, z)
 
-        write_neutral_moments_data_to_binary(moments, n_neutral_species, io_moments,
-                                             t_idx, r, z)
+        write_neutral_moments_data_to_binary(scratch, moments, n_neutral_species,
+                                             t_params, io_moments, t_idx, r, z)
 
         append_to_dynamic_var(io_moments.time_for_run, time_for_run, t_idx, parallel_io)
         append_to_dynamic_var(io_moments.step_counter, t_params.step_counter[], t_idx, parallel_io)
@@ -1934,7 +1942,7 @@ write time-dependent moments data for ions to the binary output file
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_ion_moments_data_to_binary(moments, n_ion_species,
+function write_ion_moments_data_to_binary(scratch, moments, n_ion_species, t_params,
                                           io_moments::io_moments_info, t_idx, r, z)
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -1942,14 +1950,17 @@ function write_ion_moments_data_to_binary(moments, n_ion_species,
         parallel_io = io_moments.io_input.parallel_io
 
         # add the density data at this time slice to the output file
-        append_to_dynamic_var(io_moments.density, moments.ion.dens, t_idx,
-                              parallel_io, z, r, n_ion_species)
-        append_to_dynamic_var(io_moments.parallel_flow, moments.ion.upar, t_idx,
-                              parallel_io, z, r, n_ion_species)
-        append_to_dynamic_var(io_moments.parallel_pressure, moments.ion.ppar, t_idx,
-                              parallel_io, z, r, n_ion_species)
-        append_to_dynamic_var(io_moments.perpendicular_pressure, moments.ion.pperp, t_idx,
-                              parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.density, scratch[t_params.n_rk_stages+1].density,
+                              t_idx, parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.parallel_flow,
+                              scratch[t_params.n_rk_stages+1].upar, t_idx, parallel_io, z,
+                              r, n_ion_species)
+        append_to_dynamic_var(io_moments.parallel_pressure,
+                              scratch[t_params.n_rk_stages+1].ppar, t_idx, parallel_io, z,
+                              r, n_ion_species)
+        append_to_dynamic_var(io_moments.perpendicular_pressure,
+                              scratch[t_params.n_rk_stages+1].pperp, t_idx, parallel_io,
+                              z, r, n_ion_species)
         append_to_dynamic_var(io_moments.parallel_heat_flux, moments.ion.qpar, t_idx,
                               parallel_io, z, r, n_ion_species)
         append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth, t_idx,
@@ -2026,7 +2037,7 @@ write time-dependent moments data for electrons to the binary output file
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_electron_moments_data_to_binary(moments, t_params,
+function write_electron_moments_data_to_binary(scratch, moments, t_params, electron_t_params,
                                                io_moments::Union{io_moments_info,io_initial_electron_info},
                                                t_idx, r, z)
     @serial_region begin
@@ -2034,12 +2045,21 @@ function write_electron_moments_data_to_binary(moments, t_params,
 
         parallel_io = io_moments.io_input.parallel_io
 
-        append_to_dynamic_var(io_moments.electron_density, moments.electron.dens, t_idx,
-                              parallel_io, z, r)
-        append_to_dynamic_var(io_moments.electron_parallel_flow, moments.electron.upar,
-                              t_idx, parallel_io, z, r)
+        if io_moments.electron_density !== nothing
+            append_to_dynamic_var(io_moments.electron_density,
+                                  scratch[t_params.n_rk_stages+1].electron_density, t_idx,
+                                  parallel_io, z, r)
+        end
+
+        if io_moments.electron_parallel_flow !== nothing
+            append_to_dynamic_var(io_moments.electron_parallel_flow,
+                                  scratch[t_params.n_rk_stages+1].electron_upar, t_idx,
+                                  parallel_io, z, r)
+        end
+
         append_to_dynamic_var(io_moments.electron_parallel_pressure,
-                              moments.electron.ppar, t_idx, parallel_io, z, r)
+                              scratch[t_params.n_rk_stages+1].electron_ppar, t_idx,
+                              parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_parallel_heat_flux,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
@@ -2068,19 +2088,26 @@ function write_electron_moments_data_to_binary(moments, t_params,
                               moments.electron.constraints_C_coefficient, t_idx,
                               parallel_io, z, r)
 
-        if t_params !== nothing
+        if electron_t_params !== nothing
             # Save timestepping info
-            append_to_dynamic_var(io_moments.electron_step_counter, t_params.step_counter[], t_idx, parallel_io)
-            append_to_dynamic_var(io_moments.electron_dt, t_params.dt_before_output[], t_idx, parallel_io)
-            append_to_dynamic_var(io_moments.electron_failure_counter, t_params.failure_counter[], t_idx, parallel_io)
-            append_to_dynamic_var(io_moments.electron_failure_caused_by, t_params.failure_caused_by,
-                                  t_idx, parallel_io, length(t_params.failure_caused_by);
+            append_to_dynamic_var(io_moments.electron_step_counter,
+                                  electron_t_params.step_counter[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_dt,
+                                  electron_t_params.dt_before_output[], t_idx,
+                                  parallel_io)
+            append_to_dynamic_var(io_moments.electron_failure_counter,
+                                  electron_t_params.failure_counter[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_failure_caused_by,
+                                  electron_t_params.failure_caused_by, t_idx, parallel_io,
+                                  length(electron_t_params.failure_caused_by);
                                   only_root=true)
-            append_to_dynamic_var(io_moments.electron_limit_caused_by, t_params.limit_caused_by, t_idx,
-                                  parallel_io, length(t_params.limit_caused_by);
+            append_to_dynamic_var(io_moments.electron_limit_caused_by,
+                                  electron_t_params.limit_caused_by, t_idx, parallel_io,
+                                  length(electron_t_params.limit_caused_by);
                                   only_root=true)
             append_to_dynamic_var(io_moments.electron_dt_before_last_fail,
-                                  t_params.dt_before_last_fail[], t_idx, parallel_io)
+                                  electron_t_params.dt_before_last_fail[], t_idx,
+                                  parallel_io)
         end
     end
 
@@ -2092,8 +2119,9 @@ write time-dependent moments data for neutrals to the binary output file
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_neutral_moments_data_to_binary(moments, n_neutral_species,
-                                              io_moments::io_moments_info, t_idx, r, z)
+function write_neutral_moments_data_to_binary(scratch, moments, n_neutral_species,
+                                              t_params, io_moments::io_moments_info,
+                                              t_idx, r, z)
     if n_neutral_species ≤ 0
         return nothing
     end
@@ -2103,11 +2131,14 @@ function write_neutral_moments_data_to_binary(moments, n_neutral_species,
 
         parallel_io = io_moments.io_input.parallel_io
 
-        append_to_dynamic_var(io_moments.density_neutral, moments.neutral.dens, t_idx,
+        append_to_dynamic_var(io_moments.density_neutral,
+                              scratch[t_params.n_rk_stages+1].density_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
-        append_to_dynamic_var(io_moments.uz_neutral, moments.neutral.uz, t_idx,
+        append_to_dynamic_var(io_moments.uz_neutral,
+                              scratch[t_params.n_rk_stages+1].uz_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
-        append_to_dynamic_var(io_moments.pz_neutral, moments.neutral.pz, t_idx,
+        append_to_dynamic_var(io_moments.pz_neutral,
+                              scratch[t_params.n_rk_stages+1].pz_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
         append_to_dynamic_var(io_moments.qz_neutral, moments.neutral.qz, t_idx,
                               parallel_io, z, r, n_neutral_species)
@@ -2164,10 +2195,11 @@ end
 write time-dependent distribution function data for ions, electrons and neutrals to the
 binary output file
 """
-function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
-                                       n_neutral_species, io_or_file_info_dfns, t_idx,
-                                       time_for_run, t_params, nl_solver_params, r, z,
-                                       vperp, vpa, vzeta, vr, vz)
+function write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, t,
+                                       n_ion_species, n_neutral_species,
+                                       io_or_file_info_dfns, t_idx, time_for_run,
+                                       t_params, nl_solver_params, r, z, vperp, vpa,
+                                       vzeta, vr, vz)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
@@ -2181,18 +2213,18 @@ function write_all_dfns_data_to_binary(pdf, moments, fields, t, n_ion_species,
 
         # Write the moments for this time slice to the output file.
         # This also updates the time.
-        write_all_moments_data_to_binary(moments, fields, t, n_ion_species,
+        write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_species,
                                          n_neutral_species, io_dfns.io_moments, t_idx,
                                          time_for_run, t_params, nl_solver_params, r, z)
 
         # add the distribution function data at this time slice to the output file
-        write_ion_dfns_data_to_binary(pdf.ion.norm, n_ion_species, io_dfns, t_idx, r, z,
-                                      vperp, vpa)
-        if pdf.electron !== nothing
-            write_electron_dfns_data_to_binary(pdf.electron.norm, io_dfns, t_idx, r, z,
-                                               vperp, vpa)
+        write_ion_dfns_data_to_binary(scratch, t_params, n_ion_species, io_dfns, t_idx, r,
+                                      z, vperp, vpa)
+        if scratch_electron !== nothing
+            write_electron_dfns_data_to_binary(scratch_electron, t_params, io_dfns, t_idx,
+                                               r, z, vperp, vpa)
         end
-        write_neutral_dfns_data_to_binary(pdf.neutral.norm, n_neutral_species, io_dfns,
+        write_neutral_dfns_data_to_binary(scratch, t_params, n_neutral_species, io_dfns,
                                           t_idx, r, z, vzeta, vr, vz)
 
         closefile && close(io_dfns.fid)
@@ -2205,15 +2237,15 @@ write time-dependent distribution function data for ions to the binary output fi
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_ion_dfns_data_to_binary(ff, n_ion_species, io_dfns::io_dfns_info, t_idx, r,
-                                       z, vperp, vpa)
+function write_ion_dfns_data_to_binary(scratch, t_params, n_ion_species,
+                                       io_dfns::io_dfns_info, t_idx, r, z, vperp, vpa)
     @serial_region begin
         # Only read/write from first process in each 'block'
 
         parallel_io = io_dfns.io_input.parallel_io
 
-        append_to_dynamic_var(io_dfns.f, ff, t_idx, parallel_io, vpa, vperp, z, r,
-                              n_ion_species)
+        append_to_dynamic_var(io_dfns.f, scratch[t_params.n_rk_stages+1].pdf, t_idx,
+                              parallel_io, vpa, vperp, z, r, n_ion_species)
     end
     return nothing
 end
@@ -2223,7 +2255,7 @@ write time-dependent distribution function data for electrons to the binary outp
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_electron_dfns_data_to_binary(ff_electron,
+function write_electron_dfns_data_to_binary(scratch_electron, t_params,
                                             io_dfns::Union{io_dfns_info,io_initial_electron_info},
                                             t_idx, r, z, vperp, vpa)
     @serial_region begin
@@ -2232,8 +2264,9 @@ function write_electron_dfns_data_to_binary(ff_electron,
         parallel_io = io_dfns.io_input.parallel_io
 
         if io_dfns.f_electron !== nothing
-            append_to_dynamic_var(io_dfns.f_electron, ff_electron, t_idx, parallel_io,
-                                  vpa, vperp, z, r)
+            append_to_dynamic_var(io_dfns.f_electron,
+                                  scratch_electron[t_params.electron.n_rk_stages+1].pdf_electron,
+                                  t_idx, parallel_io, vpa, vperp, z, r)
         end
     end
     return nothing
@@ -2244,7 +2277,7 @@ write time-dependent distribution function data for neutrals to the binary outpu
 
 Note: should only be called from within a function that (re-)opens the output file.
 """
-function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
+function write_neutral_dfns_data_to_binary(scratch, t_params, n_neutral_species,
                                            io_dfns::io_dfns_info, t_idx, r, z, vzeta, vr,
                                            vz)
     @serial_region begin
@@ -2253,21 +2286,22 @@ function write_neutral_dfns_data_to_binary(ff_neutral, n_neutral_species,
         parallel_io = io_dfns.io_input.parallel_io
 
         if n_neutral_species > 0
-            append_to_dynamic_var(io_dfns.f_neutral, ff_neutral, t_idx, parallel_io, vz,
-                                  vr, vzeta, z, r, n_neutral_species)
+            append_to_dynamic_var(io_dfns.f_neutral,
+                                  scratch[t_params.n_rk_stages+1].pdf_neutral, t_idx,
+                                  parallel_io, vz, vr, vzeta, z, r, n_neutral_species)
         end
     end
     return nothing
 end
 
 """
-    write_electron_state(pdf, moments, t_params, t, io_initial_electron,
+    write_electron_state(scratch_electron, moments, t_params, t, io_initial_electron,
                          t_idx, r, z, vperp, vpa)
 
 Write the electron state to an output file.
 """
-function write_electron_state(pdf, moments, t_params, t, io_or_file_info_initial_electron,
-                              t_idx, r, z, vperp, vpa)
+function write_electron_state(scratch_electron, moments, t_params, t,
+                              io_or_file_info_initial_electron, t_idx, r, z, vperp, vpa)
 
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -2285,11 +2319,11 @@ function write_electron_state(pdf, moments, t_params, t, io_or_file_info_initial
         # add the pseudo-time for this time slice to the hdf5 file
         append_to_dynamic_var(io_initial_electron.pseudotime, t, t_idx, parallel_io)
 
-        write_electron_dfns_data_to_binary(pdf, io_initial_electron, t_idx, r, z, vperp,
-                                           vpa)
+        write_electron_dfns_data_to_binary(scratch_electron, t_params,
+                                           io_initial_electron, t_idx, r, z, vperp, vpa)
 
-        write_electron_moments_data_to_binary(moments, t_params, io_initial_electron,
-                                              t_idx, r, z)
+        write_electron_moments_data_to_binary(scratch_electron, moments, t_params,
+                                              t_params, io_initial_electron, t_idx, r, z)
 
         closefile && close(io_initial_electron.fid)
     end
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index e7b1a7d2e..6111468a8 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -547,11 +547,12 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
     return nothing
 end
 
-function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp, vzeta, vr,
-                                  vz, z_spectral, vperp_spectral, vpa_spectral, z_advect,
-                                  vpa_advect, scratch_dummy, collisions, composition,
-                                  geometry, external_source_settings, num_diss_params,
-                                  t_params, t_input, io_input, input_dict)
+function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa, vperp,
+                                  vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
+                                  z_advect, vpa_advect, scratch_dummy, collisions,
+                                  composition, geometry, external_source_settings,
+                                  num_diss_params, t_params, t_input, io_input,
+                                  input_dict)
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
@@ -604,8 +605,8 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
 
         begin_serial_region()
         @serial_region begin
-            # update the electron pdf in the first scratch
-            scratch[1].pdf_electron .= pdf.electron.norm
+            # update the electron pdf in the first scratch_electron
+            scratch_electron[1].pdf_electron .= pdf.electron.norm
         end
 
         begin_r_z_region()
@@ -655,8 +656,8 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             t_params.dfns_output_times .= truncated_times
         end
         electron_pseudotime, success =
-            @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
-                                        vperp, vpa, z_spectral, vperp_spectral,
+            @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
+                                        r, z, vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                         t_params, collisions, composition,
                                         external_source_settings, num_diss_params,
@@ -676,8 +677,8 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
             println("Initializing electrons - evolving pdf_electron only to steady state")
         end
         electron_pseudotime, success =
-            @views update_electron_pdf!(scratch, pdf.electron.norm, moments, phi, r, z,
-                                        vperp, vpa, z_spectral, vperp_spectral,
+            @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
+                                        r, z, vperp, vpa, z_spectral, vperp_spectral,
                                         vpa_spectral, z_advect, vpa_advect, scratch_dummy,
                                         t_params, collisions, composition,
                                         external_source_settings, num_diss_params,
@@ -698,7 +699,7 @@ function initialize_electron_pdf!(scratch, pdf, moments, phi, r, z, vpa, vperp,
         # Write the converged initial state for the electrons to a file so that it can be
         # re-used if the simulation is re-run.
         t_params.moments_output_counter[] += 1
-        write_electron_state(pdf.electron.norm, moments, t_params, electron_pseudotime,
+        write_electron_state(scratch_electron, moments, t_params, electron_pseudotime,
                              io_initial_electron, t_params.moments_output_counter[], r, z,
                              vperp, vpa)
         finish_electron_io(io_initial_electron)
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 9051eee8c..e96ee7c41 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -1067,17 +1067,6 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
                 get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp,
                                   restart_vpa, restart_vzeta, restart_vr, restart_vz)
 
-            moments.electron.dens .=
-                reload_electron_moment("electron_density", dynamic, time_index, r, z,
-                                       r_range, z_range, restart_r, restart_r_spectral,
-                                       restart_z, restart_z_spectral,
-                                       interpolation_needed)
-            moments.electron.dens_updated[] = true
-            moments.electron.upar .=
-                reload_electron_moment("electron_parallel_flow", dynamic, time_index, r,
-                                       z, r_range, z_range, restart_r, restart_r_spectral,
-                                       restart_z, restart_z_spectral,
-                                       interpolation_needed)
             moments.electron.upar_updated[] = true
             moments.electron.ppar .=
                 reload_electron_moment("electron_parallel_pressure", dynamic, time_index,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 06d65f922..93e3d3598 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -368,10 +368,10 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
-    write_all_moments_data_to_binary(moments, fields, code_time,
+    write_all_moments_data_to_binary(scratch, moments, fields, code_time,
         composition.n_ion_species, composition.n_neutral_species, io_moments, 1, 0.0,
         t_params, nl_solver_params, r, z)
-    write_all_dfns_data_to_binary(pdf, moments, fields, code_time,
+    write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, code_time,
         composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, 0.0,
         t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 375b20cd0..d4fffa612 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -893,14 +893,22 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
                                     composition)
         end
 
-        # update scratch arrays in case they were affected by applying boundary conditions
-        # or constraints to the pdf
+        # Update scratch arrays in case they were affected by applying boundary conditions
+        # or constraints to the pdf.
+        # Also update scratch[t_params.n_rk_stages+1] as this will be used for the I/O at
+        # the initial time.
         begin_s_r_z_region()
         @loop_s_r_z is ir iz begin
             scratch[1].pdf[:,:,iz,ir,is] .= pdf.ion.norm[:,:,iz,ir,is]
             scratch[1].density[iz,ir,is] = moments.ion.dens[iz,ir,is]
             scratch[1].upar[iz,ir,is] = moments.ion.upar[iz,ir,is]
             scratch[1].ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
+            scratch[1].pperp[iz,ir,is] = moments.ion.pperp[iz,ir,is]
+            scratch[t_params.n_rk_stages+1].pdf[:,:,iz,ir,is] .= pdf.ion.norm[:,:,iz,ir,is]
+            scratch[t_params.n_rk_stages+1].density[iz,ir,is] = moments.ion.dens[iz,ir,is]
+            scratch[t_params.n_rk_stages+1].upar[iz,ir,is] = moments.ion.upar[iz,ir,is]
+            scratch[t_params.n_rk_stages+1].ppar[iz,ir,is] = moments.ion.ppar[iz,ir,is]
+            scratch[t_params.n_rk_stages+1].pperp[iz,ir,is] = moments.ion.pperp[iz,ir,is]
         end
 
         # update the electron density, parallel flow and parallel pressure (and temperature)
@@ -939,13 +947,19 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
                 moments.electron.ppar, moments.electron.upar, moments.electron.dens,
                 moments.electron, z)
         end
-        # update the electron moment entries in the scratch array
+        # Update the electron moment entries in the scratch array.
+        # Also update scratch[t_params.n_rk_stages+1] as this will be used for the I/O at
+        # the initial time.
         begin_r_z_region()
         @loop_r_z ir iz begin
             scratch[1].electron_density[iz,ir] = moments.electron.dens[iz,ir]
             scratch[1].electron_upar[iz,ir] = moments.electron.upar[iz,ir]
             scratch[1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
             scratch[1].electron_temp[iz,ir] = moments.electron.temp[iz,ir]
+            scratch[t_params.n_rk_stages+1].electron_density[iz,ir] = moments.electron.dens[iz,ir]
+            scratch[t_params.n_rk_stages+1].electron_upar[iz,ir] = moments.electron.upar[iz,ir]
+            scratch[t_params.n_rk_stages+1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+            scratch[t_params.n_rk_stages+1].electron_temp[iz,ir] = moments.electron.temp[iz,ir]
         end
 
         begin_sn_r_z_region(no_synchronize=true)
@@ -954,6 +968,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
             scratch[1].density_neutral[iz,ir,isn] = moments.neutral.dens[iz,ir,isn]
             scratch[1].uz_neutral[iz,ir,isn] = moments.neutral.uz[iz,ir,isn]
             scratch[1].pz_neutral[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
+            scratch[t_params.n_rk_stages+1].pdf_neutral[:,:,:,iz,ir,isn] .= pdf.neutral.norm[:,:,:,iz,ir,isn]
+            scratch[t_params.n_rk_stages+1].density_neutral[iz,ir,isn] = moments.neutral.dens[iz,ir,isn]
+            scratch[t_params.n_rk_stages+1].uz_neutral[iz,ir,isn] = moments.neutral.uz[iz,ir,isn]
+            scratch[t_params.n_rk_stages+1].pz_neutral[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
         end
     end
     # calculate the electron-ion parallel friction force
@@ -1752,7 +1770,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, t,
                                 composition.n_ion_species, composition.n_neutral_species,
                                 ascii_io)
-            write_all_moments_data_to_binary(moments, fields, t,
+            write_all_moments_data_to_binary(scratch, moments, fields, t,
                                              composition.n_ion_species,
                                              composition.n_neutral_species, io_moments,
                                              iwrite_moments, time_for_run, t_params,
@@ -1836,7 +1854,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                     flush(stdout)
                 end
             end
-            write_all_dfns_data_to_binary(pdf, moments, fields, t,
+            write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, t,
                                           composition.n_ion_species,
                                           composition.n_neutral_species, io_dfns,
                                           iwrite_dfns, time_for_run, t_params,

From 301fc6b8e3d8c0c5eb089f5ddfa4e298ca1c8cc7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 18:23:29 +0100
Subject: [PATCH 305/428] Allow output of Enum values by converting them to
 Strings

---
 moment_kinetics/src/file_io.jl       |  9 ++++++++-
 moment_kinetics/src/input_structs.jl | 29 +++++++++++++++++++---------
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 3b766ec98..2bffadb0e 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -493,6 +493,11 @@ attributes of the variable.
 """
 function write_single_value! end
 
+# Convert Enum values to String to be written to file
+function write_single_value!(file_or_group, name, data::Enum; kwargs...)
+    return write_single_value!(file_or_group, name, string(data); kwargs...)
+end
+
 """
 write some overview information for the simulation to the binary file
 """
@@ -559,7 +564,9 @@ function write_provenance_tracking_info!(fid, parallel_io, run_id, restart_time_
         # Convert input_dict into a TOML-formatted string so that we can store it in a
         # single variable.
         io_buffer = IOBuffer()
-        TOML.print(io_buffer, input_dict)
+        # The `mk_to_toml` function allows converting extra types (e.g. Enum) to things
+        # that can be printed to a TOML string/file.
+        TOML.print(mk_to_toml, io_buffer, input_dict)
         input_string = String(take!(io_buffer))
         write_single_value!(provenance_tracking, "input", input_string,
                             parallel_io=parallel_io,
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index acf7d0b26..6910347c4 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -8,6 +8,7 @@ export time_info
 export advection_input, advection_input_mutable
 export grid_input, grid_input_mutable
 export initial_condition_input, initial_condition_input_mutable
+export mk_to_toml
 export species_parameters, species_parameters_mutable
 export species_composition
 export drive_input, drive_input_mutable
@@ -651,19 +652,32 @@ Utility method for converting a string to an Enum when getting from a Dict, base
 type of the default value
 """
 function get(d::Dict, key, default::Enum)
-    valstring = get(d, key, nothing)
-    if valstring == nothing
+    val_maybe_string = get(d, key, nothing)
+    if val_maybe_string == nothing
         return default
+    elseif isa(val_maybe_string, Enum)
+        return val_maybe_string
     # instances(typeof(default)) gets the possible values of the Enum. Then convert to
     # Symbol, then to String.
-    elseif valstring ∈ (split(s, ".")[end] for s ∈ String.(Symbol.(instances(typeof(default)))))
-        return eval(Symbol(valstring))
+    elseif val_maybe_string ∈ Tuple(split(s, ".")[end] for s ∈ string.(instances(typeof(default))))
+        return eval(Symbol(val_maybe_string))
     else
-        error("Expected a $(typeof(default)), but '$valstring' is not in "
+        error("Expected a $(typeof(default)), but '$val_maybe_string' is not in "
               * "$(instances(typeof(default)))")
     end
 end
 
+"""
+Convert some types used by moment_kinetics to types that are supported by TOML
+"""
+function mk_to_toml(value)
+    if isa(value, Enum)
+        return string(value)
+    else
+        return value
+    end
+end
+
 """
 Set the defaults for options in the top level of the input, and check that there are not
 any unexpected options (i.e. options that have no default).
@@ -733,12 +747,9 @@ function set_defaults_and_check_section!(options::AbstractDict, section_name;
     end
 
     # Set default values if a key was not set explicitly
-    explicit_keys = keys(section)
     for (key_sym, value) ∈ kwargs
         key = String(key_sym)
-        if !(key ∈ explicit_keys)
-            section[key] = value
-        end
+        section[key] = get(section, key, value)
     end
 
     return section

From 0708a21545fa42288c5b4df2d5bc725a9bc34c73 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 18:06:24 +0100
Subject: [PATCH 306/428] Use set_defaults_and_check_section!() to read
 [output] settings

---
 moment_kinetics/src/moment_kinetics_input.jl | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index d8bcd18eb..beda2ace9 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -703,12 +703,17 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     drive_immutable = drive_input(drive.force_phi, drive.amplitude, drive.frequency, force_Er_zero)
 
     # inputs for file I/O
+    io_settings = set_defaults_and_check_section!(
+        scan_input, "output";
+        ascii_output=false,
+        binary_format=hdf5,
+        parallel_io=nothing,
+       )
+    if io_settings["parallel_io"] === nothing
+        io_settings["parallel_io"] = io_has_parallel(Val(io_settings["binary_format"]))
+    end
     # Make copy of the section to avoid modifying the passed-in Dict
-    io_settings = copy(get(scan_input, "output", Dict{String,Any}()))
-    io_settings["ascii_output"] = get(io_settings, "ascii_output", false)
-    io_settings["binary_format"] = get(io_settings, "binary_format", hdf5)
-    io_settings["parallel_io"] = get(io_settings, "parallel_io",
-                                     io_has_parallel(Val(io_settings["binary_format"])))
+    io_settings = copy(io_settings)
     run_id = string(uuid4())
     if !ignore_MPI
         # Communicate run_id to all blocks
@@ -718,8 +723,9 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         run_id = string(run_id_chars...)
     end
     io_settings["run_id"] = run_id
-    io_immutable = io_input(; output_dir=output_dir, run_name=run_name,
-                              Dict(Symbol(k)=>v for (k,v) in io_settings)...)
+    io_settings["output_dir"] = output_dir
+    io_settings["run_name"] = run_name
+    io_immutable = Dict_to_NamedTuple(io_settings)
 
     # initialize z grid and write grid point locations to file
     if ignore_MPI

From c9fe38287d13c3dd960e72225860205f9d85d0e7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 20:49:54 +0100
Subject: [PATCH 307/428] If an 'io variable' is `nothing`, just skip it

It is optional to write some variables. When this happens the
corresponding 'io variable' is set to `nothing`, so when
`append_to_dynamic_var()` is called with `data::Nothing`, just do
nothing.
---
 moment_kinetics/src/file_io.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 2bffadb0e..97fe50262 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1852,6 +1852,11 @@ each file is only written by one process).
 """
 function append_to_dynamic_var end
 
+function append_to_dynamic_var(data::Nothing, args...; kwargs...)
+    # Variable was not created to save, so nothing to do.
+    return nothing
+end
+
 @debug_shared_array begin
     function append_to_dynamic_var(data::DebugMPISharedArray, args...; kwargs...)
         return append_to_dynamic_var(data.data, args...; kwargs...)

From e62bff8192977e605f2a3b33c3083939c5c6a19c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 20:52:40 +0100
Subject: [PATCH 308/428] Option to write out info for error and/or steady
 state diagnostics

---
 moment_kinetics/src/file_io.jl                | 635 ++++++++++++++++--
 moment_kinetics/src/moment_kinetics_input.jl  |   8 +
 moment_kinetics/src/time_advance.jl           |  12 +
 .../test/recycling_fraction_tests.jl          |   4 +-
 4 files changed, 613 insertions(+), 46 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 97fe50262..08e49e635 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -71,12 +71,28 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     Ez::Tphi
     # handle for the ion species density
     density::Tmomi
+    # low-order approximation, used to diagnose timestepping error
+    density_loworder::Union{Tmomi,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    density_start_last_timestep::Union{Tmomi,Nothing}
     # handle for the ion species parallel flow
     parallel_flow::Tmomi
+    # low-order approximation, used to diagnose timestepping error
+    parallel_flow_loworder::Union{Tmomi,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    parallel_flow_start_last_timestep::Union{Tmomi,Nothing}
     # handle for the ion species parallel pressure
     parallel_pressure::Tmomi
+    # low-order approximation, used to diagnose timestepping error
+    parallel_pressure_loworder::Union{Tmomi,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    parallel_pressure_start_last_timestep::Union{Tmomi,Nothing}
     # handle for the ion species perpendicular pressure
     perpendicular_pressure::Tmomi
+    # low-order approximation, used to diagnose timestepping error
+    perpendicular_pressure_loworder::Union{Tmomi,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    perpendicular_pressure_start_last_timestep::Union{Tmomi,Nothing}
     # handle for the ion species parallel heat flux
     parallel_heat_flux::Tmomi
     # handle for the ion species thermal speed
@@ -89,10 +105,22 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     chodura_integral_upper::Tchodura_upper
     # handle for the electron species density
     electron_density::Tmome
+    # low-order approximation, used to diagnose timestepping error
+    electron_density_loworder::Union{Tmome,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_density_start_last_timestep::Union{Tmome,Nothing}
     # handle for the electron species parallel flow
     electron_parallel_flow::Tmome
+    # low-order approximation, used to diagnose timestepping error
+    electron_parallel_flow_loworder::Union{Tmome,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_parallel_flow_start_last_timestep::Union{Tmome,Nothing}
     # handle for the electron species parallel pressure
     electron_parallel_pressure::Tmome
+    # low-order approximation, used to diagnose timestepping error
+    electron_parallel_pressure_loworder::Union{Tmome,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_parallel_pressure_start_last_timestep::Union{Tmome,Nothing}
     # handle for the electron species parallel heat flux
     electron_parallel_heat_flux::Tmome
     # handle for the electron species thermal speed
@@ -100,8 +128,20 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
 
     # handle for the neutral species density
     density_neutral::Tmomn
+    # low-order approximation, used to diagnose timestepping error
+    density_neutral_loworder::Union{Tmomn,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    density_neutral_start_last_timestep::Union{Tmomn,Nothing}
     uz_neutral::Tmomn
+    # low-order approximation, used to diagnose timestepping error
+    uz_neutral_loworder::Union{Tmomn,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    uz_neutral_start_last_timestep::Union{Tmomn,Nothing}
     pz_neutral::Tmomn
+    # low-order approximation, used to diagnose timestepping error
+    pz_neutral_loworder::Union{Tmomn,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    pz_neutral_start_last_timestep::Union{Tmomn,Nothing}
     qz_neutral::Tmomn
     thermal_speed_neutral::Tmomn
 
@@ -178,10 +218,22 @@ struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments, Tinput}
     fid::Tfile
     # handle for the ion species distribution function variable
     f::Tfi
+    # low-order approximation to ion species distribution function, used to diagnose timestepping error
+    f_loworder::Union{Tfi,Nothing}
+    # ion species distribution function at the start of the last timestep before output, used to measure steady state residual
+    f_start_last_timestep::Union{Tfi,Nothing}
     # handle for the electron distribution function variable
     f_electron::Tfe
+    # low-order approximation to electron distribution function, used to diagnose timestepping error
+    f_electron_loworder::Union{Tfe,Nothing}
+    # electron distribution function at the start of the last timestep before output, used to measure steady state residual
+    f_electron_start_last_timestep::Union{Tfe,Nothing}
     # handle for the neutral species distribution function variable
     f_neutral::Tfn
+    # low-order approximation to neutral species distribution function, used to diagnose timestepping error
+    f_neutral_loworder::Union{Tfn,Nothing}
+    # neutral species distribution function at the start of the last timestep before output, used to measure steady state residual
+    f_neutral_start_last_timestep::Union{Tfn,Nothing}
 
     # Settings for I/O
     io_input::Tinput
@@ -203,12 +255,28 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3,
     pseudotime::Ttime
     # handle for the electron distribution function variable
     f_electron::Tfe
+    # low-order approximation, used to diagnose timestepping error
+    f_electron_loworder::Union{Tfe,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    f_electron_start_last_timestep::Union{Tfe,Nothing}
     # handle for the electron density variable
     electron_density::Union{Tmom,Nothing}
+    # low-order approximation, used to diagnose timestepping error
+    electron_density_loworder::Union{Tmom,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_density_start_last_timestep::Union{Tmom,Nothing}
     # handle for the electron parallel flow variable
     electron_parallel_flow::Union{Tmom,Nothing}
+    # low-order approximation, used to diagnose timestepping error
+    electron_parallel_flow_loworder::Union{Tmom,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_parallel_flow_start_last_timestep::Union{Tmom,Nothing}
     # handle for the electron parallel pressure variable
     electron_parallel_pressure::Tmom
+    # low-order approximation, used to diagnose timestepping error
+    electron_parallel_pressure_loworder::Union{Tmom,Nothing}
+    # start of the last timestep before output, used to measure steady state residual
+    electron_parallel_pressure_start_last_timestep::Union{Tmom,Nothing}
     # handle for the electron parallel heat flux variable
     electron_parallel_heat_flux::Tmom
     # handle for the electron thermal speed variable
@@ -379,8 +447,31 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
                                                  vperp, z, r;
                                                  parallel_io=parallel_io,
                                                  description="electron distribution function")
+        if io_input.write_electron_error_diagnostics
+            io_f_electron_loworder =
+                create_dynamic_variable!(dynamic, "f_electron_loworder", mk_float,
+                                         vpa, vperp, z, r;
+                                         n_ion_species=composition.n_ion_species,
+                                         parallel_io=parallel_io,
+                                         description="low-order approximation to electron distribution function, used to diagnose timestepping error")
+        else
+            io_f_electron_loworder = nothing
+        end
+        if io_input.write_electron_steady_state_diagnostics
+            io_f_electron_start_last_timestep =
+                create_dynamic_variable!(dynamic, "f_electron_start_last_timestep",
+                                         mk_float, vpa, vperp, z, r;
+                                         n_ion_species=composition.n_ion_species,
+                                         parallel_io=parallel_io,
+                                         description="electron distribution function at the start of the last electron pseudo-timestep before output, used to measure steady state residual")
+        else
+            io_f_electron_start_last_timestep = nothing
+        end
 
-        io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
+        io_electron_density, io_electron_density_loworder,
+        io_electron_density_start_last_timestep, io_electron_upar,
+        io_electron_upar_loworder, io_electron_upar_start_last_timestep, io_electron_ppar,
+        io_electron_ppar_loworder, io_electron_ppar_start_last_timestep, io_electron_qpar,
         io_electron_vth, external_source_electron_amplitude,
         external_source_electron_density_amplitude,
         external_source_electron_momentum_amplitude,
@@ -393,7 +484,9 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
                                                       evolve_ppar, kinetic_electrons,
-                                                      t_params;
+                                                      t_params,
+                                                      io_input.write_electron_error_diagnostics,
+                                                      io_input.write_electron_steady_state_diagnostics;
                                                       electron_only_io=true)
 
         close(fid)
@@ -439,9 +532,17 @@ function reopen_initial_electron_io(file_info)
             end
         end
         return io_initial_electron_info(fid, getvar("time"), getvar("f_electron"),
+                                        getvar("f_electron_loworder"),
+                                        getvar("f_electron_start_last_timestep"),
                                         getvar("electron_density"),
+                                        getvar("electron_density_loworder"),
+                                        getvar("electron_density_start_last_timestep"),
                                         getvar("electron_parallel_flow"),
+                                        getvar("electron_parallel_flow_loworder"),
+                                        getvar("electron_parallel_flow_start_last_timestep"),
                                         getvar("electron_parallel_pressure"),
+                                        getvar("electron_parallel_pressure_loworder"),
+                                        getvar("electron_parallel_pressure_start_last_timestep"),
                                         getvar("electron_parallel_heat_flux"),
                                         getvar("electron_thermal_speed"),
                                         getvar("external_source_electron_amplitude"),
@@ -892,17 +993,24 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         io_phi, io_Er, io_Ez =
             define_dynamic_em_field_variables!(fid, r, z, parallel_io)
 
-        io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
-        external_source_amplitude, external_source_density_amplitude,
-        external_source_momentum_amplitude, external_source_pressure_amplitude,
-        external_source_controller_integral, io_chodura_lower, io_chodura_upper,
-        ion_constraints_A_coefficient, ion_constraints_B_coefficient,
-        ion_constraints_C_coefficient =
+        io_density, io_density_loworder, io_density_start_last_timestep, io_upar,
+        io_upar_loworder, io_upar_start_last_timestep, io_ppar, io_ppar_loworder,
+        io_ppar_start_last_timestep, io_pperp, io_pperp_loworder,
+        io_pperp_start_last_timestep, io_qpar, io_vth, io_dSdt, external_source_amplitude,
+        external_source_density_amplitude, external_source_momentum_amplitude,
+        external_source_pressure_amplitude, external_source_controller_integral,
+        io_chodura_lower, io_chodura_upper, ion_constraints_A_coefficient,
+        ion_constraints_B_coefficient, ion_constraints_C_coefficient =
             define_dynamic_ion_moment_variables!(fid, n_ion_species, r, z, parallel_io,
                                                  external_source_settings, evolve_density,
-                                                 evolve_upar, evolve_ppar)
-
-        io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
+                                                 evolve_upar, evolve_ppar,
+                                                 io_input.write_error_diagnostics,
+                                                 io_input.write_steady_state_diagnostics)
+
+        io_electron_density, io_electron_density_loworder,
+        io_electron_density_start_last_timestep, io_electron_upar,
+        io_electron_upar_loworder, io_electron_upar_start_last_timestep, io_electron_ppar,
+        io_electron_ppar_loworder, io_electron_ppar_start_last_timestep, io_electron_qpar,
         io_electron_vth, external_source_electron_amplitude,
         external_source_electron_density_amplitude,
         external_source_electron_momentum_amplitude,
@@ -915,11 +1023,15 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
                                                       evolve_ppar, electron_physics,
-                                                      t_params.electron)
-
-        io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
-        io_thermal_speed_neutral, external_source_neutral_amplitude,
-        external_source_neutral_density_amplitude,
+                                                      t_params.electron,
+                                                      io_input.write_error_diagnostics,
+                                                      io_input.write_steady_state_diagnostics)
+
+        io_density_neutral, io_density_neutral_loworder,
+        io_density_neutral_start_last_timestep, io_uz_neutral, io_uz_neutral_loworder,
+        io_uz_neutral_start_last_timestep, io_pz_neutral, io_pz_neutral_loworder,
+        io_pz_neutral_start_last_timestep, io_qz_neutral, io_thermal_speed_neutral,
+        external_source_neutral_amplitude, external_source_neutral_density_amplitude,
         external_source_neutral_momentum_amplitude,
         external_source_neutral_pressure_amplitude,
         external_source_neutral_controller_integral, neutral_constraints_A_coefficient,
@@ -928,7 +1040,9 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                                      parallel_io,
                                                      external_source_settings,
                                                      evolve_density, evolve_upar,
-                                                     evolve_ppar)
+                                                     evolve_ppar,
+                                                     io_input.write_error_diagnostics,
+                                                     io_input.write_steady_state_diagnostics)
 
         io_time_for_run = create_dynamic_variable!(
             dynamic, "time_for_run", mk_float; parallel_io=parallel_io,
@@ -980,11 +1094,26 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                   )
             for (term, params) ∈ pairs(nl_solver_params) if params !== nothing)
 
-        return io_moments_info(fid, io_time, io_phi, io_Er, io_Ez, io_density, io_upar,
-                               io_ppar, io_pperp, io_qpar, io_vth, io_dSdt, io_chodura_lower, io_chodura_upper,
-                               io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar, io_electron_vth,
-                               io_density_neutral, io_uz_neutral,
-                               io_pz_neutral, io_qz_neutral, io_thermal_speed_neutral,
+        return io_moments_info(fid, io_time, io_phi, io_Er, io_Ez, io_density,
+                               io_density_loworder, io_density_start_last_timestep,
+                               io_upar, io_upar_loworder, io_upar_start_last_timestep,
+                               io_ppar, io_ppar_loworder, io_ppar_start_last_timestep,
+                               io_pperp, io_pperp_loworder, io_pperp_start_last_timestep,
+                               io_qpar, io_vth, io_dSdt, io_chodura_lower,
+                               io_chodura_upper,
+                               io_electron_density, io_electron_density_loworder,
+                               io_electron_density_start_last_timestep, io_electron_upar,
+                               io_electron_upar_loworder,
+                               io_electron_upar_start_last_timestep, io_electron_ppar,
+                               io_electron_ppar_loworder,
+                               io_electron_ppar_start_last_timestep, io_electron_qpar,
+                               io_electron_vth,
+                               io_density_neutral, io_density_neutral_loworder,
+                               io_density_neutral_start_last_timestep, io_uz_neutral,
+                               io_uz_neutral_loworder, io_uz_neutral_start_last_timestep,
+                               io_pz_neutral, io_pz_neutral_loworder,
+                               io_pz_neutral_start_last_timestep, io_qz_neutral,
+                               io_thermal_speed_neutral,
                                external_source_amplitude,
                                external_source_density_amplitude,
                                external_source_momentum_amplitude,
@@ -1055,7 +1184,7 @@ define dynamic (time-evolving) ion moment variables for writing to the hdf5 file
 """
 function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
         z::coordinate, parallel_io, external_source_settings, evolve_density, evolve_upar,
-        evolve_ppar)
+        evolve_ppar, write_error_diagnostics, write_steady_state_diagnostics)
 
     dynamic = get_group(fid, "dynamic_data")
 
@@ -1065,6 +1194,24 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
                                           parallel_io=parallel_io,
                                           description="ion species density",
                                           units="n_ref")
+    if write_error_diagnostics
+        io_density_loworder =
+            create_dynamic_variable!(dynamic, "density_loworder", mk_float, z, r;
+                                     n_ion_species=n_ion_species, parallel_io=parallel_io,
+                                     description="low-order approximation to ion species density, used to diagnose timestepping error",
+                                     units="n_ref")
+    else
+        io_density_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_density_start_last_timestep =
+            create_dynamic_variable!(dynamic, "density_start_last_timestep", mk_float, z, r;
+                                     n_ion_species=n_ion_species, parallel_io=parallel_io,
+                                     description="ion species density at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref")
+    else
+        io_density_start_last_timestep = nothing
+    end
 
     # io_upar is the handle for the ion parallel flow density
     io_upar = create_dynamic_variable!(dynamic, "parallel_flow", mk_float, z, r;
@@ -1072,6 +1219,25 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
                                        parallel_io=parallel_io,
                                        description="ion species parallel flow",
                                        units="c_ref = sqrt(2*T_ref/mi)")
+    if write_error_diagnostics
+        io_upar_loworder =
+            create_dynamic_variable!(dynamic, "parallel_flow_loworder", mk_float, z, r;
+                                     n_ion_species=n_ion_species, parallel_io=parallel_io,
+                                     description="low-order approximation to ion species parallel flow, used to diagnose timestepping error",
+                                     units="c_ref = sqrt(2*T_ref/mi)")
+    else
+        io_upar_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_upar_start_last_timestep =
+            create_dynamic_variable!(dynamic, "parallel_flow_start_last_timestep",
+                                     mk_float, z, r; n_ion_species=n_ion_species,
+                                     parallel_io=parallel_io,
+                                     description="ion species parallel flow at the start of the last timestep before output, used to measure steady state residual",
+                                     units="c_ref = sqrt(2*T_ref/mi)")
+    else
+        io_upar_start_last_timestep = nothing
+    end
 
     # io_ppar is the handle for the ion parallel pressure
     io_ppar = create_dynamic_variable!(dynamic, "parallel_pressure", mk_float, z, r;
@@ -1079,6 +1245,25 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
                                        parallel_io=parallel_io,
                                        description="ion species parallel pressure",
                                        units="n_ref*T_ref")
+    if write_error_diagnostics
+        io_ppar_loworder =
+            create_dynamic_variable!(dynamic, "parallel_pressure_loworder", mk_float, z, r;
+                                     n_ion_species=n_ion_species, parallel_io=parallel_io,
+                                     description="low-order approximation to ion species parallel pressure, used to diagnose timestepping error",
+                                     units="n_ref*T_ref")
+    else
+        io_ppar_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_ppar_start_last_timestep =
+            create_dynamic_variable!(dynamic, "parallel_pressure_start_last_timestep",
+                                     mk_float, z, r; n_ion_species=n_ion_species,
+                                     parallel_io=parallel_io,
+                                     description="ion species parallel pressure at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref*T_ref")
+    else
+        io_ppar_start_last_timestep = nothing
+    end
 
     # io_pperp is the handle for the ion parallel pressure
     io_pperp = create_dynamic_variable!(dynamic, "perpendicular_pressure", mk_float, z, r;
@@ -1086,6 +1271,25 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
                                         parallel_io=parallel_io,
                                         description="ion species perpendicular pressure",
                                         units="n_ref*T_ref")
+    if write_error_diagnostics
+        io_pperp_loworder =
+            create_dynamic_variable!(dynamic, "perpendicular_pressure_loworder", mk_float, z,
+                                     r; n_ion_species=n_ion_species, parallel_io=parallel_io,
+                                     description="low-order approximation to ion species perpendicular pressure, used to diagnose timestepping error",
+                                     units="n_ref*T_ref")
+    else
+        io_pperp_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_pperp_start_last_timestep =
+            create_dynamic_variable!(dynamic, "perpendicular_pressure_start_last_timestep",
+                                     mk_float, z, r; n_ion_species=n_ion_species,
+                                     parallel_io=parallel_io,
+                                     description="ion species perpendicular pressure at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref*T_ref")
+    else
+        io_pperp_start_last_timestep = nothing
+    end
 
     # io_qpar is the handle for the ion parallel heat flux
     io_qpar = create_dynamic_variable!(dynamic, "parallel_heat_flux", mk_float, z, r;
@@ -1205,7 +1409,10 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
            ion_constraints_C_coefficient = nothing
     end
 
-    return io_density, io_upar, io_ppar, io_pperp, io_qpar, io_vth, io_dSdt,
+    return io_density, io_density_loworder, io_density_start_last_timestep, io_upar,
+           io_upar_loworder, io_upar_start_last_timestep, io_ppar, io_ppar_loworder,
+           io_ppar_start_last_timestep, io_pperp, io_pperp_loworder,
+           io_pperp_start_last_timestep, io_qpar, io_vth, io_dSdt,
            external_source_amplitude, external_source_density_amplitude,
            external_source_momentum_amplitude, external_source_pressure_amplitude,
            external_source_controller_integral, io_chodura_lower, io_chodura_upper,
@@ -1218,7 +1425,8 @@ define dynamic (time-evolving) electron moment variables for writing to the hdf5
 """
 function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordinate,
         parallel_io, external_source_settings, evolve_density, evolve_upar, evolve_ppar,
-        electron_physics, t_params; electron_only_io=false)
+        electron_physics, t_params, write_error_diagnostics,
+        write_steady_state_diagnostics; electron_only_io=false)
 
     dynamic = get_group(fid, "dynamic_data")
 
@@ -1228,15 +1436,55 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                               parallel_io=parallel_io,
                                               description="electron species density",
                                               units="n_ref")
+        if write_error_diagnostics
+            io_electron_density_loworder =
+                create_dynamic_variable!(dynamic, "electron_density_loworder", mk_float, z, r;
+                                         parallel_io=parallel_io,
+                                         description="low-order approximation to electron species density, used to diagnose timestepping error",
+                                         units="n_ref")
+        else
+            io_electron_density_loworder = nothing
+        end
+        if write_steady_state_diagnostics
+            io_electron_density_start_last_timestep =
+                create_dynamic_variable!(dynamic, "electron_density_start_last_timestep",
+                                         mk_float, z, r; parallel_io=parallel_io,
+                                         description="electron species density at the start of the last timestep before output, used to measure steady state residual",
+                                         units="n_ref")
+        else
+            io_electron_density_start_last_timestep = nothing
+        end
 
         # io_electron_upar is the handle for the electron parallel flow density
         io_electron_upar = create_dynamic_variable!(dynamic, "electron_parallel_flow", mk_float, z, r;
                                            parallel_io=parallel_io,
                                            description="electron species parallel flow",
                                            units="c_ref = sqrt(2*T_ref/mi)")
+        if write_error_diagnostics
+            io_electron_upar_loworder =
+                create_dynamic_variable!(dynamic, "electron_parallel_flow_loworder", mk_float, z,
+                                         r; parallel_io=parallel_io,
+                                         description="low-order approximation to electron species parallel flow, used to diagnose timestepping error",
+                                         units="c_ref = sqrt(2*T_ref/mi)")
+        else
+            io_electron_upar_loworder = nothing
+        end
+        if write_steady_state_diagnostics
+            io_electron_upar_start_last_timestep =
+                create_dynamic_variable!(dynamic, "electron_parallel_flow_start_last_timestep",
+                                         mk_float, z, r; parallel_io=parallel_io,
+                                         description="electron species parallel flow at the start of the last timestep before output, used to measure steady state residual",
+                                         units="c_ref = sqrt(2*T_ref/mi)")
+        else
+            io_electron_upar_start_last_timestep = nothing
+        end
     else
         io_electron_density = nothing
+        io_electron_density_loworder = nothing
+        io_electron_density_start_last_timestep = nothing
         io_electron_upar = nothing
+        io_electron_upar_loworder = nothing
+        io_electron_upar_start_last_timestep = nothing
     end
 
     # io_electron_ppar is the handle for the electron parallel pressure
@@ -1244,6 +1492,25 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                        parallel_io=parallel_io,
                                        description="electron species parallel pressure",
                                        units="n_ref*T_ref")
+    if write_error_diagnostics
+        io_electron_ppar_loworder =
+            create_dynamic_variable!(dynamic, "electron_parallel_pressure_loworder", mk_float,
+                                     z, r; parallel_io=parallel_io,
+                                     description="low-order approximation to electron species parallel pressure, used to diagnose timestepping error",
+                                     units="n_ref*T_ref")
+    else
+        io_electron_ppar_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_electron_ppar_start_last_timestep =
+            create_dynamic_variable!(dynamic,
+                                     "electron_parallel_pressure_start_last_timestep",
+                                     mk_float, z, r; parallel_io=parallel_io,
+                                     description="electron species parallel pressure at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref*T_ref")
+    else
+        io_electron_ppar_start_last_timestep = nothing
+    end
 
     # io_electron_qpar is the handle for the electron parallel heat flux
     io_electron_qpar = create_dynamic_variable!(dynamic, "electron_parallel_heat_flux", mk_float, z, r;
@@ -1336,9 +1603,12 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
         io_electron_dt_before_last_fail = nothing
     end
 
-    return io_electron_density, io_electron_upar, io_electron_ppar, io_electron_qpar,
-           io_electron_vth, external_source_electron_amplitude,
-           external_source_electron_density_amplitude,
+    return io_electron_density, io_electron_density_loworder,
+           io_electron_density_start_last_timestep, io_electron_upar,
+           io_electron_upar_loworder, io_electron_upar_start_last_timestep,
+           io_electron_ppar, io_electron_ppar_loworder,
+           io_electron_ppar_start_last_timestep, io_electron_qpar, io_electron_vth,
+           external_source_electron_amplitude, external_source_electron_density_amplitude,
            external_source_electron_momentum_amplitude,
            external_source_electron_pressure_amplitude,
            electron_constraints_A_coefficient, electron_constraints_B_coefficient,
@@ -1352,7 +1622,7 @@ define dynamic (time-evolving) neutral moment variables for writing to the hdf5
 """
 function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coordinate,
         z::coordinate, parallel_io, external_source_settings, evolve_density, evolve_upar,
-        evolve_ppar)
+        evolve_ppar, write_error_diagnostics, write_steady_state_diagnostics)
 
     dynamic = get_group(fid, "dynamic_data")
 
@@ -1362,6 +1632,26 @@ function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coo
                                                   parallel_io=parallel_io,
                                                   description="neutral species density",
                                                   units="n_ref")
+    if write_error_diagnostics
+        io_density_neutral_loworder =
+            create_dynamic_variable!(dynamic, "density_neutral_loworder", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="low-order approximation to neutral species density, used to diagnose timestepping error",
+                                     units="n_ref")
+    else
+        io_density_neutral_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_density_neutral_start_last_timestep =
+            create_dynamic_variable!(dynamic, "density_neutral_start_last_timestep", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="neutral species density at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref")
+    else
+        io_density_neutral_start_last_timestep = nothing
+    end
 
     # io_uz_neutral is the handle for the neutral z momentum density
     io_uz_neutral = create_dynamic_variable!(dynamic, "uz_neutral", mk_float, z, r;
@@ -1369,6 +1659,26 @@ function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coo
                                              parallel_io=parallel_io,
                                              description="neutral species mean z velocity",
                                              units="c_ref = sqrt(2*T_ref/mi)")
+    if write_error_diagnostics
+        io_uz_neutral_loworder =
+            create_dynamic_variable!(dynamic, "uz_neutral_loworder", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="low-order approximation to neutral species mean z velocity, used to diagnose timestepping error",
+                                     units="c_ref = sqrt(2*T_ref/mi)")
+    else
+        io_uz_neutral_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_uz_neutral_start_last_timestep =
+            create_dynamic_variable!(dynamic, "uz_neutral_start_last_timestep", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="neutral species mean z velocity at the start of the last timestep before output, used to measure steady state residual",
+                                     units="c_ref = sqrt(2*T_ref/mi)")
+    else
+        io_uz_neutral_start_last_timestep = nothing
+    end
 
     # io_pz_neutral is the handle for the neutral species zz pressure
     io_pz_neutral = create_dynamic_variable!(dynamic, "pz_neutral", mk_float, z, r;
@@ -1376,6 +1686,26 @@ function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coo
                                              parallel_io=parallel_io,
                                              description="neutral species mean zz pressure",
                                              units="n_ref*T_ref")
+    if write_error_diagnostics
+        io_pz_neutral_loworder =
+            create_dynamic_variable!(dynamic, "pz_neutral_loworder", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="low-order approximation to neutral species mean zz pressure, used to diagnose timestepping error",
+                                     units="n_ref*T_ref")
+    else
+        io_pz_neutral_loworder = nothing
+    end
+    if write_steady_state_diagnostics
+        io_pz_neutral_start_last_timestep =
+            create_dynamic_variable!(dynamic, "pz_neutral_start_last_timestep", mk_float, z, r;
+                                     n_neutral_species=n_neutral_species,
+                                     parallel_io=parallel_io,
+                                     description="neutral species mean zz pressure at the start of the last timestep before output, used to measure steady state residual",
+                                     units="n_ref*T_ref")
+    else
+        io_pz_neutral_start_last_timestep = nothing
+    end
 
     # io_qz_neutral is the handle for the neutral z heat flux
     io_qz_neutral = create_dynamic_variable!(dynamic, "qz_neutral", mk_float, z, r;
@@ -1467,9 +1797,11 @@ function define_dynamic_neutral_moment_variables!(fid, n_neutral_species, r::coo
            neutral_constraints_C_coefficient = nothing
     end
 
-    return io_density_neutral, io_uz_neutral, io_pz_neutral, io_qz_neutral,
-           io_thermal_speed_neutral, external_source_neutral_amplitude,
-           external_source_neutral_density_amplitude,
+    return io_density_neutral, io_density_neutral_loworder,
+           io_density_neutral_start_last_timestep, io_uz_neutral, io_uz_neutral_loworder,
+           io_uz_neutral_start_last_timestep, io_pz_neutral, io_pz_neutral_loworder,
+           io_pz_neutral_start_last_timestep, io_qz_neutral, io_thermal_speed_neutral,
+           external_source_neutral_amplitude, external_source_neutral_density_amplitude,
            external_source_neutral_momentum_amplitude,
            external_source_neutral_pressure_amplitude,
            external_source_neutral_controller_integral, neutral_constraints_A_coefficient,
@@ -1503,6 +1835,25 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
                                         n_ion_species=composition.n_ion_species,
                                         parallel_io=parallel_io,
                                         description="ion species distribution function")
+        if io_input.write_error_diagnostics
+            io_f_loworder = create_dynamic_variable!(dynamic, "f_loworder", mk_float, vpa,
+                                                     vperp, z, r;
+                                                     n_ion_species=composition.n_ion_species,
+                                                     parallel_io=parallel_io,
+                                                     description="low-order approximation to ion species distribution function, used to diagnose timestepping error")
+        else
+            io_f_loworder = nothing
+        end
+        if io_input.write_steady_state_diagnostics
+            io_f_start_last_timestep =
+                create_dynamic_variable!(dynamic, "f_start_last_timestep", mk_float, vpa,
+                                         vperp, z, r;
+                                         n_ion_species=composition.n_ion_species,
+                                         parallel_io=parallel_io,
+                                         description="ion species distribution function at the start of the last timestep before output, used to measure steady state residual")
+        else
+            io_f_start_last_timestep = nothing
+        end
 
         if composition.electron_physics == kinetic_electrons
             # io_f_electron is the handle for the electron pdf
@@ -1510,8 +1861,28 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
                                                      vperp, z, r;
                                                      parallel_io=parallel_io,
                                                      description="electron distribution function")
+            if io_input.write_error_diagnostics
+                io_f_electron_loworder =
+                    create_dynamic_variable!(dynamic, "f_electron_loworder", mk_float,
+                                             vpa, vperp, z, r;
+                                             parallel_io=parallel_io,
+                                             description="low-order approximation to electron distribution function, used to diagnose timestepping error")
+            else
+                io_f_electron_loworder = nothing
+            end
+            if io_input.write_steady_state_diagnostics
+                io_f_electron_start_last_timestep =
+                    create_dynamic_variable!(dynamic, "f_electron_start_last_timestep",
+                                             mk_float, vpa, vperp, z, r;
+                                             parallel_io=parallel_io,
+                                             description="electron distribution function at the start of the last electron pseudo-timestep before output, used to measure steady state residual")
+            else
+                io_f_electron_start_last_timestep = nothing
+            end
         else
             io_f_electron = nothing
+            io_f_electron_loworder = nothing
+            io_f_electron_start_last_timestep = nothing
         end
 
         # io_f_neutral is the handle for the neutral pdf
@@ -1519,8 +1890,32 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
                                                 n_neutral_species=composition.n_neutral_species,
                                                 parallel_io=parallel_io,
                                                 description="neutral species distribution function")
+        if io_input.write_error_diagnostics
+            io_f_neutral_loworder =
+                create_dynamic_variable!(dynamic, "f_neutral_loworder", mk_float, vz, vr,
+                                         vzeta, z, r;
+                                         n_ion_species=composition.n_ion_species,
+                                         parallel_io=parallel_io,
+                                         description="low-order approximation to neutral species distribution function, used to diagnose timestepping error")
+        else
+            io_f_neutral_loworder = nothing
+        end
+        if io_input.write_steady_state_diagnostics
+            io_f_neutral_start_last_timestep =
+                create_dynamic_variable!(dynamic, "f_neutral_start_last_timestep",
+                                         mk_float, vz, vr, vzeta, z, r;
+                                         n_ion_species=composition.n_ion_species,
+                                         parallel_io=parallel_io,
+                                         description="neutral species distribution function at the start of the last timestep before output, used to measure steady state residual")
+        else
+            io_f_neutral_start_last_timestep = nothing
+        end
 
-        return io_dfns_info(fid, io_f, io_f_electron, io_f_neutral, parallel_io, io_moments)
+        return io_dfns_info(fid, io_f, io_f_loworder, io_f_start_last_timestep,
+                            io_f_electron, io_f_electron_loworder,
+                            io_f_electron_start_last_timestep, io_f_neutral,
+                            io_f_neutral_loworder, io_f_neutral_start_last_timestep,
+                            parallel_io, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1644,19 +2039,40 @@ function reopen_moments_io(file_info)
             end
         end
         return io_moments_info(fid, getvar("time"), getvar("phi"), getvar("Er"),
-                               getvar("Ez"), getvar("density"), getvar("parallel_flow"),
-                               getvar("parallel_pressure"), getvar("perpendicular_pressure"),
+                               getvar("Ez"), getvar("density"),
+                               getvar("density_loworder"),
+                               getvar("density_start_last_timestep"),
+                               getvar("parallel_flow"), getvar("parallel_flow_loworder"),
+                               getvar("parallel_flow_start_last_timestep"),
+                               getvar("parallel_pressure"),
+                               getvar("parallel_pressure_loworder"),
+                               getvar("parallel_pressure_start_last_timestep"),
+                               getvar("perpendicular_pressure"),
+                               getvar("perpendicular_pressure_loworder"),
+                               getvar("perpendicular_pressure_start_last_timestep"),
                                getvar("parallel_heat_flux"),
                                getvar("thermal_speed"), getvar("entropy_production"),
                                getvar("chodura_integral_lower"),
                                getvar("chodura_integral_upper"),
                                getvar("electron_density"),
+                               getvar("electron_density_loworder"),
+                               getvar("electron_density_start_last_timestep"),
                                getvar("electron_parallel_flow"),
+                               getvar("electron_parallel_flow_loworder"),
+                               getvar("electron_parallel_flow_start_last_timestep"),
                                getvar("electron_parallel_pressure"),
+                               getvar("electron_parallel_pressure_loworder"),
+                               getvar("electron_parallel_pressure_start_last_timestep"),
                                getvar("electron_parallel_heat_flux"),
                                getvar("electron_thermal_speed"),
-                               getvar("density_neutral"), getvar("uz_neutral"),
-                               getvar("pz_neutral"), getvar("qz_neutral"),
+                               getvar("density_neutral"),
+                               getvar("density_neutral_loworder"),
+                               getvar("density_neutral_start_last_timestep"),
+                               getvar("uz_neutral"), getvar("uz_neutral_loworder"),
+                               getvar("uz_neutral_start_last_timestep"),
+                               getvar("pz_neutral"), getvar("pz_neutral_loworder"),
+                               getvar("pz_neutral_start_last_timestep"),
+                               getvar("qz_neutral"),
                                getvar("thermal_speed_neutral"),
                                getvar("external_source_amplitude"),
                                getvar("external_source_density_amplitude"),
@@ -1780,18 +2196,41 @@ function reopen_dfns_io(file_info)
         end
         io_moments = io_moments_info(fid, getvar("time"), getvar("phi"), getvar("Er"),
                                      getvar("Ez"), getvar("density"),
-                                     getvar("parallel_flow"), getvar("parallel_pressure"),
+                                     getvar("density_loworder"),
+                                     getvar("density_start_last_timestep"),
+                                     getvar("parallel_flow"),
+                                     getvar("parallel_flow_loworder"),
+                                     getvar("parallel_flow_start_last_timestep"),
+                                     getvar("parallel_pressure"),
+                                     getvar("parallel_pressure_loworder"),
+                                     getvar("parallel_pressure_start_last_timestep"),
                                      getvar("perpendicular_pressure"),
-                                     getvar("parallel_heat_flux"), getvar("thermal_speed"),
-                                     getvar("entropy_production"), getvar("chodura_integral_lower"),
+                                     getvar("perpendicular_pressure_loworder"),
+                                     getvar("perpendicular_pressure_start_last_timestep"),
+                                     getvar("parallel_heat_flux"),
+                                     getvar("thermal_speed"),
+                                     getvar("entropy_production"),
+                                     getvar("chodura_integral_lower"),
                                      getvar("chodura_integral_upper"),
                                      getvar("electron_density"),
+                                     getvar("electron_density_loworder"),
+                                     getvar("electron_density_start_last_timestep"),
                                      getvar("electron_parallel_flow"),
+                                     getvar("electron_parallel_flow_loworder"),
+                                     getvar("electron_parallel_flow_start_last_timestep"),
                                      getvar("electron_parallel_pressure"),
+                                     getvar("electron_parallel_pressure_loworder"),
+                                     getvar("electron_parallel_pressure_start_last_timestep"),
                                      getvar("electron_parallel_heat_flux"),
                                      getvar("electron_thermal_speed"),
-                                     getvar("density_neutral"), getvar("uz_neutral"),
-                                     getvar("pz_neutral"), getvar("qz_neutral"),
+                                     getvar("density_neutral"),
+                                     getvar("density_neutral_loworder"),
+                                     getvar("density_neutral_start_last_timestep"),
+                                     getvar("uz_neutral"), getvar("uz_neutral_loworder"),
+                                     getvar("uz_neutral_start_last_timestep"),
+                                     getvar("pz_neutral"),getvar("pz_neutral_loworder"),
+                                     getvar("pz_neutral_start_last_timestep"),
+                                     getvar("qz_neutral"),
                                      getvar("thermal_speed_neutral"),
                                      getvar("external_source_amplitude"),
                                      getvar("external_source_density_amplitude"),
@@ -1829,8 +2268,12 @@ function reopen_dfns_io(file_info)
                                      getvar("electron_dt_before_last_fail"),
                                      getvar("nl_solver_diagnostics"), io_input)
 
-        return io_dfns_info(fid, getvar("f"), getvar("f_electron"), getvar("f_neutral"),
-                            io_input, io_moments)
+        return io_dfns_info(fid, getvar("f"), getvar("f_loworder"),
+                            getvar("f_start_last_timestep"), getvar("f_electron"),
+                            getvar("f_electron_loworder"),
+                            getvar("f_electron_start_last_timestep"), getvar("f_neutral"),
+                            getvar("f_neutral_loworder"),
+                            getvar("f_neutral_start_last_timestep"), io_input, io_moments)
     end
 
     # For processes other than the root process of each shared-memory group...
@@ -1964,15 +2407,43 @@ function write_ion_moments_data_to_binary(scratch, moments, n_ion_species, t_par
         # add the density data at this time slice to the output file
         append_to_dynamic_var(io_moments.density, scratch[t_params.n_rk_stages+1].density,
                               t_idx, parallel_io, z, r, n_ion_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.density_loworder, scratch[2].density, t_idx,
+                              parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.density_start_last_timestep, scratch[1].density,
+                              t_idx, parallel_io, z, r, n_ion_species)
+
         append_to_dynamic_var(io_moments.parallel_flow,
                               scratch[t_params.n_rk_stages+1].upar, t_idx, parallel_io, z,
                               r, n_ion_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.parallel_flow_loworder, scratch[2].upar, t_idx,
+                              parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.parallel_flow_start_last_timestep,
+                              scratch[1].upar, t_idx, parallel_io, z, r, n_ion_species)
+
         append_to_dynamic_var(io_moments.parallel_pressure,
                               scratch[t_params.n_rk_stages+1].ppar, t_idx, parallel_io, z,
                               r, n_ion_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.parallel_pressure_loworder, scratch[2].ppar,
+                              t_idx, parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.parallel_pressure_start_last_timestep,
+                              scratch[1].ppar, t_idx, parallel_io, z, r, n_ion_species)
+
         append_to_dynamic_var(io_moments.perpendicular_pressure,
                               scratch[t_params.n_rk_stages+1].pperp, t_idx, parallel_io,
                               z, r, n_ion_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.perpendicular_pressure_loworder,
+                              scratch[2].pperp, t_idx, parallel_io, z, r, n_ion_species)
+        append_to_dynamic_var(io_moments.perpendicular_pressure_start_last_timestep,
+                              scratch[1].pperp, t_idx, parallel_io, z, r, n_ion_species)
+
         append_to_dynamic_var(io_moments.parallel_heat_flux, moments.ion.qpar, t_idx,
                               parallel_io, z, r, n_ion_species)
         append_to_dynamic_var(io_moments.thermal_speed, moments.ion.vth, t_idx,
@@ -2061,17 +2532,36 @@ function write_electron_moments_data_to_binary(scratch, moments, t_params, elect
             append_to_dynamic_var(io_moments.electron_density,
                                   scratch[t_params.n_rk_stages+1].electron_density, t_idx,
                                   parallel_io, z, r)
+            # If options were not set to select the following outputs, then the io variables
+            # will be `nothing` and nothing will be written.
+            append_to_dynamic_var(io_moments.electron_density_loworder,
+                                  scratch[2].electron_density, t_idx, parallel_io, z, r)
+            append_to_dynamic_var(io_moments.electron_density_start_last_timestep,
+                                  scratch[1].electron_density, t_idx, parallel_io, z, r)
         end
 
         if io_moments.electron_parallel_flow !== nothing
             append_to_dynamic_var(io_moments.electron_parallel_flow,
                                   scratch[t_params.n_rk_stages+1].electron_upar, t_idx,
                                   parallel_io, z, r)
+            # If options were not set to select the following outputs, then the io variables
+            # will be `nothing` and nothing will be written.
+            append_to_dynamic_var(io_moments.electron_parallel_flow_loworder,
+                                  scratch[2].electron_upar, t_idx, parallel_io, z, r)
+            append_to_dynamic_var(io_moments.electron_parallel_flow_start_last_timestep,
+                                  scratch[1].electron_upar, t_idx, parallel_io, z, r)
         end
 
         append_to_dynamic_var(io_moments.electron_parallel_pressure,
                               scratch[t_params.n_rk_stages+1].electron_ppar, t_idx,
                               parallel_io, z, r)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.electron_parallel_pressure_loworder,
+                              scratch[2].electron_ppar, t_idx, parallel_io, z, r)
+        append_to_dynamic_var(io_moments.electron_parallel_pressure_start_last_timestep,
+                              scratch[1].electron_ppar, t_idx, parallel_io, z, r)
+
         append_to_dynamic_var(io_moments.electron_parallel_heat_flux,
                               moments.electron.qpar, t_idx, parallel_io, z, r)
         append_to_dynamic_var(io_moments.electron_thermal_speed, moments.electron.vth,
@@ -2146,12 +2636,39 @@ function write_neutral_moments_data_to_binary(scratch, moments, n_neutral_specie
         append_to_dynamic_var(io_moments.density_neutral,
                               scratch[t_params.n_rk_stages+1].density_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.density_neutral_loworder,
+                              scratch[2].density_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+        append_to_dynamic_var(io_moments.density_neutral_start_last_timestep,
+                              scratch[1].density_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+
         append_to_dynamic_var(io_moments.uz_neutral,
                               scratch[t_params.n_rk_stages+1].uz_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.uz_neutral_loworder,
+                              scratch[2].uz_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+        append_to_dynamic_var(io_moments.uz_neutral_start_last_timestep,
+                              scratch[1].uz_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+
         append_to_dynamic_var(io_moments.pz_neutral,
                               scratch[t_params.n_rk_stages+1].pz_neutral, t_idx,
                               parallel_io, z, r, n_neutral_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_moments.pz_neutral_loworder,
+                              scratch[2].pz_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+        append_to_dynamic_var(io_moments.pz_neutral_start_last_timestep,
+                              scratch[1].pz_neutral, t_idx, parallel_io, z, r,
+                              n_neutral_species)
+
         append_to_dynamic_var(io_moments.qz_neutral, moments.neutral.qz, t_idx,
                               parallel_io, z, r, n_neutral_species)
         append_to_dynamic_var(io_moments.thermal_speed_neutral, moments.neutral.vth,
@@ -2258,6 +2775,12 @@ function write_ion_dfns_data_to_binary(scratch, t_params, n_ion_species,
 
         append_to_dynamic_var(io_dfns.f, scratch[t_params.n_rk_stages+1].pdf, t_idx,
                               parallel_io, vpa, vperp, z, r, n_ion_species)
+        # If options were not set to select the following outputs, then the io variables
+        # will be `nothing` and nothing will be written.
+        append_to_dynamic_var(io_dfns.f_loworder, scratch[2].pdf, t_idx,
+                              parallel_io, vpa, vperp, z, r, n_ion_species)
+        append_to_dynamic_var(io_dfns.f_start_last_timestep, scratch[1].pdf, t_idx,
+                              parallel_io, vpa, vperp, z, r, n_ion_species)
     end
     return nothing
 end
@@ -2276,8 +2799,22 @@ function write_electron_dfns_data_to_binary(scratch_electron, t_params,
         parallel_io = io_dfns.io_input.parallel_io
 
         if io_dfns.f_electron !== nothing
+            if t_params.electron === nothing
+                # t_params is the t_params for electron timestepping
+                n_rk_stages = t_params.n_rk_stages
+            else
+                n_rk_stages = t_params.electron.n_rk_stages
+            end
             append_to_dynamic_var(io_dfns.f_electron,
-                                  scratch_electron[t_params.electron.n_rk_stages+1].pdf_electron,
+                                  scratch_electron[n_rk_stages+1].pdf_electron,
+                                  t_idx, parallel_io, vpa, vperp, z, r)
+            # If options were not set to select the following outputs, then the io
+            # variables will be `nothing` and nothing will be written.
+            append_to_dynamic_var(io_dfns.f_electron_loworder,
+                                  scratch_electron[2].pdf_electron,
+                                  t_idx, parallel_io, vpa, vperp, z, r)
+            append_to_dynamic_var(io_dfns.f_electron_start_last_timestep,
+                                  scratch_electron[1].pdf_electron,
                                   t_idx, parallel_io, vpa, vperp, z, r)
         end
     end
@@ -2301,6 +2838,14 @@ function write_neutral_dfns_data_to_binary(scratch, t_params, n_neutral_species,
             append_to_dynamic_var(io_dfns.f_neutral,
                                   scratch[t_params.n_rk_stages+1].pdf_neutral, t_idx,
                                   parallel_io, vz, vr, vzeta, z, r, n_neutral_species)
+            # If options were not set to select the following outputs, then the io
+            # variables will be `nothing` and nothing will be written.
+            append_to_dynamic_var(io_dfns.f_neutral_loworder, scratch[2].pdf_neutral,
+                                  t_idx, parallel_io, vz, vr, vzeta, z, r,
+                                  n_neutral_species)
+            append_to_dynamic_var(io_dfns.f_neutral_start_last_timestep,
+                                  scratch[1].pdf_neutral, t_idx, parallel_io, vz, vr,
+                                  vzeta, z, r, n_neutral_species)
         end
     end
     return nothing
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index beda2ace9..3688fd718 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -225,6 +225,8 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         implicit_ion_advance=true,
         implicit_vpa_advection=false,
         write_after_fixed_step_count=false,
+        write_error_diagnostics=false,
+        write_steady_state_diagnostics=false,
         high_precision_error_sum=false,
        )
     if timestepping_section["nwrite"] > timestepping_section["nstep"]
@@ -259,6 +261,8 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         minimum_dt=timestepping_section["minimum_dt"] * sqrt(composition.me_over_mi),
         maximum_dt=timestepping_section["maximum_dt"] * sqrt(composition.me_over_mi),
         write_after_fixed_step_count=false,
+        write_error_diagnostics=false,
+        write_steady_state_diagnostics=false,
         high_precision_error_sum=timestepping_section["high_precision_error_sum"],
         initialization_residual_value=1.0,
         no_restart=false,
@@ -725,6 +729,10 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     io_settings["run_id"] = run_id
     io_settings["output_dir"] = output_dir
     io_settings["run_name"] = run_name
+    io_settings["write_error_diagnostics"] = timestepping_section["write_error_diagnostics"]
+    io_settings["write_steady_state_diagnostics"] = timestepping_section["write_steady_state_diagnostics"]
+    io_settings["write_electron_error_diagnostics"] = timestepping_section["electron_t_input"]["write_error_diagnostics"]
+    io_settings["write_electron_steady_state_diagnostics"] = timestepping_section["electron_t_input"]["write_steady_state_diagnostics"]
     io_immutable = Dict_to_NamedTuple(io_settings)
 
     # initialize z grid and write grid point locations to file
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index d4fffa612..a557c1fb4 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -308,6 +308,18 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         # No adaptive timestep, want to use the value from the input file even when we are
         # restarting
         dt_reload = nothing
+
+        # Makes no sense to use write_error_diagnostics because non-adaptive schemes have
+        # no error estimate
+        input_dict["write_error_diagnostics"] = false
+    end
+
+    if adaptive && t_input["write_error_diagnostics"] && !t_input["write_after_fixed_step_count"]
+        println("WARNING: using adaptive timestepping, so short, random-length timesteps "
+                * "before output is written will make diagnostics from "
+                * "`write_error_diagnostics=true` hard to interpret. If these "
+                * "diagnostics are important, suggest using "
+                * "`write_after_fixed_step_count=true`.")
     end
 
     dt_shared = allocate_shared_float(1)
diff --git a/moment_kinetics/test/recycling_fraction_tests.jl b/moment_kinetics/test/recycling_fraction_tests.jl
index 93c649390..01cc344af 100644
--- a/moment_kinetics/test/recycling_fraction_tests.jl
+++ b/moment_kinetics/test/recycling_fraction_tests.jl
@@ -110,7 +110,9 @@ test_input_split3 = merge(test_input_split2,
                                "ion_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0, "vpa_dissipation_coefficient" => 1e-2),
                                "neutral_numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0, "vz_dissipation_coefficient" => 1e-2)))
 test_input_split3["timestepping"] = merge(test_input_split3["timestepping"],
-                                           Dict("dt" => 1.0e-5))
+                                           Dict("dt" => 1.0e-5,
+                                                "write_error_diagnostics" => true,
+                                                "write_steady_state_diagnostics" => true))
 
 # default inputs for adaptive timestepping tests
 test_input_adaptive = merge(test_input,

From 2a60c245d2a7479949ba6cb9ccdb2dac353519cb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 21:44:35 +0100
Subject: [PATCH 309/428] Save size of the last timestep before each output

Will be useful for some adaptive timestep diagnostics.
---
 moment_kinetics/src/file_io.jl | 51 +++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 08e49e635..3a7c681ce 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -178,6 +178,8 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     step_counter::Tint
     # current timestep size
     dt::Ttime
+    # size of last timestep before output, used for some diagnostics
+    previous_dt::Ttime
     # cumulative number of timestep failures
     failure_counter::Tint
     # cumulative count of which variable caused a timstep failure
@@ -191,6 +193,8 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     electron_step_counter::Telectronint
     # current electron pseudo-timestep size
     electron_dt::Telectrontime
+    # size of last electron pseudo-timestep before the output was written
+    electron_previous_dt::Telectrontime
     # cumulative number of electron pseudo-timestep failures
     electron_failure_counter::Telectronint
     # cumulative count of which variable caused a electron pseudo-timstep failure
@@ -198,7 +202,7 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     # cumulative count of which factors limited the electron pseudo-timestep at each step
     electron_limit_caused_by::Telectronfailcause
     # Last successful timestep before most recent electron pseudo-timestep failure, used
-    # by adaptve timestepping algorithm
+    # by adaptive timestepping algorithm
     electron_dt_before_last_fail::Telectrontime
     # Variables recording diagnostic information about non-linear solvers (used for
     # implicit parts of timestep). These are stored in nested NamedTuples so that we can
@@ -294,6 +298,8 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3,
     electron_step_counter::Telectronint
     # current electron pseudo-timestep size
     electron_dt::Telectrontime
+    # size of last electron pseudo-timestep before the output was written
+    electron_previous_dt::Telectrontime
     # cumulative number of electron pseudo-timestep failures
     electron_failure_counter::Telectronint
     # cumulative count of which variable caused a electron pseudo-timstep failure
@@ -478,8 +484,9 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
         external_source_electron_pressure_amplitude,
         electron_constraints_A_coefficient, electron_constraints_B_coefficient,
         electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
-        io_electron_failure_counter, io_electron_failure_caused_by,
-        io_electron_limit_caused_by, io_electron_dt_before_last_fail =
+        io_electron_previous_dt, io_electron_failure_counter,
+        io_electron_failure_caused_by, io_electron_limit_caused_by,
+        io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -554,6 +561,7 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_constraints_C_coefficient"),
                                         getvar("electron_step_counter"),
                                         getvar("electron_dt"),
+                                        getvar("electron_previous_dt"),
                                         getvar("electron_failure_counter"),
                                         getvar("electron_failure_caused_by"),
                                         getvar("electron_limit_caused_by"),
@@ -1017,8 +1025,9 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         external_source_electron_pressure_amplitude,
         electron_constraints_A_coefficient, electron_constraints_B_coefficient,
         electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
-        io_electron_failure_counter, io_electron_failure_caused_by,
-        io_electron_limit_caused_by, io_electron_dt_before_last_fail =
+        io_electron_previous_dt, io_electron_failure_counter,
+        io_electron_failure_caused_by, io_electron_limit_caused_by,
+        io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -1057,6 +1066,10 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
             dynamic, "dt", mk_float; parallel_io=parallel_io,
             description="current timestep size")
 
+        io_previous_dt = create_dynamic_variable!(
+            dynamic, "previous_dt", mk_float; parallel_io=parallel_io,
+            description="size of the last timestep before the output")
+
         io_failure_counter = create_dynamic_variable!(
             dynamic, "failure_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of timestep failures for the run")
@@ -1137,12 +1150,12 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                electron_constraints_A_coefficient,
                                electron_constraints_B_coefficient,
                                electron_constraints_C_coefficient,
-                               io_time_for_run, io_step_counter, io_dt,
+                               io_time_for_run, io_step_counter, io_dt, io_previous_dt,
                                io_failure_counter, io_failure_caused_by,
                                io_limit_caused_by, io_dt_before_last_fail,
                                io_electron_step_counter, io_electron_dt,
-                               io_electron_failure_counter, io_electron_failure_caused_by,
-                               io_electron_limit_caused_by,
+                               io_electron_previous_dt, io_electron_failure_counter,
+                               io_electron_failure_caused_by, io_electron_limit_caused_by,
                                io_electron_dt_before_last_fail, io_nl_solver_diagnostics,
                                io_input)
     end
@@ -1571,6 +1584,10 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
             dynamic, "electron_dt", mk_float; parallel_io=parallel_io,
             description="current electron pseudo-timestep size")
 
+        io_electron_previous_dt = create_dynamic_variable!(
+            dynamic, "electron_previous_dt", mk_float; parallel_io=parallel_io,
+            description="size of last electron pseudo-timestep before output was written")
+
         io_electron_failure_counter = create_dynamic_variable!(
             dynamic, "electron_failure_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of electron pseudo-timestep failures for the run")
@@ -1597,6 +1614,7 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
     else
         io_electron_step_counter = nothing
         io_electron_dt = nothing
+        io_electron_previous_dt = nothing
         io_electron_failure_counter = nothing
         io_electron_failure_caused_by = nothing
         io_electron_limit_caused_by = nothing
@@ -1613,8 +1631,9 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
            external_source_electron_pressure_amplitude,
            electron_constraints_A_coefficient, electron_constraints_B_coefficient,
            electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
-           io_electron_failure_counter, io_electron_failure_caused_by,
-           io_electron_limit_caused_by, io_electron_dt_before_last_fail
+           io_electron_previous_dt, io_electron_failure_counter,
+           io_electron_failure_caused_by, io_electron_limit_caused_by,
+           io_electron_dt_before_last_fail
 end
 
 """
@@ -2098,10 +2117,11 @@ function reopen_moments_io(file_info)
                                getvar("electron_constraints_B_coefficient"),
                                getvar("electron_constraints_C_coefficient"),
                                getvar("time_for_run"), getvar("step_counter"),
-                               getvar("dt"), getvar("failure_counter"),
+                               getvar("dt"), getvar("previous_dt"), getvar("failure_counter"),
                                getvar("failure_caused_by"), getvar("limit_caused_by"),
                                getvar("dt_before_last_fail"),getvar("electron_step_counter"),
-                               getvar("electron_dt"), getvar("electron_failure_counter"),
+                               getvar("electron_dt"), getvar("electron_previous_dt"),
+                               getvar("electron_failure_counter"),
                                getvar("electron_failure_caused_by"),
                                getvar("electron_limit_caused_by"),
                                getvar("electron_dt_before_last_fail"),
@@ -2256,12 +2276,14 @@ function reopen_dfns_io(file_info)
                                      getvar("electron_constraints_B_coefficient"),
                                      getvar("electron_constraints_C_coefficient"),
                                      getvar("time_for_run"), getvar("step_counter"),
-                                     getvar("dt"), getvar("failure_counter"),
+                                     getvar("dt"), getvar("previous_dt"),
+                                     getvar("failure_counter"),
                                      getvar("failure_caused_by"),
                                      getvar("limit_caused_by"),
                                      getvar("dt_before_last_fail"),
                                      getvar("electron_step_counter"),
                                      getvar("electron_dt"),
+                                     getvar("electron_previous_dt"),
                                      getvar("electron_failure_counter"),
                                      getvar("electron_failure_caused_by"),
                                      getvar("electron_limit_caused_by"),
@@ -2345,6 +2367,7 @@ function write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_spe
         append_to_dynamic_var(io_moments.time_for_run, time_for_run, t_idx, parallel_io)
         append_to_dynamic_var(io_moments.step_counter, t_params.step_counter[], t_idx, parallel_io)
         append_to_dynamic_var(io_moments.dt, t_params.dt_before_output[], t_idx, parallel_io)
+        append_to_dynamic_var(io_moments.previous_dt, t_params.previous_dt[], t_idx, parallel_io)
         append_to_dynamic_var(io_moments.failure_counter, t_params.failure_counter[], t_idx, parallel_io)
         append_to_dynamic_var(io_moments.failure_caused_by, t_params.failure_caused_by,
                               t_idx, parallel_io, length(t_params.failure_caused_by);
@@ -2597,6 +2620,8 @@ function write_electron_moments_data_to_binary(scratch, moments, t_params, elect
             append_to_dynamic_var(io_moments.electron_dt,
                                   electron_t_params.dt_before_output[], t_idx,
                                   parallel_io)
+            append_to_dynamic_var(io_moments.electron_previous_dt,
+                                  electron_t_params.previous_dt[], t_idx, parallel_io)
             append_to_dynamic_var(io_moments.electron_failure_counter,
                                   electron_t_params.failure_counter[], t_idx, parallel_io)
             append_to_dynamic_var(io_moments.electron_failure_caused_by,

From 10d5397bac9796a2d0329a65102d0bfc5a846887 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 11 Jun 2024 21:45:24 +0100
Subject: [PATCH 310/428] Post-process diagnostics for timestep error and
 steady state residual

---
 .../src/makie_post_processing.jl              | 173 +++++++++++++++++-
 moment_kinetics/src/load_data.jl              |  69 ++++++-
 2 files changed, 237 insertions(+), 5 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 53ed2585d..d97d6e67e 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -266,10 +266,10 @@ function makie_post_process(run_dir::Union{String,Tuple},
         end
     end
 
-    timestep_diagnostics(run_info; plot_prefix=plot_prefix)
+    timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix)
     if any(ri.composition.electron_physics ==
            moment_kinetics.input_structs.kinetic_electrons for ri ∈ run_info)
-        timestep_diagnostics(run_info; plot_prefix=plot_prefix, electron=true)
+        timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix, electron=true)
     end
 
     do_steady_state_residuals = any(input_dict[v]["steady_state_residual"]
@@ -807,6 +807,12 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         this_input_dict, "timestep_diagnostics";
         plot=true,
         animate_CFL=false,
+        plot_timestep_residual=false,
+        animate_timestep_residual=false,
+        plot_timestep_error=false,
+        animate_timestep_error=false,
+        plot_steady_state_residual=false,
+        animate_steady_state_residual=false,
        )
 
     return nothing
@@ -7268,7 +7274,7 @@ function manufactured_solutions_analysis_dfns(run_info; plot_prefix)
 end
 
 """
-    timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing)
+    timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=nothing)
 
 Plot a time-trace of some adaptive-timestep diagnostics: steps per output, timestep
 failures per output, how many times per output each variable caused a timestep failure,
@@ -7279,7 +7285,8 @@ will be saved with the format `plot_prefix_timestep_diagnostics.pdf`.
 
 `it` can be used to select a subset of the time points by passing a range.
 """
-function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electron=false)
+function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=nothing,
+                              electron=false)
     try
         if !isa(run_info, Tuple)
             run_info = (run_info,)
@@ -7800,6 +7807,164 @@ function timestep_diagnostics(run_info; plot_prefix=nothing, it=nothing, electro
             end
         end
 
+        if run_info_dfns[1].dfns
+            this_input_dict = input_dict_dfns
+        else
+            this_input_dict = input_dict
+        end
+        if electron
+            variable_list = (v for v ∈ union((ri.evolving_variables for ri in run_info_dfns)...)
+                             if occursin("electron", v))
+        else
+            variable_list = (v for v ∈ union((ri.evolving_variables for ri in run_info_dfns)...)
+                             if !occursin("electron", v))
+        end
+        all_variable_names = union((ri.variable_names for ri ∈ run_info_dfns)...)
+
+        if input.plot_timestep_residual
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate residual for this variable
+                    continue
+                end
+                residual_name = variable_name * "_timestep_residual"
+                if variable_name == "f_neutral"
+                    plot_vs_vz_z(run_info_dfns, residual_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * residual_name * "_vs_vz_z.pdf")
+                elseif variable_name ∈ ("f", "f_electron")
+                    plot_vs_vpa_z(run_info_dfns, residual_name;
+                                  input=this_input_dict[variable_name],
+                                  outfile=plot_prefix * residual_name * "_vs_vpa_z.pdf")
+                else
+                    plot_vs_z(run_info_dfns, residual_name;
+                              input=this_input_dict[variable_name],
+                              outfile=plot_prefix * residual_name * "_vs_z.pdf")
+                end
+            end
+        end
+
+        if input.animate_timestep_residual
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate residual for this variable
+                    continue
+                end
+                residual_name = variable_name * "_timestep_residual"
+                if variable_name == "f_neutral"
+                    animate_vs_vz_z(run_info_dfns, residual_name;
+                                    input=this_input_dict[variable_name],
+                                    outfile=plot_prefix * residual_name * "_vs_vz_z." * this_input_dict[variable_name]["animation_ext"])
+                elseif variable_name ∈ ("f", "f_electron")
+                    animate_vs_vpa_z(run_info_dfns, residual_name;
+                                     input=this_input_dict[variable_name],
+                                     outfile=plot_prefix * residual_name * "_vs_vpa_z." * this_input_dict[variable_name]["animation_ext"])
+                else
+                    animate_vs_z(run_info_dfns, residual_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * residual_name * "_vs_z." * this_input_dict[variable_name]["animation_ext"])
+                end
+            end
+        end
+
+        if input.plot_timestep_error
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate error for this variable
+                    continue
+                end
+                error_name = variable_name * "_timestep_error"
+                if variable_name == "f_neutral"
+                    plot_vs_vz_z(run_info_dfns, error_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * error_name * "_vs_vz_z.pdf")
+                elseif variable_name ∈ ("f", "f_electron")
+                    plot_vs_vpa_z(run_info_dfns, error_name;
+                                  input=this_input_dict[variable_name],
+                                  outfile=plot_prefix * error_name * "_vs_vpa_z.pdf")
+                else
+                    plot_vs_z(run_info_dfns, error_name;
+                              input=this_input_dict[variable_name],
+                              outfile=plot_prefix * error_name * "_vs_z.pdf")
+                end
+            end
+        end
+
+        if input.animate_timestep_error
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate error for this variable
+                    continue
+                end
+                error_name = variable_name * "_timestep_error"
+                if variable_name == "f_neutral"
+                    animate_vs_vz_z(run_info_dfns, error_name;
+                                    input=this_input_dict[variable_name],
+                                    outfile=plot_prefix * error_name * "_vs_vz_z." * this_input_dict[variable_name]["animation_ext"])
+                elseif variable_name ∈ ("f", "f_electron")
+                    animate_vs_vpa_z(run_info_dfns, error_name;
+                                     input=this_input_dict[variable_name],
+                                     outfile=plot_prefix * error_name * "_vs_vpa_z." * this_input_dict[variable_name]["animation_ext"])
+                else
+                    animate_vs_z(run_info_dfns, error_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * error_name * "_vs_z." * this_input_dict[variable_name]["animation_ext"])
+                end
+            end
+        end
+
+        if input.plot_steady_state_residual
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate residual for this variable
+                    continue
+                end
+                residual_name = variable_name * "_steady_state_residual"
+                if variable_name == "f_neutral"
+                    plot_vs_vz_z(run_info_dfns, residual_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * residual_name * "_vs_vz_z.pdf")
+                elseif variable_name ∈ ("f", "f_electron")
+                    plot_vs_vpa_z(run_info_dfns, residual_name;
+                                  input=this_input_dict[variable_name],
+                                  outfile=plot_prefix * residual_name * "_vs_vpa_z.pdf")
+                else
+                    plot_vs_z(run_info_dfns, residual_name;
+                              input=this_input_dict[variable_name],
+                              outfile=plot_prefix * residual_name * "_vs_z.pdf")
+                end
+            end
+        end
+
+        if input.animate_steady_state_residual
+            for variable_name ∈ variable_list
+                loworder_name = variable_name * "_loworder"
+                if loworder_name ∉ all_variable_names
+                    # No data to calculate residual for this variable
+                    continue
+                end
+                residual_name = variable_name * "_steady_state_residual"
+                if variable_name == "f_neutral"
+                    animate_vs_vz_z(run_info_dfns, residual_name;
+                                    input=this_input_dict[variable_name],
+                                    outfile=plot_prefix * residual_name * "_vs_vz_z." * this_input_dict[variable_name]["animation_ext"])
+                elseif variable_name ∈ ("f", "f_electron")
+                    animate_vs_vpa_z(run_info_dfns, residual_name;
+                                     input=this_input_dict[variable_name],
+                                     outfile=plot_prefix * residual_name * "_vs_vpa_z." * this_input_dict[variable_name]["animation_ext"])
+                else
+                    animate_vs_z(run_info_dfns, residual_name;
+                                 input=this_input_dict[variable_name],
+                                 outfile=plot_prefix * residual_name * "_vs_z." * this_input_dict[variable_name]["animation_ext"])
+                end
+            end
+        end
+
         return steps_fig, dt_fig, CFL_fig
     catch e
         println("Error in timestep_diagnostics(). Error was ", e)
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index e96ee7c41..4851e6593 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3488,6 +3488,39 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
 
     # Get variable names just from the first restart, for simplicity
     variable_names = get_variable_keys(get_group(fids0[1], "dynamic_data"))
+    if initial_electron
+        evolving_variables = ("f_electron",)
+    else
+        evolving_variables = ["f"]
+        if evolve_density
+            push!(evolving_variables, "density")
+        end
+        if evolve_upar
+            push!(evolving_variables, "parallel_flow")
+        end
+        if evolve_ppar
+            push!(evolving_variables, "parallel_pressure")
+        end
+        if composition.electron_physics == kinetic_electrons
+            push!(evolving_variables, "f_electron")
+        end
+        if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+            push!(evolving_variables, "electron_parallel_pressure")
+        end
+        if composition.n_neutral_species > 0
+            push!(evolving_variables, "f_neutral")
+            if evolve_density
+                push!(evolving_variables, "density_neutral")
+            end
+            if evolve_upar
+                push!(evolving_variables, "uz_neutral")
+            end
+            if evolve_ppar
+                push!(evolving_variables, "pz_neutral")
+            end
+        end
+        evolving_variables = Tuple(evolving_variables)
+    end
 
     if parallel_io
         files = fids0
@@ -3517,7 +3550,8 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
                 z_chunk_size=z_chunk_size, vperp_chunk_size=vperp_chunk_size,
                 vpa_chunk_size=vpa_chunk_size, vzeta_chunk_size=vzeta_chunk_size,
                 vr_chunk_size=vr_chunk_size, vz_chunk_size=vz_chunk_size,
-                variable_names=variable_names, dfns=dfns)
+                variable_names=variable_names, evolving_variables=evolving_variables,
+                dfns=dfns)
 
     return run_info
 end
@@ -4745,6 +4779,39 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             variable[it] = min_CFL
         end
         variable = select_slice_of_variable(variable; kwargs...)
+    elseif occursin("_timestep_error", variable_name)
+        prefix = split(variable_name, "_timestep_error")[1]
+        full_order = get_variable(run_info, prefix; kwargs...)
+        low_order = get_variable(run_info, prefix * "_loworder"; kwargs...)
+        variable = low_order .- full_order
+    elseif occursin("_timestep_residual", variable_name)
+        prefix = split(variable_name, "_timestep_residual")[1]
+        full_order = get_variable(run_info, prefix; kwargs...)
+        low_order = get_variable(run_info, prefix * "_loworder"; kwargs...)
+        if prefix == "pdf_electron"
+            rtol = run_info.input["electron_timestepping"]["rtol"]
+            atol = run_info.input["electron_timestepping"]["atol"]
+        else
+            rtol = run_info.input["timestepping"]["rtol"]
+            atol = run_info.input["timestepping"]["atol"]
+        end
+        variable = @. (low_order - full_order) / (rtol * abs(full_order) + atol)
+    elseif occursin("_steady_state_residual", variable_name)
+        prefix = split(variable_name, "_steady_state_residual")[1]
+        end_step = get_variable(run_info, prefix; kwargs...)
+        begin_step = get_variable(run_info, prefix * "_start_last_timestep"; kwargs...)
+        if prefix == "f_electron"
+            dt = get_variable(run_info, "electron_previous_dt"; kwargs...)
+        else
+            dt = get_variable(run_info, "previous_dt"; kwargs...)
+        end
+        dt = reshape(dt, ones(mk_int, ndims(end_step)-1)..., length(dt))
+        for i ∈ eachindex(dt)
+            if dt[i] ≤ 0.0
+                dt[i] = Inf
+            end
+        end
+        variable = (end_step .- begin_step) ./ dt
     elseif occursin("_nonlinear_iterations_per_solve", variable_name)
         prefix = split(variable_name, "_nonlinear_iterations_per_solve")[1]
         nl_nsolves = get_per_step_from_cumulative_variable(

From 6bbe188cff52b10fb0830957d8a67214eee54578 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 12 Jun 2024 12:10:55 +0100
Subject: [PATCH 311/428] Allow z-periodic boundary conditions for kinetic
 electrons

---
 moment_kinetics/src/electron_kinetic_equation.jl | 13 ++++++++++---
 moment_kinetics/src/initial_conditions.jl        | 11 ++++++++++-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index c145a902c..dd11be109 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -620,8 +620,10 @@ function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp
     A = moments.electron.constraints_A_coefficient
     B = moments.electron.constraints_B_coefficient
     C = moments.electron.constraints_C_coefficient
+    skip_first = z.irank == 0 && z.bc != "periodic"
+    skip_last = z.irank == z.nrank - 1 && z.bc != "periodic"
     @loop_r_z ir iz begin
-        if (iz == 1 && z.irank == 0) || (iz == z.n && z.irank == z.nrank - 1)
+        if (iz == 1 && skip_first) || (iz == z.n && skip_last)
             continue
         end
         (A[iz,ir], B[iz,ir], C[iz,ir]) =
@@ -652,6 +654,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
         @views enforce_vperp_boundary_condition!(pdf, vperp.bc, vperp, vperp_spectral)
     end
 
+    if z.bc == "periodic"
+        # Nothing more to do for z-periodic boundary conditions
+        return nothing
+    end
+
     # first enforce the boundary condition at z_min.
     # this involves forcing the pdf to be zero for electrons travelling faster than the max speed
     # they could attain by accelerating in the electric field between the wall and the simulation boundary;
@@ -693,7 +700,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
     if z.irank == 0
         if z.bc != "wall"
-            error("Options other than wall bc not implemented yet for electrons")
+            error("Options other than wall or z-periodic bc not implemented yet for electrons")
         end
         @loop_r ir begin
             # Impose sheath-edge boundary condition, while also imposing moment
@@ -870,7 +877,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
     
     if z.irank == z.nrank - 1
         if z.bc != "wall"
-            error("Options other than wall bc not implemented yet for electrons")
+            error("Options other than wall or z-periodic bc not implemented yet for electrons")
         end
         @loop_r ir begin
             # Impose sheath-edge boundary condition, while also imposing moment
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6111468a8..668970dc0 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -1472,7 +1472,16 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
             end
         end
     else
-        println("!!! currently, only the wall BC is supported for kinetic electrons !!!")
+        begin_r_z_region()
+        @loop_r ir begin
+            # Initialise an unshifted Maxwellian as a first step
+            @loop_z iz begin
+                vpa_over_vth = @. vpa.scratch3 = vpa.grid + upar[iz,ir] / vth[iz,ir]
+                @loop_vperp ivperp begin
+                    @. pdf[:,ivperp,iz,ir] = exp(-vpa_over_vth^2)
+                end
+            end
+        end
     end
 end
 

From 7a4452b4657e8c139e2f10efd5360ecfbfc4d09c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 12 Jun 2024 12:12:09 +0100
Subject: [PATCH 312/428] Example kinetic electron simulation in 1D periodic
 box

---
 .../periodic_split3_kinetic.toml              | 131 +++++++++++++++++
 ...ic_split3_kinetic_high-collisionality.toml | 134 ++++++++++++++++++
 2 files changed, 265 insertions(+)
 create mode 100644 examples/kinetic-electrons/periodic_split3_kinetic.toml
 create mode 100644 examples/kinetic-electrons/periodic_split3_kinetic_high-collisionality.toml

diff --git a/examples/kinetic-electrons/periodic_split3_kinetic.toml b/examples/kinetic-electrons/periodic_split3_kinetic.toml
new file mode 100644
index 000000000..c2c734a78
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_kinetic.toml
@@ -0,0 +1,131 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "kinetic_electrons"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+#type = "KennedyCarpenterARK324"
+type = "Fekete4(3)"
+implicit_ion_advance = false
+implicit_vpa_advection = false
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-7
+rtol = 1.0e-7
+max_increase_factor_near_last_fail = 1.001
+last_fail_proximity_factor = 1.1
+max_increase_factor = 1.05
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+#[timestepping]
+##type = "KennedyCarpenterARK324"
+#type = "Fekete4(3)"
+#implicit_ion_advance = false
+#implicit_vpa_advection = false
+#nstep = 1000000
+#dt = 1.0e-6
+#minimum_dt = 1.0e-7
+#rtol = 1.0e-7
+#max_increase_factor_near_last_fail = 1.001
+#last_fail_proximity_factor = 1.1
+#max_increase_factor = 1.05
+#nwrite = 100
+#nwrite_dfns = 100
+#steady_state_residual = true
+#converged_residual_value = 1.0e-3
+#write_after_fixed_step_count = true
+#write_error_diagnostics = true
+#write_steady_state_diagnostics = true
+
+[electron_timestepping]
+nstep = 5000000
+#nstep = 1
+dt = 2.0e-8
+maximum_dt = 1.0
+nwrite = 10000
+nwrite_dfns = 100000
+#type = "SSPRK4"
+type = "Fekete4(3)"
+rtol = 1.0e-6
+atol = 1.0e-14
+minimum_dt = 1.0e-10
+initialization_residual_value = 2.5
+converged_residual_value = 0.1 #1.0e-3
+#debug_io = 10000
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient = 2.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true
diff --git a/examples/kinetic-electrons/periodic_split3_kinetic_high-collisionality.toml b/examples/kinetic-electrons/periodic_split3_kinetic_high-collisionality.toml
new file mode 100644
index 000000000..133d76b16
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_kinetic_high-collisionality.toml
@@ -0,0 +1,134 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "kinetic_electrons"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[reference_params]
+Tref = 20.0
+
+[timestepping]
+#type = "KennedyCarpenterARK324"
+type = "Fekete4(3)"
+implicit_ion_advance = false
+implicit_vpa_advection = false
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-7
+rtol = 1.0e-7
+max_increase_factor_near_last_fail = 1.001
+last_fail_proximity_factor = 1.1
+max_increase_factor = 1.05
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+#[timestepping]
+##type = "KennedyCarpenterARK324"
+#type = "Fekete4(3)"
+#implicit_ion_advance = false
+#implicit_vpa_advection = false
+#nstep = 1000000
+#dt = 1.0e-6
+#minimum_dt = 1.0e-7
+#rtol = 1.0e-7
+#max_increase_factor_near_last_fail = 1.001
+#last_fail_proximity_factor = 1.1
+#max_increase_factor = 1.05
+#nwrite = 100
+#nwrite_dfns = 100
+#steady_state_residual = true
+#converged_residual_value = 1.0e-3
+#write_after_fixed_step_count = true
+#write_error_diagnostics = true
+#write_steady_state_diagnostics = true
+
+[electron_timestepping]
+nstep = 5000000
+#nstep = 1
+dt = 2.0e-8
+maximum_dt = 1.0
+nwrite = 10000
+nwrite_dfns = 100000
+#type = "SSPRK4"
+type = "Fekete4(3)"
+rtol = 1.0e-6
+atol = 1.0e-14
+minimum_dt = 1.0e-10
+initialization_residual_value = 2.5
+converged_residual_value = 0.1 #1.0e-3
+#debug_io = 10000
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient = 2.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true

From 6802613f4dce6f96d016e867352ebf1ceeafa84d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 19 Jun 2024 11:27:56 +0100
Subject: [PATCH 313/428] Example periodic inputs for Boltzmann and non-IMEX
 Braginskii

---
 .../periodic_split3_boltzmann.toml            |  91 +++++++++++++++
 .../periodic_split3_braginskii.toml           | 104 ++++++++++++++++++
 2 files changed, 195 insertions(+)
 create mode 100644 examples/kinetic-electrons/periodic_split3_boltzmann.toml
 create mode 100644 examples/kinetic-electrons/periodic_split3_braginskii.toml

diff --git a/examples/kinetic-electrons/periodic_split3_boltzmann.toml b/examples/kinetic-electrons/periodic_split3_boltzmann.toml
new file mode 100644
index 000000000..27a3eff33
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_boltzmann.toml
@@ -0,0 +1,91 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "boltzmann_electron_response"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "Fekete4(3)"
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-7
+rtol = 1.0e-7
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+
+[nonlinear_solver]
+nonlinear_max_iterations = 100
+#rtol = 1.0e-9
+#atol = 1.0e-12
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true
diff --git a/examples/kinetic-electrons/periodic_split3_braginskii.toml b/examples/kinetic-electrons/periodic_split3_braginskii.toml
new file mode 100644
index 000000000..31741e462
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_braginskii.toml
@@ -0,0 +1,104 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "braginskii_fluid"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "Fekete4(3)"
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-9
+rtol = 1.0e-7
+nwrite = 10000
+nwrite_dfns = 100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+#[timestepping]
+#type = "Fekete4(3)"
+#nstep = 1000000
+#dt = 1.0e-6
+#minimum_dt = 1.0e-9
+#rtol = 1.0e-7
+#nwrite = 1000
+#nwrite_dfns = 1000
+#steady_state_residual = true
+#converged_residual_value = 1.0e-3
+#write_after_fixed_step_count = true
+#write_error_diagnostics = true
+#write_steady_state_diagnostics = true
+
+[nonlinear_solver]
+nonlinear_max_iterations = 100
+#rtol = 1.0e-9
+#atol = 1.0e-12
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true

From d2e0d24c9549fd142bb13f7192c73274e7215761 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 13 Jun 2024 13:09:56 +0100
Subject: [PATCH 314/428] Delay overwriting of scratch_electron[1]

We want to have the option to save the state in `scratch_electron[1]`,
in order to calculate steady-state residual diagnostics in
post-processing.  To allow this, we need to start the electron
pseudo-time-step with the state in
`scratch_electron[t_params.electron.n_rk_stages+1]` rather than the
state in `scratch_electron[1]`, and copy that state into
`scratch_electron[1]` at the beginning of each pseudo-time-step.
---
 .../src/electron_kinetic_equation.jl          | 43 +++++++++----------
 moment_kinetics/src/initial_conditions.jl     |  6 ++-
 2 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index dd11be109..841b0fad5 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -169,7 +169,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             0.5 * moments.electron.vth[iz,ir] *
             (moments.electron.dppar_dz[iz,ir] / moments.electron.ppar[iz,ir] -
              moments.electron.ddens_dz[iz,ir] / moments.electron.dens[iz,ir])
-        scratch[1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
+        scratch[t_params.n_rk_stages+1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
     end
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
@@ -223,7 +223,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 # need to exit or handle this appropriately
                 @loop_vpa ivpa begin
                     @loop_z iz begin
-                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", scratch[1].pdf_electron[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", scratch[t_params.n_rk_stages+1].pdf_electron[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
                     end
                     println(io_pdf,"")
                 end
@@ -246,10 +246,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     t_params.moments_output_counter[] += 1
     @serial_region begin
         if io_electron !== nothing
-            # Copy scratch arrays in order to save initial state - not optimal, but should
-            # be done at most once or twice while initialising a simulation
-            scratch[t_params.n_rk_stages+1].pdf_electron .= scratch[1].pdf_electron
-            scratch[t_params.n_rk_stages+1].electron_ppar .= scratch[1].electron_ppar
             write_electron_state(scratch, moments, t_params, time, io_electron,
                                  t_params.moments_output_counter[], r, z, vperp, vpa)
         end
@@ -258,6 +254,24 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     while (!electron_pdf_converged
            && (t_params.step_counter[] - initial_step_counter < max_electron_pdf_iterations)
            && t_params.dt[] > 0.0 && !isnan(t_params.dt[]))
+
+        # Set the initial values for the next step to the final values from the previous
+        # step
+        begin_r_z_vperp_vpa_region()
+        new_pdf = scratch[1].pdf_electron
+        old_pdf = scratch[t_params.n_rk_stages+1].pdf_electron
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            new_pdf[ivpa,ivperp,iz,ir] = old_pdf[ivpa,ivperp,iz,ir]
+        end
+        if evolve_ppar
+            begin_r_z_region()
+            new_ppar = scratch[1].electron_ppar
+            old_ppar = scratch[t_params.n_rk_stages+1].electron_ppar
+            @loop_r_z ir iz begin
+                new_ppar[iz,ir] = old_ppar[iz,ir]
+            end
+        end
+
         for istage ∈ 1:t_params.n_rk_stages
             # Set the initial values for this stage to the final values from the previous
             # stage
@@ -430,23 +444,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             electron_pdf_converged = MPI.Bcast(electron_pdf_converged, 0, comm_world)
         end
 
-        # Set the initial values for the next step to the final values from the previous
-        # step
-        begin_r_z_vperp_vpa_region()
-        new_pdf = scratch[1].pdf_electron
-        old_pdf = scratch[t_params.n_rk_stages+1].pdf_electron
-        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
-            new_pdf[ivpa,ivperp,iz,ir] = old_pdf[ivpa,ivperp,iz,ir]
-        end
-        if evolve_ppar
-            begin_r_z_region()
-            new_ppar = scratch[1].electron_ppar
-            old_ppar = scratch[t_params.n_rk_stages+1].electron_ppar
-            @loop_r_z ir iz begin
-                new_ppar[iz,ir] = old_ppar[iz,ir]
-            end
-        end
-
         if text_output
             if (mod(t_params.step_counter[] - initial_step_counter, t_params.nwrite_moments)==1)
                 begin_serial_region()
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 668970dc0..6f62d38f6 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -605,8 +605,10 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
 
         begin_serial_region()
         @serial_region begin
-            # update the electron pdf in the first scratch_electron
-            scratch_electron[1].pdf_electron .= pdf.electron.norm
+            # update the electron pdf in the last scratch_electron (which will be copied
+            # to the first entry as part of the pseudo-time-loop in
+            # update_electron_pdf!()).
+            scratch_electron[t_params.n_rk_stages+1].pdf_electron .= pdf.electron.norm
         end
 
         begin_r_z_region()

From 9c7d28d6b84e40bbcebfc7950159a284819fee2a Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 15 Jun 2024 19:01:40 +0100
Subject: [PATCH 315/428] Treat initial_electron restarts separately in
 get_run_info_no_setup()

There can be a different number of `*.initial_electron.h5` files than
there are `*.moments.h5` and `*.dfns.h5` files, because the electron
initialisation and the actual run (ion/neutral timestepping) restart
separately. This needs to be taken into account when setting up
`run_prefixes`.
---
 moment_kinetics/src/load_data.jl | 40 +++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 4851e6593..15d8f18ed 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3376,18 +3376,36 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
         # Find output files from all restarts in the directory
         counter = 1
         run_prefixes = Vector{String}()
-        while true
-            # Test if output files exist for this value of counter
-            prefix_with_count = base_prefix * "_$counter"
-            if length(glob(basename(prefix_with_count) * ".*.h5", dirname(prefix_with_count))) > 0 ||
-                length(glob(basename(prefix_with_count) * ".*.cdf", dirname(prefix_with_count))) > 0
-
-                push!(run_prefixes, prefix_with_count)
-            else
-                # No more output files found
-                break
+        if initial_electron
+            while true
+                # Test if output files exist for this value of counter
+                prefix_with_count = base_prefix * "_$counter"
+                if length(glob(basename(prefix_with_count) * ".initial_elctron*.h5", dirname(prefix_with_count))) > 0 ||
+                    length(glob(basename(prefix_with_count) * ".initial_elctron*.cdf", dirname(prefix_with_count))) > 0
+
+                    push!(run_prefixes, prefix_with_count)
+                else
+                    # No more output files found
+                    break
+                end
+                counter += 1
+            end
+        else
+            while true
+                # Test if output files exist for this value of counter
+                prefix_with_count = base_prefix * "_$counter"
+                if length(glob(basename(prefix_with_count) * ".dfns*.h5", dirname(prefix_with_count))) > 0 ||
+                    length(glob(basename(prefix_with_count) * ".dfns*.cdf", dirname(prefix_with_count))) > 0 ||
+                    length(glob(basename(prefix_with_count) * ".moments*.h5", dirname(prefix_with_count))) > 0 ||
+                    length(glob(basename(prefix_with_count) * ".moments*.cdf", dirname(prefix_with_count))) > 0
+
+                    push!(run_prefixes, prefix_with_count)
+                else
+                    # No more output files found
+                    break
+                end
+                counter += 1
             end
-            counter += 1
         end
         # Add the final run which does not have a '_$counter' suffix
         push!(run_prefixes, base_prefix)

From 9f828fdcb069e501613c95c6997897ebcee55b4e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 13:12:10 +0100
Subject: [PATCH 316/428] Functions to get or modify an attribute's value in an
 HDF5 or NetCDF file

---
 moment_kinetics/ext/file_io_netcdf.jl | 15 ++++++++++++---
 moment_kinetics/src/file_io.jl        |  5 +++++
 moment_kinetics/src/file_io_hdf5.jl   |  5 +++++
 moment_kinetics/src/load_data.jl      |  5 +++++
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/ext/file_io_netcdf.jl b/moment_kinetics/ext/file_io_netcdf.jl
index 9ca463674..bb7a86f2b 100644
--- a/moment_kinetics/ext/file_io_netcdf.jl
+++ b/moment_kinetics/ext/file_io_netcdf.jl
@@ -8,9 +8,10 @@ module file_io_netcdf
 
 import moment_kinetics.file_io: io_has_parallel, open_output_file_implementation,
                                 create_io_group, get_group, is_group, get_subgroup_keys,
-                                get_variable_keys, add_attribute!, write_single_value!,
-                                create_dynamic_variable!, append_to_dynamic_var
-import moment_kinetics.load_data: open_file_to_read, load_variable, load_slice
+                                get_variable_keys, add_attribute!, modify_attribute!,
+                                write_single_value!, create_dynamic_variable!,
+                                append_to_dynamic_var
+import moment_kinetics.load_data: open_file_to_read, get_attribute, load_variable, load_slice
 using moment_kinetics.coordinates: coordinate
 using moment_kinetics.input_structs: netcdf
 
@@ -76,6 +77,14 @@ function add_attribute!(var::NCDatasets.CFVariable, name, value)
     var.attrib[name] = value
 end
 
+function modify_attribute!(file_or_group_or_var::Union{NCDataset,NCDatasets.CFVariable}, name, value)
+    file_or_group_or_var.attrib[name] = value
+end
+
+function get_attribute(file_or_group_or_var::Union{NCDataset,NCDatasets.CFVariable}, name)
+    return var.attrib[name]
+end
+
 function maybe_create_netcdf_dim(file_or_group::NCDataset, name, size)
     if !(name ∈ keys(file_or_group.dim))
         defDim(file_or_group, name, size)
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 3a7c681ce..a30f70a0d 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1946,6 +1946,11 @@ Add an attribute to a file, group or variable
 """
 function add_attribute! end
 
+"""
+Modify an attribute to a file, group or variable
+"""
+function modify_attribute! end
+
 """
 Low-level function to open a binary output file
 
diff --git a/moment_kinetics/src/file_io_hdf5.jl b/moment_kinetics/src/file_io_hdf5.jl
index e39066f4c..271d79fee 100644
--- a/moment_kinetics/src/file_io_hdf5.jl
+++ b/moment_kinetics/src/file_io_hdf5.jl
@@ -54,6 +54,11 @@ function add_attribute!(var::HDF5.Dataset, name, value)
     attributes(var)[name] = value
 end
 
+# HDF5.H5DataStore is the supertype for HDF5.File and HDF5.Group
+function modify_attribute!(file_or_group_or_var::Union{HDF5.H5DataStore,HDF5.Dataset}, name, value)
+    attrs(file_or_group_or_var)[name] = value
+end
+
 function get_group(file_or_group::HDF5.H5DataStore, name::String)
     # This overload deals with cases where fid is an HDF5 `File` or `Group` (`H5DataStore`
     # is the abstract super-type for both
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 15d8f18ed..785c759de 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -82,6 +82,11 @@ function open_file_to_read(::Val{hdf5}, filename)
     return h5open(filename, "r")
 end
 
+function get_attribute end
+function get_attribute(file_or_group_or_var::Union{HDF5.H5DataStore,HDF5.Dataset}, name)
+    return attrs(file_or_group_or_var)[name]
+end
+
 """
 """
 function open_readonly_output_file(run_name, ext; iblock=0, printout=false)

From 58df6efb098175a3df3e31e3f46d55a850aec675 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 13:16:17 +0100
Subject: [PATCH 317/428] Record whether pdf_electron converged in
 *.initial_electron.h5

If the saved initial pdf did converge, we do not need to restart the
initialisation, we can just use the loaded values directly.
---
 moment_kinetics/src/file_io.jl            |  11 +-
 moment_kinetics/src/initial_conditions.jl | 146 +++++++++++++---------
 moment_kinetics/src/load_data.jl          |   5 +-
 moment_kinetics/src/moment_kinetics.jl    |   5 +-
 moment_kinetics/src/utils.jl              |   2 +-
 5 files changed, 102 insertions(+), 67 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index a30f70a0d..1c12e16f5 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -429,6 +429,7 @@ function setup_electron_io(io_input, vpa, vperp, z, r, composition, collisions,
         # write a header to the output file
         add_attribute!(fid, "file_info",
                        "Output initial electron state from the moment_kinetics code")
+        add_attribute!(fid, "pdf_electron_converged", false)
 
         # write some overview information to the output file
         write_overview!(fid, composition, collisions, parallel_io, evolve_density,
@@ -2883,12 +2884,13 @@ end
 
 """
     write_electron_state(scratch_electron, moments, t_params, t, io_initial_electron,
-                         t_idx, r, z, vperp, vpa)
+                         t_idx, r, z, vperp, vpa; pdf_electron_converged=false)
 
 Write the electron state to an output file.
 """
 function write_electron_state(scratch_electron, moments, t_params, t,
-                              io_or_file_info_initial_electron, t_idx, r, z, vperp, vpa)
+                              io_or_file_info_initial_electron, t_idx, r, z, vperp, vpa;
+                              pdf_electron_converged=false)
 
     @serial_region begin
         # Only read/write from first process in each 'block'
@@ -2912,6 +2914,11 @@ function write_electron_state(scratch_electron, moments, t_params, t,
         write_electron_moments_data_to_binary(scratch_electron, moments, t_params,
                                               t_params, io_initial_electron, t_idx, r, z)
 
+        if pdf_electron_converged
+            modify_attribute!(io_initial_electron.fid, "pdf_electron_converged",
+                              pdf_electron_converged)
+        end
+
         closefile && close(io_initial_electron.fid)
     end
 
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 6f62d38f6..b2a773200 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -41,7 +41,8 @@ using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..electron_kinetic_equation: update_electron_pdf!, enforce_boundary_condition_on_electron_pdf!
 using ..input_structs: boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons
 using ..derivatives: derivative_z!
-using ..utils: get_default_restart_filename, get_prefix_iblock_and_move_existing_file
+using ..utils: get_default_restart_filename, get_prefix_iblock_and_move_existing_file,
+               get_backup_filename
 
 using ..manufactured_solns: manufactured_solutions
 
@@ -573,35 +574,45 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
             previous_runs_info = nothing
             code_time = 0.0
             restart_time_index = -1
+            pdf_electron_converged = false
         else
-            if global_rank[] == 0
-                println("Restarting electrons from $restart_filename")
-            end
             # Previously-created electron distribution function exists, so use it as
             # the initial guess.
-            backup_prefix_iblock =
+            backup_prefix_iblock, initial_electrons_filename,
+            backup_initial_electrons_filename =
                 get_prefix_iblock_and_move_existing_file(restart_filename,
                                                          io_input.output_dir)
 
             # Reload pdf and moments from an existing output file
-            code_time, previous_runs_info, restart_time_index =
+            code_time, pdf_electron_converged, previous_runs_info, restart_time_index =
                 reload_electron_data!(pdf, moments, t_params, backup_prefix_iblock, -1,
                                       geometry, r, z, vpa, vperp, vzeta, vr, vz)
 
-            # Broadcast code_time from the root process of each shared-memory block (on which it
-            # might have been loaded from a restart file).
+            # Broadcast code_time and pdf_electron_converged from the root process of each
+            # shared-memory block (on which it might have been loaded from a restart
+            # file).
             code_time = MPI.Bcast(code_time, 0, comm_block[])
+            pdf_electron_converged = MPI.Bcast(pdf_electron_converged, 0, comm_block[])
+
+            if pdf_electron_converged
+                if global_rank[] == 0
+                    println("Reading initial electron state from $restart_filename")
+                end
+                # Move the *.initial_electron.h5 file back to its original location, as we
+                # do not need to create a new output file.
+                MPI.Barrier(comm_world)
+                if global_rank[] == 0
+                    if initial_electrons_filename != backup_initial_electrons_filename
+                        mv(backup_initial_electrons_filename, initial_electrons_filename)
+                    end
+                end
+                MPI.Barrier(comm_world)
+            else
+                if global_rank[] == 0
+                    println("Restarting electron initialisation from $restart_filename")
+                end
+            end
         end
-        # Setup I/O for initial electron state
-        io_initial_electron = setup_electron_io(io_input, vpa, vperp, z, r,
-                                                composition, collisions,
-                                                moments.evolve_density,
-                                                moments.evolve_upar,
-                                                moments.evolve_ppar,
-                                                external_source_settings, t_params,
-                                                input_dict, restart_time_index,
-                                                previous_runs_info,
-                                                "initial_electron")
 
         begin_serial_region()
         @serial_region begin
@@ -637,9 +648,6 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
         max_electron_pdf_iterations = 2000000
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
-        if global_rank[] == 0
-            println("Initializing electrons - evolving both pdf_electron and electron_ppar")
-        end
         if t_params.debug_io !== nothing
             io_electron = setup_electron_io(t_params.debug_io[1], vpa, vperp, z, r,
                                             composition, collisions,
@@ -657,39 +665,63 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
             resize!(t_params.dfns_output_times, n_truncated)
             t_params.dfns_output_times .= truncated_times
         end
-        electron_pseudotime, success =
-            @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
-                                        r, z, vperp, vpa, z_spectral, vperp_spectral,
-                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                        t_params, collisions, composition,
-                                        external_source_settings, num_diss_params,
-                                        max_electron_pdf_iterations;
-                                        io_electron=io_initial_electron,
-                                        initial_time=code_time,
-                                        residual_tolerance=t_input["initialization_residual_value"],
-                                        evolve_ppar=true)
-        if success != ""
-            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
-                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
-        end
+        if !pdf_electron_converged 
+            if global_rank[] == 0
+                println("Initializing electrons - evolving both pdf_electron and electron_ppar")
+            end
+            # Setup I/O for initial electron state
+            io_initial_electron = setup_electron_io(io_input, vpa, vperp, z, r,
+                                                    composition, collisions,
+                                                    moments.evolve_density,
+                                                    moments.evolve_upar,
+                                                    moments.evolve_ppar,
+                                                    external_source_settings, t_params,
+                                                    input_dict, restart_time_index,
+                                                    previous_runs_info,
+                                                    "initial_electron")
+
+            electron_pseudotime, success =
+                @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
+                                            r, z, vperp, vpa, z_spectral, vperp_spectral,
+                                            vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                            t_params, collisions, composition,
+                                            external_source_settings, num_diss_params,
+                                            max_electron_pdf_iterations;
+                                            io_electron=io_initial_electron,
+                                            initial_time=code_time,
+                                            residual_tolerance=t_input["initialization_residual_value"],
+                                            evolve_ppar=true)
+            if success != ""
+                error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+                      * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+            end
 
-        # Now run without evolve_ppar=true to get pdf_electron fully to steady state,
-        # ready for the start of the ion time advance.
-        if global_rank[] == 0
-            println("Initializing electrons - evolving pdf_electron only to steady state")
-        end
-        electron_pseudotime, success =
-            @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
-                                        r, z, vperp, vpa, z_spectral, vperp_spectral,
-                                        vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                        t_params, collisions, composition,
-                                        external_source_settings, num_diss_params,
-                                        max_electron_pdf_iterations;
-                                        io_electron=io_initial_electron,
-                                        initial_time=electron_pseudotime)
-        if success != ""
-            error("!!!max number of iterations for electron pdf update exceeded!!!\n"
-                  * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+            # Now run without evolve_ppar=true to get pdf_electron fully to steady state,
+            # ready for the start of the ion time advance.
+            if global_rank[] == 0
+                println("Initializing electrons - evolving pdf_electron only to steady state")
+            end
+            electron_pseudotime, success =
+                @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
+                                            r, z, vperp, vpa, z_spectral, vperp_spectral,
+                                            vpa_spectral, z_advect, vpa_advect, scratch_dummy,
+                                            t_params, collisions, composition,
+                                            external_source_settings, num_diss_params,
+                                            max_electron_pdf_iterations;
+                                            io_electron=io_initial_electron,
+                                            initial_time=electron_pseudotime)
+            if success != ""
+                error("!!!max number of iterations for electron pdf update exceeded!!!\n"
+                      * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
+            end
+
+            # Write the converged initial state for the electrons to a file so that it can be
+            # re-used if the simulation is re-run.
+            t_params.moments_output_counter[] += 1
+            write_electron_state(scratch_electron, moments, t_params, electron_pseudotime,
+                                 io_initial_electron, t_params.moments_output_counter[], r, z,
+                                 vperp, vpa; pdf_electron_converged=true)
+            finish_electron_io(io_initial_electron)
         end
 
         begin_r_z_vperp_vpa_region()
@@ -698,14 +730,6 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
                 pdf.electron.norm[ivpa,ivperp,iz,ir]
         end
 
-        # Write the converged initial state for the electrons to a file so that it can be
-        # re-used if the simulation is re-run.
-        t_params.moments_output_counter[] += 1
-        write_electron_state(scratch_electron, moments, t_params, electron_pseudotime,
-                             io_initial_electron, t_params.moments_output_counter[], r, z,
-                             vperp, vpa)
-        finish_electron_io(io_initial_electron)
-
         # No need to do electron I/O (apart from possibly debug I/O) any more, so if
         # adaptive timestep is used, it does not need to adjust to output times.
         resize!(t_params.moments_output_times, 0)
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 785c759de..24e5b12a5 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -1032,6 +1032,7 @@ Reload electron pdf and moments from an existing output file.
 function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, time_index,
                                geometry, r, z, vpa, vperp, vzeta, vr, vz)
     code_time = 0.0
+    pdf_electron_converged = false
     previous_runs_info = nothing
     begin_serial_region()
     @serial_region begin
@@ -1068,6 +1069,8 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
 
             code_time = load_slice(dynamic, "time", time_index)
 
+            pdf_electron_converged = get_attribute(fid, "pdf_electron_converged")
+
             r_range, z_range, vperp_range, vpa_range, vzeta_range, vr_range, vz_range =
                 get_reload_ranges(parallel_io, restart_r, restart_z, restart_vperp,
                                   restart_vpa, restart_vzeta, restart_vr, restart_vz)
@@ -1116,7 +1119,7 @@ function reload_electron_data!(pdf, moments, t_params, restart_prefix_iblock, ti
         end
     end
 
-    return code_time, previous_runs_info, time_index
+    return code_time, pdf_electron_converged, previous_runs_info, time_index
 end
 
 function load_restart_coordinates(fid, r, z, vperp, vpa, vzeta, vr, vz, parallel_io)
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 93e3d3598..dea0f2967 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -303,8 +303,9 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             restart_filename = restart
         end
 
-        backup_prefix_iblock = get_prefix_iblock_and_move_existing_file(restart_filename,
-                                                                        io_input.output_dir)
+        backup_prefix_iblock, _, _ =
+            get_prefix_iblock_and_move_existing_file(restart_filename,
+                                                     io_input.output_dir)
 
         # Reload pdf and moments from an existing output file
         code_time, dt, dt_before_last_fail, electron_dt, electron_dt_before_last_fail,
diff --git a/moment_kinetics/src/utils.jl b/moment_kinetics/src/utils.jl
index 1007f7425..2118c0407 100644
--- a/moment_kinetics/src/utils.jl
+++ b/moment_kinetics/src/utils.jl
@@ -274,7 +274,7 @@ function get_prefix_iblock_and_move_existing_file(restart_filename, output_dir)
     # Ensure files have been moved before any process tries to read from them
     MPI.Barrier(comm_world)
 
-    return backup_prefix_iblock
+    return backup_prefix_iblock, dfns_filename, backup_dfns_filename
 end
 
 """

From 08c626bf5600f88fba4d6c757becb555e43cde3d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 15:55:38 +0100
Subject: [PATCH 318/428] Fix region in distributed_dot_s_r_z_vperp_vpa()

---
 moment_kinetics/src/nonlinear_solvers.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl
index 7219c1530..e0872b102 100644
--- a/moment_kinetics/src/nonlinear_solvers.jl
+++ b/moment_kinetics/src/nonlinear_solvers.jl
@@ -565,7 +565,7 @@ function distributed_dot_s_r_z_vperp_vpa(v::AbstractArray{mk_float, 5},
     vperp = coords.vperp
     vpa = coords.vpa
 
-    begin_z_region()
+    begin_s_r_z_vperp_vpa_region()
 
     local_dot = 0.0
     if r.irank < r.nrank - 1

From 850fe0958c6c810389acd448db9afe7c0b661be1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 17:04:18 +0100
Subject: [PATCH 319/428] Move creation of scratch_dummy back into
 setup_time_advance!()

Was moved earlier when electron initialisation solve was called earlier,
but that has moved back into setup_time_advance() now, and it will be
useful to move the call to setup_dummy_and_buffer_arrays() later again.
---
 moment_kinetics/src/initial_conditions.jl |  4 ++--
 moment_kinetics/src/moment_kinetics.jl    | 18 ++++++------------
 moment_kinetics/src/time_advance.jl       |  7 +++++--
 3 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index b2a773200..2133156d6 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -131,8 +131,8 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                                vperp_spectral, vpa_spectral, vz_spectral, species,
                                collisions, external_source_settings,
-                               manufactured_solns_input, scratch_dummy, t_input,
-                               num_diss_params, advection_structs, io_input, input_dict)
+                               manufactured_solns_input, t_input, num_diss_params,
+                               advection_structs, io_input, input_dict)
     if manufactured_solns_input.use_for_init
         init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z,
                                              r, composition.n_ion_species,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index dea0f2967..e1a329336 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -101,7 +101,6 @@ using .type_definitions: mk_int
 using .utils: to_minutes, get_default_restart_filename,
               get_prefix_iblock_and_move_existing_file
 using .em_fields: setup_em_fields
-using .time_advance: setup_dummy_and_buffer_arrays
 using .time_advance: allocate_advection_structs
 
 @debug_detect_redundant_block_synchronize using ..communication: debug_detect_redundant_is_active
@@ -271,11 +270,6 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     # NB: the returned advection_structs are yet to be initialized
     advection_structs = allocate_advection_structs(composition, z, r, vpa, vperp, vz, vr, vzeta)
 
-    # setup dummy arrays & buffer arrays for z r MPI                             
-    n_neutral_species_alloc = max(1, composition.n_neutral_species)
-    scratch_dummy = setup_dummy_and_buffer_arrays(r.n, z.n, vpa.n, vperp.n, vz.n, vr.n, vzeta.n, 
-        composition.n_ion_species, n_neutral_species_alloc)
-
     if restart === false
         restarting = false
         # initialize f(z,vpa) and the lowest three v-space moments (density(z), upar(z) and ppar(z)),
@@ -284,8 +278,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
                               composition, r, z, vperp, vpa, vzeta, vr, vz,
                               z_spectral, r_spectral, vperp_spectral, vpa_spectral,
                               vz_spectral, species, collisions, external_source_settings,
-                              manufactured_solns_input, scratch_dummy, t_input,
-                              num_diss_params, advection_structs, io_input, input_dict)
+                              manufactured_solns_input, t_input, num_diss_params,
+                              advection_structs, io_input, input_dict)
         # initialize time variable
         code_time = 0.
         dt = nothing
@@ -343,16 +337,16 @@ function setup_moment_kinetics(input_dict::AbstractDict;
     # create arrays and do other work needed to setup
     # the main time advance loop -- including normalisation of f by density if requested
 
-    moments, spectral_objects, scratch, scratch_implicit, scratch_electron, advance,
-    advance_implicit, t_params, fp_arrays, gyroavs, manufactured_source_list,
+    moments, spectral_objects, scratch, scratch_implicit, scratch_electron, scratch_dummy,
+    advance, advance_implicit, t_params, fp_arrays, gyroavs, manufactured_source_list,
     nl_solver_params =
         setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrophase,
             vz_spectral, vr_spectral, vzeta_spectral, vpa_spectral, vperp_spectral,
             z_spectral, r_spectral, composition, moments, t_input, code_time, dt,
             dt_before_last_fail, electron_dt, electron_dt_before_last_fail, collisions,
             species, geometry, boundary_distributions, external_source_settings,
-            num_diss_params, manufactured_solns_input, advection_structs, scratch_dummy,
-            io_input, restarting, restart_electron_physics, input_dict)
+            num_diss_params, manufactured_solns_input, advection_structs, io_input,
+            restarting, restart_electron_physics, input_dict)
 
     # This is the closest we can get to the end time of the setup before writing it to the
     # output file
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index a557c1fb4..5f8e618c6 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -442,7 +442,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
                              electron_dt_before_last_fail_reload, collisions, species,
                              geometry, boundary_distributions, external_source_settings,
                              num_diss_params, manufactured_solns_input,
-                             advection_structs, scratch_dummy, io_input, restarting,
+                             advection_structs, io_input, restarting,
                              restart_electron_physics, input_dict)
     # define some local variables for convenience/tidiness
     n_ion_species = composition.n_ion_species
@@ -677,6 +677,9 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     end
     # setup dummy arrays & buffer arrays for z r MPI
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
+    scratch_dummy = setup_dummy_and_buffer_arrays(r.n, z.n, vpa.n, vperp.n, vz.n, vr.n,
+                                                  vzeta.n, composition.n_ion_species,
+                                                  n_neutral_species_alloc)
     # create arrays for Fokker-Planck collisions 
     if advance.explicit_weakform_fp_collisions
         fp_arrays = init_fokker_planck_collisions_weak_form(vpa,vperp,vpa_spectral,vperp_spectral; precompute_weights=true)
@@ -1006,7 +1009,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     _block_synchronize()
 
     return moments, spectral_objects, scratch, scratch_implicit, scratch_electron,
-           advance, advance_implicit, t_params, fp_arrays, gyroavs,
+           scratch_dummy, advance, advance_implicit, t_params, fp_arrays, gyroavs,
            manufactured_source_list, nl_solver_params
 end
 

From 31636dbf6b55ba9ec131d44fb59e169e4731a104 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 17:08:35 +0100
Subject: [PATCH 320/428] Functions in nonlinear_solvers for kinetic electron
 implicit solve

---
 moment_kinetics/src/nonlinear_solvers.jl | 276 ++++++++++++++++++++++-
 1 file changed, 268 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl
index e0872b102..441e29b38 100644
--- a/moment_kinetics/src/nonlinear_solvers.jl
+++ b/moment_kinetics/src/nonlinear_solvers.jl
@@ -39,6 +39,7 @@ using LinearAlgebra
 using MINPACK
 using MPI
 using SparseArrays
+using StatsBase: mean
 
 struct nl_solver_info{TH,TV,Tlig,Tprecon}
     rtol::mk_float
@@ -75,7 +76,7 @@ for example a preconditioner object for each point in that outer loop.
 """
 function setup_nonlinear_solve(input_dict, coords, outer_coords=(); default_rtol=1.0e-5,
                                default_atol=1.0e-12, serial_solve=false,
-                               preconditioner_type=nothing)
+                               electron_ppar_pdf_solve=false, preconditioner_type=nothing)
     nl_solver_section = set_defaults_and_check_section!(
         input_dict, "nonlinear_solver";
         rtol=default_rtol,
@@ -100,6 +101,19 @@ function setup_nonlinear_solve(input_dict, coords, outer_coords=(); default_rtol
         V = allocate_float(reverse(coord_sizes)..., linear_restart+1)
         H .= 0.0
         V .= 0.0
+    elseif electron_ppar_pdf_solve
+        H = allocate_shared_float(linear_restart + 1, linear_restart)
+        V_ppar = allocate_shared_float(coords.z.n, coords.r.n, linear_restart+1)
+        V_pdf = allocate_shared_float(reverse(coord_sizes)..., linear_restart+1)
+
+        begin_serial_region()
+        @serial_region begin
+            H .= 0.0
+            V_ppar .= 0.0
+            V_pdf .= 0.0
+        end
+
+        V = (V_ppar, V_pdf)
     else
         H = allocate_shared_float(linear_restart + 1, linear_restart)
         V = allocate_shared_float(reverse(coord_sizes)..., linear_restart+1)
@@ -379,6 +393,10 @@ function get_distributed_norm(coords, rtol, atol, x)
         this_norm = distributed_norm_z
     elseif dims == (:vpa,)
         this_norm = distributed_norm_vpa
+    elseif dims == (:r, :z, :vperp, :vpa)
+        # Intended for implicit solve combining electron_ppar and pdf_electron, so will
+        # not work for a single variable.
+        this_norm = distributed_norm_r_z_vperp_vpa
     elseif dims == (:s, :r, :z, :vperp, :vpa)
         this_norm = distributed_norm_s_r_z_vperp_vpa
     else
@@ -439,6 +457,75 @@ function distributed_norm_vpa(residual::AbstractArray{mk_float, 1}; coords, rtol
     return residual_norm
 end
 
+function distributed_norm_r_z_vperp_vpa(residual::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}};
+                                        coords, rtol, atol, x)
+    ppar_residual, pdf_residual = residual
+    x_ppar, x_pdf = x
+    r = coords.r
+    z = coords.z
+    vperp = coords.vperp
+    vpa = coords.vpa
+
+    if r.irank < r.nrank - 1
+        rend = r.n
+    else
+        rend = r.n + 1
+    end
+    if z.irank < z.nrank - 1
+        zend = z.n
+    else
+        zend = z.n + 1
+    end
+
+    begin_r_z_region()
+
+    ppar_local_norm_square = 0.0
+    @loop_r_z ir iz begin
+        if ir == rend || iz == zend
+            continue
+        end
+        ppar_local_norm_square += (ppar_residual[iz,ir] / (rtol * abs(x_ppar[iz,ir]) + atol))^2
+    end
+
+    _block_synchronize()
+    ppar_block_norm_square = MPI.Reduce(ppar_local_norm_square, +, comm_block[])
+
+    if block_rank[] == 0
+        ppar_global_norm_square = MPI.Allreduce(ppar_block_norm_square, +, comm_inter_block[])
+        ppar_global_norm_square = ppar_global_norm_square / (r.n_global * z.n_global)
+    else
+        ppar_global_norm_square = nothing
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    pdf_local_norm_square = 0.0
+    @loop_r_z ir iz begin
+        if ir == rend || iz == zend
+            continue
+        end
+        @loop_vperp_vpa ivperp ivpa begin
+            pdf_local_norm_square += (pdf_residual[ivpa,ivperp,iz,ir] / (rtol * abs(x_pdf[ivpa,ivperp,iz,ir]) + atol))^2
+        end
+    end
+
+    _block_synchronize()
+    pdf_block_norm_square = MPI.Reduce(pdf_local_norm_square, +, comm_block[])
+
+    if block_rank[] == 0
+        pdf_global_norm_square = MPI.Allreduce(pdf_block_norm_square, +, comm_inter_block[])
+        pdf_global_norm_square = pdf_global_norm_square / (r.n_global * z.n_global * vperp.n_global * vpa.n_global)
+
+        global_norm = sqrt(mean((ppar_global_norm_square, pdf_global_norm_square)))
+    else
+        global_norm = nothing
+    end
+
+    global_norm = MPI.bcast(global_norm, comm_block[]; root=0)
+
+    return global_norm
+end
+
 function distributed_norm_s_r_z_vperp_vpa(residual::AbstractArray{mk_float, 5};
                                           coords, rtol, atol, x)
     n_ion_species = coords.s
@@ -495,6 +582,10 @@ function get_distributed_dot(coords, rtol, atol, x)
         this_dot = distributed_dot_z
     elseif dims == (:vpa,)
         this_dot = distributed_dot_vpa
+    elseif dims == (:r, :z, :vperp, :vpa)
+        # Intended for implicit solve combining electron_ppar and pdf_electron, so will
+        # not work for a single variable.
+        this_dot = distributed_dot_r_z_vperp_vpa
     elseif dims == (:s, :r, :z, :vperp, :vpa)
         this_dot = distributed_dot_s_r_z_vperp_vpa
     else
@@ -556,6 +647,74 @@ function distributed_dot_vpa(v::AbstractArray{mk_float, 1}, w::AbstractArray{mk_
     return local_dot
 end
 
+function distributed_dot_r_z_vperp_vpa(v::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}},
+                                       w::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}};
+                                       coords, atol, rtol, x)
+    v_ppar, v_pdf = v
+    w_ppar, w_pdf = w
+    x_ppar, x_pdf = x
+
+    r = coords.r
+    z = coords.z
+    vperp = coords.vperp
+    vpa = coords.vpa
+
+    if r.irank < r.nrank - 1
+        rend = r.n
+    else
+        rend = r.n + 1
+    end
+    if z.irank < z.nrank - 1
+        zend = z.n
+    else
+        zend = z.n + 1
+    end
+
+    begin_r_z_region()
+
+    ppar_local_dot = 0.0
+    @loop_r_z ir iz begin
+        if ir == rend || iz == zend
+            continue
+        end
+        ppar_local_dot += v_ppar[iz,ir] * w_ppar[iz,ir] / (rtol * abs(x_ppar[iz,ir]) + atol)^2
+    end
+
+    _block_synchronize()
+    ppar_block_dot = MPI.Reduce(ppar_local_dot, +, comm_block[])
+
+    if block_rank[] == 0
+        ppar_global_dot = MPI.Allreduce(ppar_block_dot, +, comm_inter_block[])
+        ppar_global_dot = ppar_global_dot / (r.n_global * z.n_global * vperp.n_global * vpa.n_global)
+    else
+        ppar_global_dot = nothing
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    pdf_local_dot = 0.0
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        if ir == rend || iz == zend
+            continue
+        end
+        pdf_local_dot += v_pdf[ivpa,ivperp,iz,ir] * w_pdf[ivpa,ivperp,iz,ir] / (rtol * abs(x_pdf[ivpa,ivperp,iz,ir]) + atol)^2
+    end
+
+    _block_synchronize()
+    pdf_block_dot = MPI.Reduce(pdf_local_dot, +, comm_block[])
+
+    if block_rank[] == 0
+        pdf_global_dot = MPI.Allreduce(pdf_block_dot, +, comm_inter_block[])
+        pdf_global_dot = pdf_global_dot / (r.n_global * z.n_global * vperp.n_global * vpa.n_global)
+
+        global_dot = mean((ppar_global_dot, pdf_global_dot))
+    else
+        global_dot = nothing
+    end
+
+    return global_dot
+end
+
 function distributed_dot_s_r_z_vperp_vpa(v::AbstractArray{mk_float, 5},
                                          w::AbstractArray{mk_float, 5};
                                          coords, atol, rtol, x)
@@ -611,6 +770,10 @@ function get_parallel_map(coords)
         return parallel_map_z
     elseif dims == (:vpa,)
         return parallel_map_vpa
+    elseif dims == (:r, :z, :vperp, :vpa)
+        # Intended for implicit solve combining electron_ppar and pdf_electron, so will
+        # not work for a single variable.
+        return parallel_map_r_z_vperp_vpa
     elseif dims == (:s, :r, :z, :vperp, :vpa)
         return parallel_map_s_r_z_vperp_vpa
     else
@@ -678,6 +841,64 @@ function parallel_map_vpa(func, result::AbstractArray{mk_float, 1}, x1, x2)
     return nothing
 end
 
+function parallel_map_r_z_vperp_vpa(func, result::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}})
+
+    result_ppar, result_pdf = result
+
+    begin_r_z_region()
+
+    @loop_r_z ir iz begin
+        result_ppar[iz,ir] = func()
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        result_pdf[ivpa,ivperp,iz,ir] = func()
+    end
+
+    return nothing
+end
+function parallel_map_r_z_vperp_vpa(func, result::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}}, x1)
+
+    result_ppar, result_pdf = result
+    x1_ppar, x1_pdf = x1
+
+    begin_r_z_region()
+
+    @loop_r_z ir iz begin
+        result_ppar[iz,ir] = func(x1_ppar[iz,ir])
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        result_pdf[ivpa,ivperp,iz,ir] = func(x1_pdf[ivpa,ivperp,iz,ir])
+    end
+
+    return nothing
+end
+function parallel_map_r_z_vperp_vpa(func, result::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}}, x1, x2)
+
+    result_ppar, result_pdf = result
+    x1_ppar, x1_pdf = x1
+    x2_ppar, x2_pdf = x2
+
+    begin_r_z_region()
+
+    @loop_r_z ir iz begin
+        result_ppar[iz,ir] = func(x1_ppar[iz,ir], x2_ppar[iz,ir])
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        result_pdf[ivpa,ivperp,iz,ir] = func(x1_pdf[ivpa,ivperp,iz,ir], x2_pdf[ivpa,ivperp,iz,ir])
+    end
+
+    return nothing
+end
+
 function parallel_map_s_r_z_vperp_vpa(func, result::AbstractArray{mk_float, 5})
 
     begin_s_r_z_vperp_vpa_region()
@@ -721,6 +942,10 @@ function get_parallel_delta_x_calc(coords)
         return parallel_delta_x_calc_z
     elseif dims == (:vpa,)
         return parallel_delta_x_calc_vpa
+    elseif dims == (:r, :z, :vperp, :vpa)
+        # Intended for implicit solve combining electron_ppar and pdf_electron, so will
+        # not work for a single variable.
+        return parallel_delta_x_calc_r_z_vperp_vpa
     elseif dims == (:s, :r, :z, :vperp, :vpa)
         return parallel_delta_x_calc_s_r_z_vperp_vpa
     else
@@ -755,6 +980,33 @@ function parallel_delta_x_calc_vpa(delta_x::AbstractArray{mk_float, 1}, V, y)
     return nothing
 end
 
+function parallel_delta_x_calc_r_z_vperp_vpa(delta_x::Tuple{AbstractArray{mk_float, 2},AbstractArray{mk_float, 4}}, V, y)
+
+    delta_x_ppar, delta_x_pdf = delta_x
+    V_ppar, V_pdf = V
+    y_ppar, y_pdf = y
+
+    ny = length(y)
+
+    begin_r_z_region()
+
+    @loop_r_z ir iz begin
+        for iy ∈ 1:ny
+            delta_x_ppar[iz,ir] += y[iy] * V_ppar[iz,ir,iy]
+        end
+    end
+
+    begin_r_z_vperp_vpa_region()
+
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        for iy ∈ 1:ny
+            delta_x_pdf[ivpa,ivperp,iz,ir] += y[iy] * V_pdf[ivpa,ivperp,iz,ir,iy]
+        end
+    end
+
+    return nothing
+end
+
 function parallel_delta_x_calc_s_r_z_vperp_vpa(delta_x::AbstractArray{mk_float, 5}, V, y)
 
     begin_s_r_z_vperp_vpa_region()
@@ -769,6 +1021,14 @@ function parallel_delta_x_calc_s_r_z_vperp_vpa(delta_x::AbstractArray{mk_float,
     return nothing
 end
 
+# Utility function for neatness handling that V may be an array or a Tuple of arrays
+function select_from_V(V::Tuple, i)
+    return Tuple(selectdim(Vpart,ndims(Vpart),i) for Vpart ∈ V)
+end
+function select_from_V(V, i)
+    return selectdim(V,ndims(V),i)
+end
+
 """
 Apply the GMRES algorithm to solve the 'linear problem' J.δx^n = R(x^n), which is needed
 at each step of the outer Newton iteration (in `newton_solve!()`).
@@ -804,7 +1064,7 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
     # Now we actually set 'w' as the first Krylov vector, and normalise it.
     parallel_map((residual0, v) -> -residual0 - v, w, residual0, v)
     beta = distributed_norm(w)
-    parallel_map((w) -> w/beta, selectdim(V,ndims(V),1), w)
+    parallel_map((w) -> w/beta, select_from_V(V, 1), w)
 
     # Set tolerance for GMRES iteration to rtol times the initial residual, unless this is
     # so small that it is smaller than atol, in which case use atol instead.
@@ -820,12 +1080,12 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
             #println("Linear ", counter)
 
             # Compute next Krylov vector
-            parallel_map((V) -> V, w, selectdim(V,ndims(V),i))
+            parallel_map((V) -> V, w, select_from_V(V, i))
             approximate_Jacobian_vector_product!(w)
 
             # Gram-Schmidt orthogonalization
             for j ∈ 1:i
-                parallel_map((V) -> V, v, selectdim(V,ndims(V),j))
+                parallel_map((V) -> V, v, select_from_V(V, j))
                 w_dot_Vj = distributed_dot(w, v)
                 if serial_solve
                     H[j,i] = w_dot_Vj
@@ -835,7 +1095,7 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
                         H[j,i] = w_dot_Vj
                     end
                 end
-                parallel_map((w, V) -> w - H[j,i] * V, w, w, selectdim(V,ndims(V),j))
+                parallel_map((w, V) -> w - H[j,i] * V, w, w, select_from_V(V, j))
             end
             norm_w = distributed_norm(w)
             if serial_solve
@@ -846,7 +1106,7 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
                     H[i+1,i] = norm_w
                 end
             end
-            parallel_map((w) -> w / H[i+1,i], selectdim(V,ndims(V),i+1), w)
+            parallel_map((w) -> w / H[i+1,i], select_from_V(V, i+1), w)
 
             function temporary_residual!(result, guess)
                 #println("temporary residual ", size(result), " ", size(@view(H[1:i+1,1:i])), " ", size(guess))
@@ -912,9 +1172,9 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
         parallel_map((residual0, v) -> -residual0 - v, v, residual0, v)
         beta = distributed_norm(v)
         for i ∈ 2:length(y)
-            parallel_map(() -> 0.0, selectdim(V,ndims(V),i))
+            parallel_map(() -> 0.0, select_from_V(V, i))
         end
-        parallel_map((v) -> v/beta, selectdim(V,ndims(V),1), v)
+        parallel_map((v) -> v/beta, select_from_V(V, 1), v)
     end
 
     return counter

From b199e6dc778cb3373347fa9a8ec43420611b76e7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 12:36:14 +0100
Subject: [PATCH 321/428] Implicit solve for kinetic electron distribution
 function and ppar

---
 .../src/makie_post_processing.jl              |   5 +-
 .../src/electron_fluid_equations.jl           |  36 +++++
 .../src/electron_kinetic_equation.jl          | 141 ++++++++++++++++-
 moment_kinetics/src/input_structs.jl          |   1 +
 moment_kinetics/src/moment_kinetics_input.jl  |   3 +-
 moment_kinetics/src/time_advance.jl           | 144 +++++++++++++++---
 6 files changed, 298 insertions(+), 32 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index d97d6e67e..9d88f94ad 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -267,8 +267,9 @@ function makie_post_process(run_dir::Union{String,Tuple},
     end
 
     timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix)
-    if any(ri.composition.electron_physics ==
-           moment_kinetics.input_structs.kinetic_electrons for ri ∈ run_info)
+    if any((ri.composition.electron_physics ==
+                moment_kinetics.input_structs.kinetic_electrons
+            && !ri.t_input["implicit_electron_advance"]) for ri ∈ run_info)
         timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix, electron=true)
     end
 
diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 74816d1f2..55d14655b 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -232,6 +232,42 @@ function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron
     return nothing
 end
 
+"""
+    electron_energy_residual!(residual, electron_ppar_out, fvec_in, moments,
+                              collisions, composition, external_source_settings,
+                              num_diss_params, z, dt)
+
+The residual is a function whose input is `electron_ppar`, so that when it's output
+`residual` is zero, electron_ppar is the result of a backward-Euler timestep:
+  (f_out - f_in) / dt = RHS(f_out)
+⇒ (f_out - f_in) - dt*RHS(f_out) = 0
+
+This function assumes any needed moment derivatives are already calculated using
+`electron_ppar_out` and stored in `moments.electron`.
+"""
+function electron_energy_residual!(residual, electron_ppar_out, fvec_in, moments,
+                                   collisions, composition, external_source_settings,
+                                   num_diss_params, z, dt)
+    begin_r_z_region()
+    electron_ppar_in = fvec_in.electron_ppar
+    @loop_r_z ir iz begin
+        residual[iz,ir] = electron_ppar_in[iz,ir]
+    end
+    electron_energy_equation!(residual, electron_ppar_out,
+                              fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
+                              fvec_in.ppar, fvec_in.density_neutral,
+                              fvec_in.uz_neutral, fvec_in.pz_neutral,
+                              moments.electron, collisions, dt, composition,
+                              external_source_settings.electron, num_diss_params, z)
+    # Now
+    #   residual = f_in + dt*RHS(f_out)
+    # so update to desired residual
+    begin_r_z_region()
+    @loop_r_z ir iz begin
+        residual[iz,ir] = (electron_ppar_out[iz,ir] - residual[iz,ir])
+    end
+end
+
 """
 Add just the braginskii conduction contribution to the electron pressure, and assume that
 we have to calculate qpar and dqpar_dz from ppar within this function (they are not
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 841b0fad5..377d310f7 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -15,19 +15,24 @@ using ..communication
 using ..interpolation: interpolate_to_grid_1d!
 using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float
-using ..electron_fluid_equations: update_electron_vth_temperature!,
-                                  calculate_electron_qpar_from_pdf!
-using ..electron_fluid_equations: electron_energy_equation!
+using ..electron_fluid_equations: calculate_electron_moments!,
+                                  update_electron_vth_temperature!,
+                                  calculate_electron_qpar_from_pdf!,
+                                  calculate_electron_parallel_friction_force!
+using ..electron_fluid_equations: electron_energy_equation!, electron_energy_residual!
 using ..electron_z_advection: electron_z_advection!, update_electron_speed_z!
 using ..electron_vpa_advection: electron_vpa_advection!, update_electron_speed_vpa!
+using ..em_fields: update_phi!
 using ..external_sources: external_electron_source!
 using ..file_io: get_electron_io_info, write_electron_state, finish_electron_io
 using ..krook_collisions: electron_krook_collisions!
 using ..moment_constraints: hard_force_moment_constraints!
+using ..moment_kinetics_structs: scratch_pdf, scratch_electron_pdf, electron_pdf_substruct
+using ..nonlinear_solvers: newton_solve!
 using ..runge_kutta: rk_update_variable!, rk_loworder_solution!, local_error_norm,
                      adaptive_timestep_update_t_params!
 using ..utils: get_minimum_CFL_z, get_minimum_CFL_vpa
-using ..velocity_moments: integrate_over_vspace
+using ..velocity_moments: integrate_over_vspace, calculate_electron_moment_derivatives!
 
 """
 update_electron_pdf is a function that uses the electron kinetic equation 
@@ -553,6 +558,134 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     return time, success
 end
 
+"""
+    implicit_electron_advance!()
+
+Do an implicit solve which finds: the steady-state electron shape function \$g_e\$; the
+backward-Euler advanced electron pressure which is updated using \$g_e\$ at the new
+time-level.
+
+Implicit electron solve includes r-dimension. For 1D runs this makes no difference. In 2D
+it might or might not be necessary. If the r-dimension is not needed in the implicit
+solve, we would need to work on the parallelisation. The simplest option would be a
+non-parallelised outer loop over r, with each nonlinear solve being parallelised over
+{z,vperp,vpa}. More efficient might be to add an equivalent to the 'anyv' parallelisation
+used for the collision operator (e.g. 'anyzv'?) to allow the outer r-loop to be
+parallelised.
+"""
+function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, moments,
+                                    fields, collisions, composition,
+                                    external_source_settings, num_diss_params, r, z,
+                                    vperp, vpa, r_spectral, z_spectral, vpa_spectral,
+                                    z_advect, vpa_advect, gyroavs, scratch_dummy, dt,
+                                    nl_solver_params)
+
+    electron_ppar_out = fvec_out.electron_ppar
+    # Store the solved-for pdf in n_rk_stages+1, because this was the version that gets
+    # written to output for the explicit-electron-timestepping version.
+    pdf_electron_out = scratch_electron.pdf_electron
+
+    # Do a forward-Euler step for electron_ppar to get the initial guess.
+    # No equivalent for f_electron, as f_electron obeys a steady-state equation.
+    electron_energy_equation!(electron_ppar_out, fvec_in.electron_ppar,
+                              fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
+                              fvec_in.ppar, fvec_in.density_neutral,
+                              fvec_in.uz_neutral, fvec_in.pz_neutral,
+                              moments.electron, collisions, dt, composition,
+                              external_source_settings.electron, num_diss_params, z)
+
+    function residual_func!(residual, new_variables)
+        electron_ppar_residual, f_electron_residual = residual
+        electron_ppar_new, f_electron_new = new_variables
+
+        new_scratch = scratch_pdf(fvec_in.pdf, fvec_in.density, fvec_in.upar, fvec_in.ppar,
+                                  fvec_in.pperp, fvec_in.temp_z_s,
+                                  fvec_in.electron_density, fvec_in.electron_upar,
+                                  electron_ppar_new, fvec_in.electron_pperp,
+                                  fvec_in.electron_temp, fvec_in.pdf_neutral,
+                                  fvec_in.density_neutral, fvec_in.uz_neutral,
+                                  fvec_in.pz_neutral)
+        # Only the first entry in the `electron_pdf_substruct` will be used, so does not
+        # matter what we put in the second and third except that they have the right type.
+        new_pdf = (electron=electron_pdf_substruct(f_electron_new, f_electron_new,
+                                                   f_electron_new,),)
+        # Calculate heat flux and derivatives using new_variables
+        calculate_electron_moments!(new_scratch, new_pdf, moments, composition,
+                                    collisions, r, z, vpa)
+        calculate_electron_moment_derivatives!(moments, new_scratch, scratch_dummy, z,
+                                               z_spectral,
+                                               num_diss_params.electron.moment_dissipation_coefficient,
+                                               composition.electron_physics)
+
+        electron_energy_residual!(electron_ppar_residual, electron_ppar_new, fvec_in,
+                                  moments, collisions, composition,
+                                  external_source_settings, num_diss_params, z, dt)
+
+        # electron_kinetic_equation_euler_update!() just adds dt*d(g_e)/dt to the
+        # electron_pdf member of the first argument, so if we set the electron_pdf member
+        # of the first argument to zero, and pass dt=1, then it will evaluate the time
+        # derivative, which is the residual for a steady-state solution.
+        begin_r_z_vperp_vpa_region()
+        @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+            f_electron_residual[ivpa,ivperp,iz,ir] = 0.0
+        end
+        residual_scratch_electron = scratch_electron_pdf(f_electron_residual,
+                                                         electron_ppar_residual)
+        new_scratch_electron = scratch_electron_pdf(f_electron_new, electron_ppar_new)
+        electron_kinetic_equation_euler_update!(residual_scratch_electron,
+                                                new_scratch_electron, moments, z, vperp,
+                                                vpa, z_spectral, vpa_spectral, z_advect,
+                                                vpa_advect, scratch_dummy, collisions,
+                                                composition, external_source_settings,
+                                                num_diss_params, 1.0)
+        return nothing
+    end
+
+    residual = (scratch_dummy.implicit_buffer_zr_1,
+                scratch_dummy.implicit_buffer_vpavperpzr_1)
+    delta_x = (scratch_dummy.implicit_buffer_zr_2,
+               scratch_dummy.implicit_buffer_vpavperpzr_2)
+    rhs_delta = (scratch_dummy.implicit_buffer_zr_3,
+                 scratch_dummy.implicit_buffer_vpavperpzr_3)
+    v = (scratch_dummy.implicit_buffer_zr_4,
+         scratch_dummy.implicit_buffer_vpavperpzr_4)
+    w = (scratch_dummy.implicit_buffer_zr_5,
+         scratch_dummy.implicit_buffer_vpavperpzr_5)
+
+    newton_success = newton_solve!((electron_ppar_out, pdf_electron_out), residual_func!,
+                                   residual, delta_x, rhs_delta, v, w, nl_solver_params;
+                                   left_preconditioner=nothing,
+                                   right_preconditioner=nothing,
+                                   coords=(r=r, z=z, vperp=vperp, vpa=vpa))
+
+    # Fill pdf.electron.norm
+    non_scratch_pdf = pdf.electron.norm
+    begin_r_z_vperp_vpa_region()
+    @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
+        non_scratch_pdf[ivpa,ivperp,iz,ir] = pdf_electron_out[ivpa,ivperp,iz,ir]
+    end
+
+    # Update the electron parallel friction force.
+    # This does not actually do anything for kinetic electron runs at the moment, but
+    # include as a reminder to update this if/when we do include e-i collisions for
+    # kinetic electrons.
+    calculate_electron_parallel_friction_force!(
+        moments.electron.parallel_friction, fvec_out.electron_density,
+        fvec_out.electron_upar, fvec_out.upar, moments.electron.dT_dz,
+        composition.me_over_mi, collisions.nu_ei, composition.electron_physics)
+
+    # Solve for EM fields now that electrons are updated.
+    update_phi!(fields, fvec_out, vperp, z, r, composition, collisions, moments,
+                z_spectral, r_spectral, scratch_dummy, gyroavs)
+
+    if !newton_success
+        success = "kinetic-electrons"
+    else
+        success = ""
+    end
+    return success
+end
+
 function speedup_hack!(fvec_out, fvec_in, z_speedup_fac, z, vpa; evolve_ppar=false)
     # Divide by wpa to relax CFL condition at large wpa - only looking for steady
     # state here, so does not matter that this makes time evolution incorrect.
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 6910347c4..7bf4c63c7 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -81,6 +81,7 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     minimum_dt::mk_float
     maximum_dt::mk_float
     implicit_braginskii_conduction::Bool
+    implicit_electron_advance::Bool
     implicit_ion_advance::Bool
     implicit_vpa_advection::Bool
     write_after_fixed_step_count::Bool
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 3688fd718..b36880fc5 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -222,7 +222,8 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         minimum_dt=0.0,
         maximum_dt=Inf,
         implicit_braginskii_conduction=true,
-        implicit_ion_advance=true,
+        implicit_electron_advance=true,
+        implicit_ion_advance=false,
         implicit_vpa_advection=false,
         write_after_fixed_step_count=false,
         write_error_diagnostics=false,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 5f8e618c6..6dbd19481 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -46,7 +46,7 @@ using ..charge_exchange: ion_charge_exchange_collisions_1V!,
                          neutral_charge_exchange_collisions_1V!,
                          ion_charge_exchange_collisions_3V!,
                          neutral_charge_exchange_collisions_3V!
-using ..electron_kinetic_equation: update_electron_pdf!
+using ..electron_kinetic_equation: update_electron_pdf!, implicit_electron_advance!
 using ..ionization: ion_ionization_collisions_1V!, neutral_ionization_collisions_1V!,
                     ion_ionization_collisions_3V!, neutral_ionization_collisions_3V!,
                     constant_ionization_source!
@@ -153,6 +153,20 @@ struct scratch_dummy_arrays
     # needs to be shared memory
     buffer_vpavperpzrs_1::MPISharedArray{mk_float,5}
     buffer_vpavperpzrs_2::MPISharedArray{mk_float,5}
+    # buffers to hold moment quantities for implicit solves
+    implicit_buffer_zr_1::MPISharedArray{mk_float,2}
+    implicit_buffer_zr_2::MPISharedArray{mk_float,2}
+    implicit_buffer_zr_3::MPISharedArray{mk_float,2}
+    implicit_buffer_zr_4::MPISharedArray{mk_float,2}
+    implicit_buffer_zr_5::MPISharedArray{mk_float,2}
+    implicit_buffer_zr_6::MPISharedArray{mk_float,2}
+    # buffers to hold electron for implicit solves
+    implicit_buffer_vpavperpzr_1::MPISharedArray{mk_float,4}
+    implicit_buffer_vpavperpzr_2::MPISharedArray{mk_float,4}
+    implicit_buffer_vpavperpzr_3::MPISharedArray{mk_float,4}
+    implicit_buffer_vpavperpzr_4::MPISharedArray{mk_float,4}
+    implicit_buffer_vpavperpzr_5::MPISharedArray{mk_float,4}
+    implicit_buffer_vpavperpzr_6::MPISharedArray{mk_float,4}
     # buffers to hold ion pdf for implicit solves
     implicit_buffer_vpavperpzrs_1::MPISharedArray{mk_float,5}
     implicit_buffer_vpavperpzrs_2::MPISharedArray{mk_float,5}
@@ -296,8 +310,8 @@ If something is passed in `electron`, it is stored in the `electron_t_params` me
 the returned `time_info`.
 """
 function setup_time_info(t_input, n_variables, code_time, dt_reload,
-                         dt_before_last_fail_reload, manufactured_solns_input, io_input,
-                         input_dict; electron=nothing)
+                         dt_before_last_fail_reload, composition,
+                         manufactured_solns_input, io_input, input_dict; electron=nothing)
     rk_coefs, rk_coefs_implicit, implicit_coefficient_is_zero, n_rk_stages, rk_order,
     adaptive, low_storage, CFL_prefactor =
         setup_runge_kutta_coefficients!(t_input["type"],
@@ -374,8 +388,16 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
     if rk_coefs_implicit === nothing
         # Not an IMEX scheme, so cannot have any implicit terms
         t_input["implicit_braginskii_conduction"] = false
+        t_input["implicit_electron_advance"] = false
         t_input["implicit_ion_advance"] = false
         t_input["implicit_vpa_advection"] = false
+    else
+        if composition.electron_physics != braginskii_fluid
+            t_input["implicit_braginskii_conduction"] = false
+        end
+        if composition.electron_physics != kinetic_electrons
+            t_input["implicit_electron_advance"] = false
+        end
     end
 
     if t_input["high_precision_error_sum"]
@@ -418,6 +440,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      t_input["last_fail_proximity_factor"], t_input["minimum_dt"],
                      t_input["maximum_dt"],
                      electron !== nothing && t_input["implicit_braginskii_conduction"],
+                     electron !== nothing && t_input["implicit_electron_advance"],
                      electron !== nothing && t_input["implicit_ion_advance"],
                      electron !== nothing && t_input["implicit_vpa_advection"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
@@ -463,8 +486,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         electron_t_params = setup_time_info(t_input["electron_t_input"], 2, 0.0,
                                             electron_dt_reload,
                                             electron_dt_before_last_fail_reload,
-                                            manufactured_solns_input, io_input,
-                                            input_dict)
+                                            composition, manufactured_solns_input,
+                                            io_input, input_dict)
         # Make Vectors that count which variable caused timestep limits and timestep failures
         # the right length. Do this setup even when not using adaptive timestepping, because
         # it is easier than modifying the file I/O according to whether we are using adaptive
@@ -521,8 +544,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         end
     end
     t_params = setup_time_info(t_input, n_variables, code_time, dt_reload,
-                               dt_before_last_fail_reload, manufactured_solns_input,
-                               io_input, input_dict;
+                               dt_before_last_fail_reload, composition,
+                               manufactured_solns_input, io_input, input_dict;
                                electron=electron_t_params)
 
     # Make Vectors that count which variable caused timestep limits and timestep failures
@@ -624,6 +647,18 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     else
        electron_conduction_nl_solve_parameters = nothing
     end
+    if t_params.implicit_electron_advance
+        nl_solver_electron_advance_params =
+            setup_nonlinear_solve(input_dict,
+                                  (r=r, z=z, vperp=vperp, vpa=vpa),
+                                  ();
+                                  default_rtol=t_params.rtol / 10.0,
+                                  default_atol=t_params.atol / 10.0,
+                                  electron_ppar_pdf_solve=true,
+                                  preconditioner_type="lu")
+    else
+        nl_solver_electron_advance_params = nothing
+    end
     if t_params.implicit_ion_advance
         # Implicit solve for vpa_advection term should be done in serial, as it will be
         # called within a parallelised s_r_z_vperp loop.
@@ -656,6 +691,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
               * "time")
     end
     nl_solver_params = (electron_conduction=electron_conduction_nl_solve_parameters,
+                        electron_advance=nl_solver_electron_advance_params,
                         ion_advance=nl_solver_ion_advance_params,
                         vpa_advection=nl_solver_vpa_advection_params,)
 
@@ -679,7 +715,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     n_neutral_species_alloc = max(1,composition.n_neutral_species)
     scratch_dummy = setup_dummy_and_buffer_arrays(r.n, z.n, vpa.n, vperp.n, vz.n, vr.n,
                                                   vzeta.n, composition.n_ion_species,
-                                                  n_neutral_species_alloc)
+                                                  n_neutral_species_alloc, t_params)
     # create arrays for Fokker-Planck collisions 
     if advance.explicit_weakform_fp_collisions
         fp_arrays = init_fokker_planck_collisions_weak_form(vpa,vperp,vpa_spectral,vperp_spectral; precompute_weights=true)
@@ -1166,8 +1202,10 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         # if treating the electrons as a fluid with Braginskii closure, or
         # moment-kinetically then advance the electron energy equation
         if composition.electron_physics == kinetic_electrons
-            advance_electron_energy = true
-            advance_electron_conduction = true
+            if !t_params.implicit_electron_advance
+                advance_electron_energy = true
+                advance_electron_conduction = true
+            end
         elseif composition.electron_physics == braginskii_fluid
             if t_params.implicit_braginskii_conduction
                 # if treating the electrons as a fluid with Braginskii closure, and using
@@ -1327,6 +1365,11 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
         advance_electron_conduction = true
     end
 
+    if t_params.implicit_electron_advance
+        advance_electron_energy = true
+        advance_electron_conduction = true
+    end
+
     manufactured_solns_test = manufactured_solns_input.use_for_advance
 
     return advance_info(advance_vpa_advection, advance_vperp_advection, advance_z_advection, advance_r_advection,
@@ -1347,7 +1390,8 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
                         vpa_diffusion, vperp_diffusion, vz_diffusion)
 end
 
-function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies_ion,nspecies_neutral)
+function setup_dummy_and_buffer_arrays(nr, nz, nvpa, nvperp, nvz, nvr, nvzeta,
+                                       nspecies_ion, nspecies_neutral, t_params)
 
     dummy_s = allocate_float(nspecies_ion)
     dummy_sr = allocate_float(nr, nspecies_ion)
@@ -1422,12 +1466,51 @@ function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies
     buffer_vpavperpr_5 = allocate_shared_float(nvpa,nvperp,nr)
     buffer_vpavperpr_6 = allocate_shared_float(nvpa,nvperp,nr)
 
-    implicit_buffer_vpavperpzrs_1 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
-    implicit_buffer_vpavperpzrs_2 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
-    implicit_buffer_vpavperpzrs_3 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
-    implicit_buffer_vpavperpzrs_4 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
-    implicit_buffer_vpavperpzrs_5 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
-    implicit_buffer_vpavperpzrs_6 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+    if t_params.implicit_electron_advance
+        implicit_buffer_zr_1 = allocate_shared_float(nz,nr)
+        implicit_buffer_zr_2 = allocate_shared_float(nz,nr)
+        implicit_buffer_zr_3 = allocate_shared_float(nz,nr)
+        implicit_buffer_zr_4 = allocate_shared_float(nz,nr)
+        implicit_buffer_zr_5 = allocate_shared_float(nz,nr)
+        implicit_buffer_zr_6 = allocate_shared_float(nz,nr)
+
+        implicit_buffer_vpavperpzr_1 = allocate_shared_float(nvpa,nvperp,nz,nr)
+        implicit_buffer_vpavperpzr_2 = allocate_shared_float(nvpa,nvperp,nz,nr)
+        implicit_buffer_vpavperpzr_3 = allocate_shared_float(nvpa,nvperp,nz,nr)
+        implicit_buffer_vpavperpzr_4 = allocate_shared_float(nvpa,nvperp,nz,nr)
+        implicit_buffer_vpavperpzr_5 = allocate_shared_float(nvpa,nvperp,nz,nr)
+        implicit_buffer_vpavperpzr_6 = allocate_shared_float(nvpa,nvperp,nz,nr)
+    else
+        implicit_buffer_zr_1 = allocate_shared_float(0,0)
+        implicit_buffer_zr_2 = allocate_shared_float(0,0)
+        implicit_buffer_zr_3 = allocate_shared_float(0,0)
+        implicit_buffer_zr_4 = allocate_shared_float(0,0)
+        implicit_buffer_zr_5 = allocate_shared_float(0,0)
+        implicit_buffer_zr_6 = allocate_shared_float(0,0)
+
+        implicit_buffer_vpavperpzr_1 = allocate_shared_float(0,0,0,0)
+        implicit_buffer_vpavperpzr_2 = allocate_shared_float(0,0,0,0)
+        implicit_buffer_vpavperpzr_3 = allocate_shared_float(0,0,0,0)
+        implicit_buffer_vpavperpzr_4 = allocate_shared_float(0,0,0,0)
+        implicit_buffer_vpavperpzr_5 = allocate_shared_float(0,0,0,0)
+        implicit_buffer_vpavperpzr_6 = allocate_shared_float(0,0,0,0)
+    end
+
+    if t_params.implicit_ion_advance
+        implicit_buffer_vpavperpzrs_1 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+        implicit_buffer_vpavperpzrs_2 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+        implicit_buffer_vpavperpzrs_3 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+        implicit_buffer_vpavperpzrs_4 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+        implicit_buffer_vpavperpzrs_5 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+        implicit_buffer_vpavperpzrs_6 = allocate_shared_float(nvpa,nvperp,nz,nr,nspecies_ion)
+    else
+        implicit_buffer_vpavperpzrs_1 = allocate_shared_float(0,0,0,0,0)
+        implicit_buffer_vpavperpzrs_2 = allocate_shared_float(0,0,0,0,0)
+        implicit_buffer_vpavperpzrs_3 = allocate_shared_float(0,0,0,0,0)
+        implicit_buffer_vpavperpzrs_4 = allocate_shared_float(0,0,0,0,0)
+        implicit_buffer_vpavperpzrs_5 = allocate_shared_float(0,0,0,0,0)
+        implicit_buffer_vpavperpzrs_6 = allocate_shared_float(0,0,0,0,0)
+    end
 
     buffer_vzvrvzetazsn_1 = allocate_shared_float(nvz,nvr,nvzeta,nz,nspecies_neutral)
     buffer_vzvrvzetazsn_2 = allocate_shared_float(nvz,nvr,nvzeta,nz,nspecies_neutral)
@@ -1464,6 +1547,8 @@ function setup_dummy_and_buffer_arrays(nr,nz,nvpa,nvperp,nvz,nvr,nvzeta,nspecies
         buffer_vpavperpzs_1,buffer_vpavperpzs_2,buffer_vpavperpzs_3,buffer_vpavperpzs_4,buffer_vpavperpzs_5,buffer_vpavperpzs_6,
         buffer_vpavperprs_1,buffer_vpavperprs_2,buffer_vpavperprs_3,buffer_vpavperprs_4,buffer_vpavperprs_5,buffer_vpavperprs_6,
         buffer_vpavperpzrs_1,buffer_vpavperpzrs_2,
+        implicit_buffer_zr_1,implicit_buffer_zr_2,implicit_buffer_zr_3,implicit_buffer_zr_4,implicit_buffer_zr_5,implicit_buffer_zr_6,
+        implicit_buffer_vpavperpzr_1,implicit_buffer_vpavperpzr_2,implicit_buffer_vpavperpzr_3,implicit_buffer_vpavperpzr_4,implicit_buffer_vpavperpzr_5,implicit_buffer_vpavperpzr_6,
         implicit_buffer_vpavperpzrs_1,implicit_buffer_vpavperpzrs_2,implicit_buffer_vpavperpzrs_3,implicit_buffer_vpavperpzrs_4,implicit_buffer_vpavperpzrs_5,implicit_buffer_vpavperpzrs_6,
         buffer_vzvrvzetazsn_1,buffer_vzvrvzetazsn_2,buffer_vzvrvzetazsn_3,buffer_vzvrvzetazsn_4,buffer_vzvrvzetazsn_5,buffer_vzvrvzetazsn_6,
         buffer_vzvrvzetarsn_1,buffer_vzvrvzetarsn_2,buffer_vzvrvzetarsn_3,buffer_vzvrvzetarsn_4,buffer_vzvrvzetarsn_5,buffer_vzvrvzetarsn_6,
@@ -2260,7 +2345,7 @@ function apply_all_bcs_constraints_update_moments!(
         # to the beginning of the ion/neutral timestep, so the electron solution
         # calculated here would be discarded - we might as well skip calculating it in
         # that case.
-        if update_electrons && success == ""
+        if update_electrons && !t_params.implicit_electron_advance && success == ""
             _, kinetic_electron_success = update_electron_pdf!(
                scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
                z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect,
@@ -2739,6 +2824,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # The diagonal elements are equal to the Butcher 'a' coefficients
                 # rk_coefs_implicit[istage,istage]=a[istage,istage].
                 nl_success = backward_euler!(scratch_implicit[istage], scratch[istage],
+                                             scratch_electron[t_params.electron.n_rk_stages+1],
                                              pdf, fields, moments, advect_objects, vz, vr,
                                              vzeta, vpa, vperp, gyrophase, z, r, t,
                                              t_params.dt[] *
@@ -3182,19 +3268,27 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     return nothing
 end
 
-function backward_euler!(fvec_out, fvec_in, pdf, fields, moments, advect_objects, vz, vr,
-                         vzeta, vpa, vperp, gyrophase, z, r, t, dt, spectral_objects,
-                         composition, collisions, geometry, scratch_dummy,
-                         manufactured_source_list, external_source_settings,
-                         num_diss_params, gyroavs, nl_solver_params, advance, fp_arrays,
-                         istage)
+function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, moments,
+                         advect_objects, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, t,
+                         dt, spectral_objects, composition, collisions, geometry,
+                         scratch_dummy, manufactured_source_list,
+                         external_source_settings, num_diss_params, gyroavs,
+                         nl_solver_params, advance, fp_arrays, istage)
 
     vpa_spectral, vperp_spectral, r_spectral, z_spectral = spectral_objects.vpa_spectral, spectral_objects.vperp_spectral, spectral_objects.r_spectral, spectral_objects.z_spectral
     vz_spectral, vr_spectral, vzeta_spectral = spectral_objects.vz_spectral, spectral_objects.vr_spectral, spectral_objects.vzeta_spectral
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
-    if advance.electron_conduction
+    if composition.electron_physics == kinetic_electrons && advance.electron_energy
+        success = implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron,
+                                             moments, fields, collisions, composition,
+                                             external_source_settings, num_diss_params, r,
+                                             z, vperp, vpa, r_spectral, z_spectral,
+                                             vpa_spectral, z_advect, vpa_advect, gyroavs,
+                                             scratch_dummy, dt,
+                                             nl_solver_params.electron_advance)
+    elseif advance.electron_conduction
         success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,
                                                   z_spectral, composition, collisions,
                                                   scratch_dummy,

From 0a91f1fe5bc23c605e792bc047bfc9638341eaaa Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 13:43:24 +0100
Subject: [PATCH 322/428] Example input file for implicit kinetic electron
 solve in periodic box

---
 .../periodic_split3_kinetic-IMEX.toml         | 118 ++++++++++++++++++
 1 file changed, 118 insertions(+)
 create mode 100644 examples/kinetic-electrons/periodic_split3_kinetic-IMEX.toml

diff --git a/examples/kinetic-electrons/periodic_split3_kinetic-IMEX.toml b/examples/kinetic-electrons/periodic_split3_kinetic-IMEX.toml
new file mode 100644
index 000000000..ea91e79fa
--- /dev/null
+++ b/examples/kinetic-electrons/periodic_split3_kinetic-IMEX.toml
@@ -0,0 +1,118 @@
+#runtime_plots = true
+n_ion_species = 1
+n_neutral_species = 1
+electron_physics = "kinetic_electrons"
+evolve_moments_density = true
+evolve_moments_parallel_flow = true
+evolve_moments_parallel_pressure = true
+evolve_moments_conservation = true
+recycling_fraction = 0.5
+T_e = 1.0
+T_wall = 0.1
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.1
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 1.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.1
+z_IC_temperature_phase1 = 1.0
+vpa_IC_option1 = "gaussian"
+vpa_IC_density_amplitude1 = 1.0
+vpa_IC_density_phase1 = 0.0
+vpa_IC_upar_amplitude1 = 0.0
+vpa_IC_upar_phase1 = 0.0
+vpa_IC_temperature_amplitude1 = 0.0
+vpa_IC_temperature_phase1 = 0.0
+initial_density2 = 1.0
+initial_temperature2 = 1.0
+z_IC_option2 = "sinusoid"
+z_IC_density_amplitude2 = 0.001
+z_IC_density_phase2 = 0.0
+z_IC_upar_amplitude2 = 0.0
+z_IC_upar_phase2 = 0.0
+z_IC_temperature_amplitude2 = 0.0
+z_IC_temperature_phase2 = 0.0
+vpa_IC_option2 = "gaussian"
+vpa_IC_density_amplitude2 = 1.0
+vpa_IC_density_phase2 = 0.0
+vpa_IC_upar_amplitude2 = 0.0
+vpa_IC_upar_phase2 = 0.0
+vpa_IC_temperature_amplitude2 = 0.0
+vpa_IC_temperature_phase2 = 0.0
+charge_exchange_frequency = 0.75
+ionization_frequency = 0.0
+constant_ionization_rate = false
+nu_ei = 1000.0
+r_ngrid = 1
+r_nelement = 1
+z_ngrid = 17
+z_nelement = 16
+#z_nelement_local = 16
+z_bc = "periodic"
+z_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 6
+vpa_nelement = 31
+vpa_L = 12.0
+vpa_bc = "zero"
+vpa_discretization = "chebyshev_pseudospectral"
+vz_ngrid = 6
+vz_nelement = 31
+vz_L = 12.0
+vz_bc = "zero"
+vz_discretization = "chebyshev_pseudospectral"
+
+[timestepping]
+type = "KennedyCarpenterARK324"
+implicit_electron_advance = true
+implicit_ion_advance = false
+implicit_vpa_advection = false
+nstep = 1000000
+dt = 1.0e-6
+minimum_dt = 1.0e-7
+#maximum_dt = 2.0e-5
+rtol = 1.0e-7
+max_increase_factor_near_last_fail = 1.001
+last_fail_proximity_factor = 1.1
+max_increase_factor = 1.05
+nwrite = 1000 #10000
+nwrite_dfns = 1000 #100000
+steady_state_residual = true
+converged_residual_value = 1.0e-3
+
+[electron_timestepping]
+nstep = 5000000
+#nstep = 1
+dt = 2.0e-8
+maximum_dt = 1.0
+nwrite = 10000
+nwrite_dfns = 100000
+#type = "SSPRK4"
+type = "Fekete4(3)"
+rtol = 1.0e-6
+atol = 1.0e-14
+minimum_dt = 1.0e-10
+initialization_residual_value = 2.5
+converged_residual_value = 0.1 #1.0e-3
+#debug_io = 10000
+
+[nonlinear_solver]
+nonlinear_max_iterations = 100
+#rtol = 1.0e-9
+#atol = 1.0e-12
+
+[ion_numerical_dissipation]
+vpa_dissipation_coefficient = 1.0e0
+force_minimum_pdf_value = 0.0
+
+[electron_numerical_dissipation]
+vpa_dissipation_coefficient = 2.0e0
+force_minimum_pdf_value = 0.0
+
+[neutral_numerical_dissipation]
+vz_dissipation_coefficient = 1.0e-1
+force_minimum_pdf_value = 0.0
+
+[krook_collisions]
+use_krook = true

From 8ffa6dc832c384228d8d5498566a33c26f54295c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 20 Jun 2024 20:20:20 +0100
Subject: [PATCH 323/428] Initialise kinetic electron pdf using implicit solve

...when using implicit timestepping for the electrons.
---
 moment_kinetics/src/initial_conditions.jl | 149 +++++++++++++++-------
 moment_kinetics/src/time_advance.jl       |   4 +-
 2 files changed, 103 insertions(+), 50 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 2133156d6..548a2d60f 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -22,11 +22,13 @@ using ..communication
 using ..external_sources
 using ..interpolation: interpolate_to_grid_1d!
 using ..looping
+using ..electron_kinetic_equation: implicit_electron_advance!
 using ..em_fields: update_phi!
 using ..file_io: setup_electron_io, write_electron_state, finish_electron_io
 using ..load_data: reload_electron_data!
 using ..moment_kinetics_structs: scratch_pdf, pdf_substruct, electron_pdf_substruct,
                                  pdf_struct, moments_struct, boundary_distributions_struct
+using ..nonlinear_solvers: nl_solver_info
 using ..velocity_moments: integrate_over_vspace, integrate_over_neutral_vspace
 using ..velocity_moments: integrate_over_positive_vz, integrate_over_negative_vz
 using ..velocity_moments: create_moments_ion, create_moments_electron, create_moments_neutral
@@ -260,8 +262,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
                                vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                                vperp_spectral, vpa_spectral, collisions, gyroavs,
                                external_source_settings, scratch_dummy, scratch,
-                               scratch_electron, t_params, t_input, num_diss_params,
-                               advection_structs, io_input, input_dict;
+                               scratch_electron, nl_solver_params, t_params, t_input,
+                               num_diss_params, advection_structs, io_input, input_dict;
                                restart_electron_physics)
     
     moments.electron.dens_updated[] = false
@@ -442,6 +444,11 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         scratch[1].electron_ppar .= moments.electron.ppar
         scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
         scratch[1].electron_temp .= moments.electron.temp
+        scratch[t_params.electron.n_rk_stages+1].electron_density .= moments.electron.dens
+        scratch[t_params.electron.n_rk_stages+1].electron_upar .= moments.electron.upar
+        scratch[t_params.electron.n_rk_stages+1].electron_ppar .= moments.electron.ppar
+        scratch[t_params.electron.n_rk_stages+1].electron_pperp .= 0.0 #moments.electron.pperp
+        scratch[t_params.electron.n_rk_stages+1].electron_temp .= moments.electron.temp
     end
     if scratch_electron !== nothing
         begin_serial_region()
@@ -451,12 +458,12 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     end
 
     # initialize the electron pdf that satisfies the electron kinetic equation
-    initialize_electron_pdf!(scratch_electron, pdf, moments, fields.phi, r, z, vpa, vperp,
-                             vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
-                             advection_structs.electron_z_advect,
+    initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, fields, r, z, vpa,
+                             vperp, vzeta, vr, vz, r_spectral, z_spectral, vperp_spectral,
+                             vpa_spectral, advection_structs.electron_z_advect,
                              advection_structs.electron_vpa_advect, scratch_dummy,
                              collisions, composition, geometry, external_source_settings,
-                             num_diss_params, t_params.electron,
+                             num_diss_params, gyroavs, nl_solver_params, t_params,
                              t_input["electron_t_input"], io_input, input_dict)
 
     return nothing
@@ -548,11 +555,12 @@ function initialize_pdf!(pdf, moments, boundary_distributions, composition, r, z
     return nothing
 end
 
-function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa, vperp,
-                                  vzeta, vr, vz, z_spectral, vperp_spectral, vpa_spectral,
-                                  z_advect, vpa_advect, scratch_dummy, collisions,
-                                  composition, geometry, external_source_settings,
-                                  num_diss_params, t_params, t_input, io_input,
+function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, fields, r, z,
+                                  vpa, vperp, vzeta, vr, vz, r_spectral, z_spectral,
+                                  vperp_spectral, vpa_spectral, z_advect, vpa_advect,
+                                  scratch_dummy, collisions, composition, geometry,
+                                  external_source_settings, num_diss_params, gyroavs,
+                                  nl_solver_params, t_params, t_input, io_input,
                                   input_dict)
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
@@ -585,8 +593,9 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
 
             # Reload pdf and moments from an existing output file
             code_time, pdf_electron_converged, previous_runs_info, restart_time_index =
-                reload_electron_data!(pdf, moments, t_params, backup_prefix_iblock, -1,
-                                      geometry, r, z, vpa, vperp, vzeta, vr, vz)
+                reload_electron_data!(pdf, moments, t_params.electron,
+                                      backup_prefix_iblock, -1, geometry, r, z, vpa,
+                                      vperp, vzeta, vr, vz)
 
             # Broadcast code_time and pdf_electron_converged from the root process of each
             # shared-memory block (on which it might have been loaded from a restart
@@ -619,7 +628,7 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
             # update the electron pdf in the last scratch_electron (which will be copied
             # to the first entry as part of the pseudo-time-loop in
             # update_electron_pdf!()).
-            scratch_electron[t_params.n_rk_stages+1].pdf_electron .= pdf.electron.norm
+            scratch_electron[t_params.electron.n_rk_stages+1].pdf_electron .= pdf.electron.norm
         end
 
         begin_r_z_region()
@@ -648,22 +657,23 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
         max_electron_pdf_iterations = 2000000
         #max_electron_pdf_iterations = 500000
         #max_electron_pdf_iterations = 10000
-        if t_params.debug_io !== nothing
-            io_electron = setup_electron_io(t_params.debug_io[1], vpa, vperp, z, r,
-                                            composition, collisions,
+        if t_params.electron.debug_io !== nothing
+            io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z,
+                                            r, composition, collisions,
                                             moments.evolve_density, moments.evolve_upar,
                                             moments.evolve_ppar, external_source_settings,
-                                            t_params, t_params.debug_io[2], -1, nothing,
+                                            t_params.electron,
+                                            t_params.electron.debug_io[2], -1, nothing,
                                             "electron_debug")
         end
         if code_time > 0.0
-            tind = searchsortedfirst(t_params.moments_output_times, code_time)
-            n_truncated = length(t_params.moments_output_times) - tind
-            truncated_times = t_params.moments_output_times[tind+1:end]
-            resize!(t_params.moments_output_times, n_truncated)
-            t_params.moments_output_times .= truncated_times
-            resize!(t_params.dfns_output_times, n_truncated)
-            t_params.dfns_output_times .= truncated_times
+            tind = searchsortedfirst(t_params.electron.moments_output_times, code_time)
+            n_truncated = length(t_params.electron.moments_output_times) - tind
+            truncated_times = t_params.electron.moments_output_times[tind+1:end]
+            resize!(t_params.electron.moments_output_times, n_truncated)
+            t_params.electron.moments_output_times .= truncated_times
+            resize!(t_params.electron.dfns_output_times, n_truncated)
+            t_params.electron.dfns_output_times .= truncated_times
         end
         if !pdf_electron_converged 
             if global_rank[] == 0
@@ -675,16 +685,18 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
                                                     moments.evolve_density,
                                                     moments.evolve_upar,
                                                     moments.evolve_ppar,
-                                                    external_source_settings, t_params,
-                                                    input_dict, restart_time_index,
+                                                    external_source_settings,
+                                                    t_params.electron, input_dict,
+                                                    restart_time_index,
                                                     previous_runs_info,
                                                     "initial_electron")
 
             electron_pseudotime, success =
-                @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
-                                            r, z, vperp, vpa, z_spectral, vperp_spectral,
-                                            vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                            t_params, collisions, composition,
+                @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments,
+                                            fields.phi, r, z, vperp, vpa, z_spectral,
+                                            vperp_spectral, vpa_spectral, z_advect,
+                                            vpa_advect, scratch_dummy, t_params.electron,
+                                            collisions, composition,
                                             external_source_settings, num_diss_params,
                                             max_electron_pdf_iterations;
                                             io_electron=io_initial_electron,
@@ -701,15 +713,55 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
             if global_rank[] == 0
                 println("Initializing electrons - evolving pdf_electron only to steady state")
             end
-            electron_pseudotime, success =
-                @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments, phi,
-                                            r, z, vperp, vpa, z_spectral, vperp_spectral,
-                                            vpa_spectral, z_advect, vpa_advect, scratch_dummy,
-                                            t_params, collisions, composition,
-                                            external_source_settings, num_diss_params,
-                                            max_electron_pdf_iterations;
-                                            io_electron=io_initial_electron,
-                                            initial_time=electron_pseudotime)
+            if t_params.implicit_electron_advance
+                # Create new nl_solver_info ojbect with higher maximum iterations for
+                # initialisation.
+                initialisation_nl_solver_params =
+                    nl_solver_info(nl_solver_params.electron_advance.rtol,
+                                   nl_solver_params.electron_advance.atol,
+                                   100000,
+                                   nl_solver_params.electron_advance.linear_rtol,
+                                   nl_solver_params.electron_advance.linear_atol,
+                                   nl_solver_params.electron_advance.linear_restart,
+                                   nl_solver_params.electron_advance.linear_max_restarts,
+                                   nl_solver_params.electron_advance.H,
+                                   nl_solver_params.electron_advance.V,
+                                   nl_solver_params.electron_advance.linear_initial_guess,
+                                   nl_solver_params.electron_advance.n_solves,
+                                   nl_solver_params.electron_advance.nonlinear_iterations,
+                                   nl_solver_params.electron_advance.linear_iterations,
+                                   nl_solver_params.electron_advance.global_n_solves,
+                                   nl_solver_params.electron_advance.global_nonlinear_iterations,
+                                   nl_solver_params.electron_advance.global_linear_iterations,
+                                   nl_solver_params.electron_advance.stage_counter,
+                                   nl_solver_params.electron_advance.serial_solve,
+                                   nl_solver_params.electron_advance.max_nonlinear_iterations_this_step,
+                                   nl_solver_params.electron_advance.preconditioner_update_interval,
+                                   nl_solver_params.electron_advance.preconditioners,
+                                  )
+                # Run implicit solve with dt=0 so that we don't update electron_ppar here
+                success =
+                    implicit_electron_advance!(scratch[t_params.n_rk_stages+1],
+                                               scratch[1], pdf,
+                                               scratch_electron[t_params.electron.n_rk_stages+1],
+                                               moments, fields, collisions, composition,
+                                               external_source_settings, num_diss_params,
+                                               r, z, vperp, vpa, r_spectral, z_spectral,
+                                               vpa_spectral, z_advect, vpa_advect,
+                                               gyroavs, scratch_dummy, 0.0,
+                                               initialisation_nl_solver_params)
+            else
+                electron_pseudotime, success =
+                    update_electron_pdf!(scratch_electron, pdf.electron.norm, moments,
+                                         fields.phi, r, z, vperp, vpa, z_spectral,
+                                         vperp_spectral, vpa_spectral, z_advect,
+                                         vpa_advect, scratch_dummy, t_params.electron,
+                                         collisions, composition,
+                                         external_source_settings, num_diss_params,
+                                         max_electron_pdf_iterations;
+                                         io_electron=io_initial_electron,
+                                         initial_time=electron_pseudotime)
+            end
             if success != ""
                 error("!!!max number of iterations for electron pdf update exceeded!!!\n"
                       * "Stopping at $(Dates.format(now(), dateformat"H:MM:SS"))")
@@ -717,10 +769,11 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
 
             # Write the converged initial state for the electrons to a file so that it can be
             # re-used if the simulation is re-run.
-            t_params.moments_output_counter[] += 1
-            write_electron_state(scratch_electron, moments, t_params, electron_pseudotime,
-                                 io_initial_electron, t_params.moments_output_counter[], r, z,
-                                 vperp, vpa; pdf_electron_converged=true)
+            t_params.electron.moments_output_counter[] += 1
+            write_electron_state(scratch_electron, moments, t_params.electron,
+                                 electron_pseudotime, io_initial_electron,
+                                 t_params.electron.moments_output_counter[], r, z, vperp,
+                                 vpa; pdf_electron_converged=true)
             finish_electron_io(io_initial_electron)
         end
 
@@ -732,10 +785,10 @@ function initialize_electron_pdf!(scratch_electron, pdf, moments, phi, r, z, vpa
 
         # No need to do electron I/O (apart from possibly debug I/O) any more, so if
         # adaptive timestep is used, it does not need to adjust to output times.
-        resize!(t_params.moments_output_times, 0)
-        resize!(t_params.dfns_output_times, 0)
-        t_params.moments_output_counter[] = 1
-        t_params.dfns_output_counter[] = 1
+        resize!(t_params.electron.moments_output_times, 0)
+        resize!(t_params.electron.dfns_output_times, 0)
+        t_params.electron.moments_output_counter[] = 1
+        t_params.electron.dfns_output_counter[] = 1
 
     end
     return nothing
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 6dbd19481..6f802d6de 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -740,8 +740,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
                               vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
                               vperp_spectral, vpa_spectral, collisions, gyroavs,
                               external_source_settings, scratch_dummy, scratch,
-                              scratch_electron, t_params, t_input, num_diss_params,
-                              advection_structs, io_input, input_dict;
+                              scratch_electron, nl_solver_params, t_params, t_input,
+                              num_diss_params, advection_structs, io_input, input_dict;
                               restart_electron_physics=restart_electron_physics)
     elseif restarting && composition.electron_physics == kinetic_electrons &&
            t_params.electron.debug_io !== nothing

From 75e6c96f2be64dae1bfd5bcd974775af697ecf9e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 21 Jun 2024 15:13:43 +0100
Subject: [PATCH 324/428] Apply bc and constraints to f_new in implicit
 electron solve

---
 .../src/electron_kinetic_equation.jl          | 22 +++++++++++++------
 moment_kinetics/src/initial_conditions.jl     |  4 ++--
 moment_kinetics/src/time_advance.jl           |  4 ++--
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 377d310f7..ad6c86330 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -326,7 +326,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             apply_electron_bc_and_constraints!(scratch[istage+1], phi, moments, z, vperp,
                                                vpa, vperp_spectral, vpa_spectral,
                                                vpa_advect, num_diss_params, composition,
-                                               t_params)
+                                               t_params.rtol)
 
             latest_pdf = scratch[istage+1].pdf_electron
             
@@ -576,9 +576,9 @@ parallelised.
 function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, moments,
                                     fields, collisions, composition,
                                     external_source_settings, num_diss_params, r, z,
-                                    vperp, vpa, r_spectral, z_spectral, vpa_spectral,
-                                    z_advect, vpa_advect, gyroavs, scratch_dummy, dt,
-                                    nl_solver_params)
+                                    vperp, vpa, r_spectral, z_spectral, vperp_spectral,
+                                    vpa_spectral, z_advect, vpa_advect, gyroavs,
+                                    scratch_dummy, dt, nl_solver_params)
 
     electron_ppar_out = fvec_out.electron_ppar
     # Store the solved-for pdf in n_rk_stages+1, because this was the version that gets
@@ -605,6 +605,13 @@ function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, mo
                                   fvec_in.electron_temp, fvec_in.pdf_neutral,
                                   fvec_in.density_neutral, fvec_in.uz_neutral,
                                   fvec_in.pz_neutral)
+        new_scratch_electron = scratch_electron_pdf(f_electron_new, electron_ppar_new)
+
+        apply_electron_bc_and_constraints!(new_scratch_electron, fields.phi, moments, z,
+                                           vperp, vpa, vperp_spectral, vpa_spectral,
+                                           vpa_advect, num_diss_params, composition,
+                                           nl_solver_params.rtol)
+
         # Only the first entry in the `electron_pdf_substruct` will be used, so does not
         # matter what we put in the second and third except that they have the right type.
         new_pdf = (electron=electron_pdf_substruct(f_electron_new, f_electron_new,
@@ -730,7 +737,7 @@ end
 
 function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp, vpa,
                                             vperp_spectral, vpa_spectral, vpa_advect,
-                                            num_diss_params, composition, t_params)
+                                            num_diss_params, composition, rtol)
     latest_pdf = this_scratch.pdf_electron
 
     begin_r_z_vperp_vpa_region()
@@ -744,7 +751,7 @@ function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp
                                                 vperp_spectral, vpa_spectral, vpa_advect,
                                                 moments,
                                                 num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                composition.me_over_mi, 0.1 * t_params.rtol)
+                                                composition.me_over_mi, 0.1 * rtol)
 
     begin_r_z_region()
     A = moments.electron.constraints_A_coefficient
@@ -1257,7 +1264,8 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
     end
     apply_electron_bc_and_constraints!(scratch[t_params.n_rk_stages+1], phi, moments, z,
                                        vperp, vpa, vperp_spectral, vpa_spectral,
-                                       vpa_advect, num_diss_params, composition, t_params)
+                                       vpa_advect, num_diss_params, composition,
+                                       t_params.rtol)
     if evolve_ppar
         # Reset vth in the `moments` struct to the result consistent with full-accuracy RK
         # solution.
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 548a2d60f..3a0a7f198 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -747,8 +747,8 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                                moments, fields, collisions, composition,
                                                external_source_settings, num_diss_params,
                                                r, z, vperp, vpa, r_spectral, z_spectral,
-                                               vpa_spectral, z_advect, vpa_advect,
-                                               gyroavs, scratch_dummy, 0.0,
+                                               vperp_spectral, vpa_spectral, z_advect,
+                                               vpa_advect, gyroavs, scratch_dummy, 0.0,
                                                initialisation_nl_solver_params)
             else
                 electron_pseudotime, success =
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 6f802d6de..881d962df 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -3285,8 +3285,8 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                              moments, fields, collisions, composition,
                                              external_source_settings, num_diss_params, r,
                                              z, vperp, vpa, r_spectral, z_spectral,
-                                             vpa_spectral, z_advect, vpa_advect, gyroavs,
-                                             scratch_dummy, dt,
+                                             vperp_spectral, vpa_spectral, z_advect,
+                                             vpa_advect, gyroavs, scratch_dummy, dt,
                                              nl_solver_params.electron_advance)
     elseif advance.electron_conduction
         success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,

From 51916c80ad9f34b230fcc42a84e6596001238057 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 21 Jun 2024 15:15:15 +0100
Subject: [PATCH 325/428] Calculate electron moment derivatives at start of
 electron implicit solve

Ensures `moments.electron.dupar_dz` is initialised correctly.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index ad6c86330..55ae19b2f 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -587,6 +587,9 @@ function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, mo
 
     # Do a forward-Euler step for electron_ppar to get the initial guess.
     # No equivalent for f_electron, as f_electron obeys a steady-state equation.
+    calculate_electron_moment_derivatives!(moments, fvec_in, scratch_dummy, z, z_spectral,
+                                           num_diss_params.electron.moment_dissipation_coefficient,
+                                           composition.electron_physics)
     electron_energy_equation!(electron_ppar_out, fvec_in.electron_ppar,
                               fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
                               fvec_in.ppar, fvec_in.density_neutral,

From 4a99612ba44b2712ee14a8c6430296f1efc261b7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 21 Jun 2024 15:18:40 +0100
Subject: [PATCH 326/428] Apply bcs and constraints to residual in implicit
 electron solve

---
 .../src/electron_kinetic_equation.jl          | 61 ++++++++++++++++++-
 1 file changed, 60 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 55ae19b2f..e85b7f154 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -26,7 +26,8 @@ using ..em_fields: update_phi!
 using ..external_sources: external_electron_source!
 using ..file_io: get_electron_io_info, write_electron_state, finish_electron_io
 using ..krook_collisions: electron_krook_collisions!
-using ..moment_constraints: hard_force_moment_constraints!
+using ..moment_constraints: hard_force_moment_constraints!,
+                            moment_constraints_on_residual!
 using ..moment_kinetics_structs: scratch_pdf, scratch_electron_pdf, electron_pdf_substruct
 using ..nonlinear_solvers: newton_solve!
 using ..runge_kutta: rk_update_variable!, rk_loworder_solution!, local_error_norm,
@@ -648,6 +649,64 @@ function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, mo
                                                 vpa_advect, scratch_dummy, collisions,
                                                 composition, external_source_settings,
                                                 num_diss_params, 1.0)
+
+        # Set residual to zero where pdf_electron is determined by boundary conditions.
+        if vpa.n > 1
+            begin_r_z_vperp_region()
+            @loop_r_z_vperp ir iz ivperp begin
+                @views enforce_v_boundary_condition_local!(f_electron_residual[:,ivperp,iz,ir], vpa.bc,
+                                                           vpa_advect[1].speed[:,ivperp,iz,ir],
+                                                           num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
+                                                           vpa, vpa_spectral)
+            end
+        end
+        if vperp.n > 1
+            begin_r_z_vpa_region()
+            enforce_vperp_boundary_condition!(f_electron_residual, vperp.bc, vperp, vperp_spectral,
+                                              vperp_adv, vperp_diffusion)
+        end
+        if z.bc == "wall" && (z.irank == 0 || z.irank == z.nrank - 1)
+            # Wall boundary conditions. Note that as density, upar, ppar do not
+            # change in this implicit step, f_new, f_old, and residual should all
+            # be zero at exactly the same set of grid points, so it is reasonable
+            # to zero-out `residual` to impose the boundary condition. We impose
+            # this after subtracting f_old in case rounding errors, etc. mean that
+            # at some point f_old had a different boundary condition cut-off
+            # index.
+            begin_r_vperp_vpa_region()
+            v_unnorm = vpa.scratch
+            zero = 1.0e-14
+            if z.irank == 0
+                iz = 1
+                @loop_r ir begin
+                    v_unnorm .= vpagrid_to_dzdt(vpa.grid, moments.electron.vth[iz,ir],
+                                                fvec_in.electron_upar[iz,ir], true, true)
+                    @loop_vperp_vpa ivperp ivpa begin
+                        if v_unnorm > -zero
+                            f_electron_residual[ivpa,ivperp,iz,ir] .= 0.0
+                        end
+                    end
+                end
+            end
+            if z.irank == z.nrank - 1
+                iz = z.n
+                @loop_r ir begin
+                    v_unnorm .= vpagrid_to_dzdt(vpa.grid, moments.electron.vth[iz,ir],
+                                                fvec_in.electron_upar[iz,ir], true, true)
+                    @loop_vperp_vpa ivpa ivperp begin
+                        if v_unnorm < zero
+                            f_electron_residual[ivpa,ivperp,iz,ir] .= 0.0
+                        end
+                    end
+                end
+            end
+        end
+        begin_r_z_region()
+        @loop_r_z ir iz begin
+            @views moment_constraints_on_residual!(f_electron_residual[:,:,iz,ir], f_electron_new[:,:,iz,ir],
+                                                   (evolve_density=true, evolve_upar=true, evolve_ppar=true),
+                                                   vpa)
+        end
         return nothing
     end
 

From 83c93accefdee4be76cc33684f865f17e37a6172 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 11:32:58 +0100
Subject: [PATCH 327/428] Add vperp.bc option "zero-no-regularity", intended
 for use with the collision operator only (or any diffusion operator that has
 no flux at vperp = 0 -- so not the numerical vperp diffusion at present --
 where zero boundary conditions are imposed at vperp = L but no condition is
 imposed at vperp = 0. Initial short tests suggests that Maxwellian relaxation
 problems might run successfully without the regularity conditions -TBC.

---
 moment_kinetics/src/boundary_conditions.jl    |  6 ++++--
 moment_kinetics/src/fokker_planck_calculus.jl | 10 +++++++---
 moment_kinetics/src/moment_kinetics_input.jl  |  4 +++-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index 61dff1a08..7082c1383 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -972,7 +972,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vpe
 end
 
 function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral, vperp_advect, diffusion)
-    if bc == "zero"
+    if bc == "zero" || bc == "zero-no-regularity"
         nvperp = vperp.n
         ngrid = vperp.ngrid
         # set zero boundary condition
@@ -982,7 +982,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vpe
             end
         end
         # set regularity condition d F / d vperp = 0 at vperp = 0
-        if vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral"
+        if bc == "zero" && (vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral")
             D0 = vperp_spectral.radau.D0
             buffer = @view vperp.scratch[1:ngrid-1]
             @loop_r_z_vpa ir iz ivpa begin
@@ -992,6 +992,8 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vpe
                     f[ivpa,1,iz,ir] = -sum(buffer)/D0[1]
                 end
             end
+        elseif bc == "zero-no-regularity"
+            # do nothing
         else
             println("vperp.bc=\"$bc\" not supported by discretization "
                     * "$(vperp.discretization)")
diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl
index 300e0ca6e..7838608fe 100644
--- a/moment_kinetics/src/fokker_planck_calculus.jl
+++ b/moment_kinetics/src/fokker_planck_calculus.jl
@@ -2297,11 +2297,15 @@ function enforce_vpavperp_BCs!(pdf,vpa,vperp,vpa_spectral,vperp_spectral)
     # set regularity condition d F / d vperp = 0 at vperp = 0
     # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
     begin_anyv_vpa_region()
-    buffer = @view vperp.scratch[1:ngrid_vperp-1]
     @loop_vpa ivpa begin
         pdf[ivpa,nvperp] = 0.0
-        @views @. buffer = D0[2:ngrid_vperp] * pdf[ivpa,2:ngrid_vperp]
-        pdf[ivpa,1] = -sum(buffer)/D0[1]
+    end
+    if vperp.bc == "zero"
+        buffer = @view vperp.scratch[1:ngrid_vperp-1]
+        @loop_vpa ivpa begin
+            @views @. buffer = D0[2:ngrid_vperp] * pdf[ivpa,2:ngrid_vperp]
+            pdf[ivpa,1] = -sum(buffer)/D0[1]
+        end
     end
 end
 
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 56358a27b..8e0772442 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1196,6 +1196,8 @@ function check_coordinate_input(coord, coord_name, io)
         println(io,">$coord_name.bc = 'constant'.  enforcing constant incoming BC in $coord_name.")
     elseif coord.bc == "zero"
         println(io,">$coord_name.bc = 'zero'.  enforcing zero incoming BC in $coord_name. Enforcing zero at both boundaries if diffusion operator is present.")
+    elseif coord.bc == "zero-no-regularity"
+        println(io,">$coord_name.bc = 'zero'.  enforcing zero incoming BC in $coord_name. Enforcing zero at both boundaries if diffusion operator is present. Do not enforce dF/dcoord = 0 at origin if coord = vperp.")
     elseif coord.bc == "zero_gradient"
         println(io,">$coord_name.bc = 'zero_gradient'.  enforcing zero gradients at both limits of $coord_name domain.")
     elseif coord.bc == "both_zero"
@@ -1216,7 +1218,7 @@ function check_coordinate_input(coord, coord_name, io)
 
         if coord.bc != "zero" && coord.n_global > 1 && global_rank[] == 0
             println("WARNING: regularity condition (df/dvperp=0 at vperp=0) not being "
-                    * "imposed. Collisions or vperp-diffusion will be unstable.")
+                    * "imposed. Collisions or vperp-diffusion may be unstable.")
         end
     else
         println(io,">using ", coord.ngrid, " grid points per $coord_name element on ",

From 1d38f20ac6e79d3814f0d7a79c565b7655fae0cf Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 12:54:15 +0100
Subject: [PATCH 328/428] Attempt to add new test for the fokker planck time
 evolution, currently failing.

---
 .../fokker_planck_time_evolution_tests.jl     | 117 ++++++++++++++----
 1 file changed, 96 insertions(+), 21 deletions(-)

diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index 08239c0f7..e11563d24 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -34,7 +34,7 @@ struct expected_data
     f_ion::Array{mk_float, 3} # vpa, vperp, time
 end
 
-const expected =
+const expected_base =
   expected_data(
    [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
    [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
@@ -81,34 +81,97 @@ const expected =
     0.005877384425022965 0.00466291282312835 0.0028530945920350282 0.0008130106752860425 2.4399432485774198e-5 -9.743480904212476e-5 0.0;
     0.00017108987342944414 7.097261590252536e-5 -7.822316408657793e-5 -0.0001534897546375058 -9.087984332447822e-5 -3.3079379573126077e-5 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 0.0])
+const expected_no_regularity = 
+expected_data(
+   [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
+   [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
+   # Expected phi:
+   [-1.249777922692213, -1.249777922692476],
+   # Expected n_ion:
+   [0.286568430139637, 0.286568430139562],
+   # Expected upar_ion:
+   [0.0, 0.0],
+   # Expected ppar_ion:
+   [0.182900778868619, 0.157103495982094],
+   # Expected pperp_ion
+   [0.143950337301367, 0.156848978744563],
+   # Expected qpar_ion
+   [0.0, 0.0],
+   # Expected v_t_ion
+   [1.046547864034549, 1.046547864034539],
+   # Expected dSdt
+   [0.000000000000000, -0.000000000409521],
+   # Expected f_ion:
+   [0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
+    0.000709315517525 0.000479326544463 0.000267291885601 0.000076580407386 0.000013307679383 0.000001402619088 0.000000000000000 ;
+    0.006729798254906 0.004547723633218 0.002535994801781 0.000726574675523 0.000126259746577 0.000013307679383 0.000000000000000 ;
+    0.038727315040126 0.026170342585036 0.014593642470202 0.004181148571375 0.000726574675400 0.000076580407373 0.000000000000000 ;
+    0.135165490857824 0.091339334982523 0.050934510844468 0.014592981682667 0.002535879973697 0.000267279782810 0.000000000000000 ;
+    0.262731408806302 0.177543188036297 0.099005269067555 0.028365484502288 0.004929182099845 0.000519531971048 0.000000000000000 ;
+    0.000976245858025 0.000659707199564 0.000367879441171 0.000105399224562 0.000018315638889 0.000001930454136 0.000000000000000 ;
+    0.262731408806302 0.177543188036297 0.099005269067555 0.028365484502288 0.004929182099845 0.000519531971048 0.000000000000000 ;
+    0.135165490857824 0.091339334982523 0.050934510844468 0.014592981682667 0.002535879973697 0.000267279782810 0.000000000000000 ;
+    0.038727315040126 0.026170342585036 0.014593642470202 0.004181148571375 0.000726574675400 0.000076580407373 0.000000000000000 ;
+    0.006729798254906 0.004547723633218 0.002535994801781 0.000726574675523 0.000126259746577 0.000013307679383 0.000000000000000 ;
+    0.000709315517525 0.000479326544463 0.000267291885601 0.000076580407386 0.000013307679383 0.000001402619088 0.000000000000000 ;
+    0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;;;
+    0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
+    0.000474751734681 0.000190566474175 -0.000068453615894 -0.000220474853312 -0.000166467621731 -0.000070269966866 0.000000000000000 ;
+    0.008012848857986 0.005724808854515 0.003340228975343 0.000928890127606 -0.000038607624651 -0.000175709447738 0.000000000000000 ;
+    0.034159249923911 0.024646971077049 0.014803095427209 0.004994733815267 0.000871813908241 -0.000279350037847 0.000000000000000 ;
+    0.095764438891876 0.069477778980504 0.042139603659514 0.014773029429542 0.003247498909427 -0.000201036843458 0.000000000000000 ;
+    0.177293203999877 0.129346984538368 0.079203821037167 0.028272781052345 0.006554590511338 0.000118708200311 0.000000000000000 ;
+    0.216600640826244 0.158225686082234 0.097107273632873 0.034812857798807 0.008154333928791 0.000303091030245 0.000000000000000 ;
+    0.177293203999877 0.129346984538368 0.079203821037167 0.028272781052345 0.006554590511338 0.000118708200311 0.000000000000000 ;
+    0.095764438891876 0.069477778980504 0.042139603659514 0.014773029429542 0.003247498909427 -0.000201036843458 0.000000000000000 ;
+    0.034159249923911 0.024646971077049 0.014803095427209 0.004994733815267 0.000871813908241 -0.000279350037847 0.000000000000000 ;
+    0.008012848857986 0.005724808854515 0.003340228975343 0.000928890127606 -0.000038607624651 -0.000175709447738 0.000000000000000 ;
+    0.000474751734681 0.000190566474175 -0.000068453615894 -0.000220474853312 -0.000166467621731 -0.000070269966866 0.000000000000000 ;
+    0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000])
 ###########################################################################################
 # to modify the test, with a new expected f, print the new f using the following commands
 # in an interative Julia REPL. The path is the path to the .dfns file. 
 ########################################################################################## 
 """
+using moment_kinetics.load_data: open_readonly_output_file, load_pdf_data, load_ion_moments_data, load_fields_data
 fid = open_readonly_output_file(path, "dfns")
 f_ion_vpavperpzrst = load_pdf_data(fid)
 f_ion = f_ion_vpavperpzrst[:,:,1,1,1,:]
 ntind = 2
 nvpa = 13  #subject to grid choices
 nvperp = 7 #subject to grid choices
+# pdf
+for k in 1:ntind
+     for i in 1:nvpa-1
+         for j in 1:nvperp-1
+            @printf("%.15f ", f_ion[i,j,k])
+         end
+         @printf("%.15f ", f_ion[i,nvperp,k])
+         print(";\n")
+     end
+     for j in 1:nvperp-1
+       @printf("%.15f ", f_ion[nvpa,j,k])
+     end
+     @printf("%.15f ", f_ion[nvpa,nvperp,k])
+     if k < ntind
+         print(";;;\n")
+     end
+end
+# a moment
+n_ion_zrst, upar_ion_zrst, ppar_ion_zrst, pperp_ion_zrst, qpar_ion_zrst, v_t_ion_zrst, dSdt_zrst = load_ion_moments_data(fid,extended_moments=true)
 for k in 1:ntind
-  for j in 1:nvperp-1
-      for i in 1:nvpa-1
-         @printf("%.15f ", f_ion[i,j,k])
-         print("; ")
-      end
-      @printf("%.15f ", f_ion[nvpa,j,k])
-      print(";;\n")
-  end
-  for i in 1:nvpa-1
-    @printf("%.15f ", f_ion[i,nvperp,k])
-    print("; ")
-  end
-  @printf("%.15f ", f_ion[nvpa,nvperp,k])
-  if k < ntind
-      print(";;;\n")
-  end  
+   @printf("%.15f", n_ion_zrst[1,1,1,k])
+   if k < ntind
+       print(", ")
+   end
+end
+# a field
+phi_zrt, Er_zrt, Ez_zrt = load_fields_data(fid)
+for k in 1:ntind
+   @printf("%.15f", phi_zrt[1,1,k])
+   if k < ntind
+       print(", ")
+   end
 end
 """
 # default inputs for tests
@@ -175,19 +238,21 @@ test_input_gauss_legendre = Dict("run_name" => "gausslegendre_pseudospectral",
 Run a sound-wave test for a single set of parameters
 """
 # Note 'name' should not be shared by any two tests in this file
-function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
+function run_test(test_input, run_name, vperp_bc, expected, rtol, atol, upar_rtol=nothing; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
     # Make a copy to make sure nothing modifies the input Dicts defined in this test
     # script.
     test_input = deepcopy(test_input)
-
+    test_input["vperp_bc"] = vperp_bc
+    
     if upar_rtol === nothing
         upar_rtol = rtol
     end
 
     # Convert keyword arguments to a unique name
+    test_input["run_name"] = run_name
     name = test_input["run_name"]
     if length(args) > 0
         name = string(name, "_", (string(k, "-", v, "_") for (k, v) in args)...)
@@ -207,7 +272,7 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
     #input = test_input
 
     input["run_name"] = name
-
+    print(input)
     # Suppress console output while running
     quietoutput() do
         # run simulation
@@ -296,6 +361,8 @@ function run_test(test_input, rtol, atol, upar_rtol=nothing; args...)
                 ######################################
 
                 @test isapprox(expected.n_ion[tind], n_ion[tind], atol=atol)
+                println(expected.n_ion[tind])
+                println(n_ion[tind])
                 @test isapprox(expected.upar_ion[tind], upar_ion[tind], atol=atol)
                 @test isapprox(expected.ppar_ion[tind], ppar_ion[tind], atol=atol)
                 @test isapprox(expected.pperp_ion[tind], pperp_ion[tind], atol=atol)
@@ -325,7 +392,15 @@ function runtests()
         # GaussLegendre pseudospectral
         # Benchmark data is taken from this run (GaussLegendre)
         @testset "Gauss Legendre base" begin
-            run_test(test_input_gauss_legendre, 1.e-14, 1.0e-14 )
+            run_name = "gausslegendre_pseudospectral"
+            vperp_bc = "zero"
+            run_test(test_input_gauss_legendre, run_name, vperp_bc, expected_base, 1.0e-14, 1.0e-14 )
+        end
+        @testset "Gauss Legendre no enforced regularity condition at vperp = 0" begin
+            run_name = "gausslegendre_pseudospectral_no_regularity"
+            vperp_bc = "zero-no-regularity"
+            run_test(test_input_gauss_legendre, run_name, vperp_bc, expected_no_regularity,
+             1.0e-13, 1.0e-13)
         end
     end
 end

From 6ec3cefcbdce49817a0cf8911a3b1a99ab0fc040 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 14:31:28 +0100
Subject: [PATCH 329/428] Correct test expected values, refactor test inputs to
 make better use of existing functionality.

---
 .../fokker_planck_time_evolution_tests.jl     | 75 +++++++++----------
 1 file changed, 36 insertions(+), 39 deletions(-)

diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index e11563d24..669af32de 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -86,47 +86,47 @@ expected_data(
    [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
    [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
    # Expected phi:
-   [-1.249777922692213, -1.249777922692476],
+   [-1.254259088243025, -1.254259088243286],
    # Expected n_ion:
-   [0.286568430139637, 0.286568430139562],
+   [0.285287142537587, 0.285287142537513],
    # Expected upar_ion:
    [0.0, 0.0],
    # Expected ppar_ion:
-   [0.182900778868619, 0.157103495982094],
+   [0.182220654804438, 0.156448883483764],
    # Expected pperp_ion
-   [0.143950337301367, 0.156848978744563],
+   [0.143306715174515, 0.156192600834786],
    # Expected qpar_ion
    [0.0, 0.0],
    # Expected v_t_ion
-   [1.046547864034549, 1.046547864034539],
+   [1.046701532502699, 1.046701532502689],
    # Expected dSdt
-   [0.000000000000000, -0.000000000409521],
+   [0.000000000000000, -0.000000000865997],
    # Expected f_ion:
    [0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
-    0.000709315517525 0.000479326544463 0.000267291885601 0.000076580407386 0.000013307679383 0.000001402619088 0.000000000000000 ;
-    0.006729798254906 0.004547723633218 0.002535994801781 0.000726574675523 0.000126259746577 0.000013307679383 0.000000000000000 ;
-    0.038727315040126 0.026170342585036 0.014593642470202 0.004181148571375 0.000726574675400 0.000076580407373 0.000000000000000 ;
-    0.135165490857824 0.091339334982523 0.050934510844468 0.014592981682667 0.002535879973697 0.000267279782810 0.000000000000000 ;
-    0.262731408806302 0.177543188036297 0.099005269067555 0.028365484502288 0.004929182099845 0.000519531971048 0.000000000000000 ;
-    0.000976245858025 0.000659707199564 0.000367879441171 0.000105399224562 0.000018315638889 0.000001930454136 0.000000000000000 ;
-    0.262731408806302 0.177543188036297 0.099005269067555 0.028365484502288 0.004929182099845 0.000519531971048 0.000000000000000 ;
-    0.135165490857824 0.091339334982523 0.050934510844468 0.014592981682667 0.002535879973697 0.000267279782810 0.000000000000000 ;
-    0.038727315040126 0.026170342585036 0.014593642470202 0.004181148571375 0.000726574675400 0.000076580407373 0.000000000000000 ;
-    0.006729798254906 0.004547723633218 0.002535994801781 0.000726574675523 0.000126259746577 0.000013307679383 0.000000000000000 ;
-    0.000709315517525 0.000479326544463 0.000267291885601 0.000076580407386 0.000013307679383 0.000001402619088 0.000000000000000 ;
+    0.000706724195475 0.000477575434536 0.000266315395816 0.000076300638379 0.000013259062819 0.000001397494940 0.000000000000000 ;
+    0.006705212475828 0.004531109564814 0.002526730124657 0.000723920301085 0.000125798485463 0.000013259062819 0.000000000000000 ;
+    0.038585833650058 0.026074735222541 0.014540327934436 0.004165873701136 0.000723920300963 0.000076300638366 0.000000000000000 ;
+    0.134671678398729 0.091005636629981 0.050748427134097 0.014539667807029 0.002526615411287 0.000266303305116 0.000000000000000 ;
+    0.261709398369233 0.176852554997681 0.098620144126568 0.028255144359305 0.004910007858078 0.000517511020834 0.000000000000000 ;
+    0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
+    0.261709398369233 0.176852554997682 0.098620144126568 0.028255144359305 0.004910007858078 0.000517511020834 0.000000000000000 ;
+    0.134671678398729 0.091005636629981 0.050748427134097 0.014539667807029 0.002526615411287 0.000266303305116 0.000000000000000 ;
+    0.038585833650058 0.026074735222541 0.014540327934436 0.004165873701136 0.000723920300963 0.000076300638366 0.000000000000000 ;
+    0.006705212475828 0.004531109564814 0.002526730124657 0.000723920301085 0.000125798485463 0.000013259062819 0.000000000000000 ;
+    0.000706724195475 0.000477575434536 0.000266315395816 0.000076300638379 0.000013259062819 0.000001397494940 0.000000000000000 ;
     0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;;;
     0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
-    0.000474751734681 0.000190566474175 -0.000068453615894 -0.000220474853312 -0.000166467621731 -0.000070269966866 0.000000000000000 ;
-    0.008012848857986 0.005724808854515 0.003340228975343 0.000928890127606 -0.000038607624651 -0.000175709447738 0.000000000000000 ;
-    0.034159249923911 0.024646971077049 0.014803095427209 0.004994733815267 0.000871813908241 -0.000279350037847 0.000000000000000 ;
-    0.095764438891876 0.069477778980504 0.042139603659514 0.014773029429542 0.003247498909427 -0.000201036843458 0.000000000000000 ;
-    0.177293203999877 0.129346984538368 0.079203821037167 0.028272781052345 0.006554590511338 0.000118708200311 0.000000000000000 ;
-    0.216600640826244 0.158225686082234 0.097107273632873 0.034812857798807 0.008154333928791 0.000303091030245 0.000000000000000 ;
-    0.177293203999877 0.129346984538368 0.079203821037167 0.028272781052345 0.006554590511338 0.000118708200311 0.000000000000000 ;
-    0.095764438891876 0.069477778980504 0.042139603659514 0.014773029429542 0.003247498909427 -0.000201036843458 0.000000000000000 ;
-    0.034159249923911 0.024646971077049 0.014803095427209 0.004994733815267 0.000871813908241 -0.000279350037847 0.000000000000000 ;
-    0.008012848857986 0.005724808854515 0.003340228975343 0.000928890127606 -0.000038607624651 -0.000175709447738 0.000000000000000 ;
-    0.000474751734681 0.000190566474175 -0.000068453615894 -0.000220474853312 -0.000166467621731 -0.000070269966866 0.000000000000000 ;
+    0.000473874379555 0.000190726186343 -0.000067413540342 -0.000219129923463 -0.000165609965457 -0.000069940808382 0.000000000000000 ;
+    0.007980077110978 0.005701898246888 0.003327438535529 0.000925955494251 -0.000037963012435 -0.000174809935222 0.000000000000000 ;
+    0.034010103035913 0.024541441396039 0.014741892515397 0.004975785949797 0.000869171347288 -0.000277723510529 0.000000000000000 ;
+    0.095320265225675 0.069160937049138 0.041953067849510 0.014711859721633 0.003235165444330 -0.000199356792174 0.000000000000000 ;
+    0.176443375955397 0.128736406466194 0.078839749105573 0.028150317244036 0.006528121645593 0.000119350436884 0.000000000000000 ;
+    0.215552595150568 0.157471466895046 0.096656199821596 0.034660213647059 0.008120924795998 0.000303073103709 0.000000000000000 ;
+    0.176443375955397 0.128736406466194 0.078839749105573 0.028150317244036 0.006528121645593 0.000119350436884 0.000000000000000 ;
+    0.095320265225675 0.069160937049138 0.041953067849510 0.014711859721633 0.003235165444330 -0.000199356792174 0.000000000000000 ;
+    0.034010103035913 0.024541441396039 0.014741892515397 0.004975785949797 0.000869171347288 -0.000277723510529 0.000000000000000 ;
+    0.007980077110978 0.005701898246888 0.003327438535529 0.000925955494251 -0.000037963012435 -0.000174809935222 0.000000000000000 ;
+    0.000473874379555 0.000190726186343 -0.000067413540342 -0.000219129923463 -0.000165609965457 -0.000069940808382 0.000000000000000 ;
     0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000])
 ###########################################################################################
 # to modify the test, with a new expected f, print the new f using the following commands
@@ -191,7 +191,7 @@ test_input_gauss_legendre = Dict("run_name" => "gausslegendre_pseudospectral",
                               "vperp_nelement" => 3,
                               "vperp_L" => 3.0,
                               "vperp_discretization" => "gausslegendre_pseudospectral",
-                              "split_operators" => false,
+                              #"split_operators" => false,
                               "ionization_frequency" => 0.0,
                               "charge_exchange_frequency" => 0.0,
                               "constant_ionization_rate" => false,
@@ -238,21 +238,19 @@ test_input_gauss_legendre = Dict("run_name" => "gausslegendre_pseudospectral",
 Run a sound-wave test for a single set of parameters
 """
 # Note 'name' should not be shared by any two tests in this file
-function run_test(test_input, run_name, vperp_bc, expected, rtol, atol, upar_rtol=nothing; args...)
+function run_test(test_input, expected, rtol, atol, upar_rtol=nothing; args...)
     # by passing keyword arguments to run_test, args becomes a Dict which can be used to
     # update the default inputs
 
     # Make a copy to make sure nothing modifies the input Dicts defined in this test
     # script.
     test_input = deepcopy(test_input)
-    test_input["vperp_bc"] = vperp_bc
     
     if upar_rtol === nothing
         upar_rtol = rtol
     end
 
     # Convert keyword arguments to a unique name
-    test_input["run_name"] = run_name
     name = test_input["run_name"]
     if length(args) > 0
         name = string(name, "_", (string(k, "-", v, "_") for (k, v) in args)...)
@@ -269,10 +267,8 @@ function run_test(test_input, run_name, vperp_bc, expected, rtol, atol, upar_rto
 
     # Update default inputs with values to be changed
     input = merge(test_input, modified_inputs)
-    #input = test_input
 
     input["run_name"] = name
-    print(input)
     # Suppress console output while running
     quietoutput() do
         # run simulation
@@ -361,8 +357,6 @@ function run_test(test_input, run_name, vperp_bc, expected, rtol, atol, upar_rto
                 ######################################
 
                 @test isapprox(expected.n_ion[tind], n_ion[tind], atol=atol)
-                println(expected.n_ion[tind])
-                println(n_ion[tind])
                 @test isapprox(expected.upar_ion[tind], upar_ion[tind], atol=atol)
                 @test isapprox(expected.ppar_ion[tind], ppar_ion[tind], atol=atol)
                 @test isapprox(expected.pperp_ion[tind], pperp_ion[tind], atol=atol)
@@ -394,13 +388,16 @@ function runtests()
         @testset "Gauss Legendre base" begin
             run_name = "gausslegendre_pseudospectral"
             vperp_bc = "zero"
-            run_test(test_input_gauss_legendre, run_name, vperp_bc, expected_base, 1.0e-14, 1.0e-14 )
+            run_test(test_input_gauss_legendre,
+             expected_base, 1.0e-14, 1.0e-14;
+             vperp_bc="zero")
         end
         @testset "Gauss Legendre no enforced regularity condition at vperp = 0" begin
             run_name = "gausslegendre_pseudospectral_no_regularity"
             vperp_bc = "zero-no-regularity"
-            run_test(test_input_gauss_legendre, run_name, vperp_bc, expected_no_regularity,
-             1.0e-13, 1.0e-13)
+            run_test(test_input_gauss_legendre,
+            expected_no_regularity,
+             1.0e-14, 1.0e-14; vperp_bc="zero-no-regularity")
         end
     end
 end

From adc1facafdacf62d5e4256767bea498dccc0c665 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 14:34:15 +0100
Subject: [PATCH 330/428] Change input to match check-in test.

---
 examples/fokker-planck/fokker-planck-relaxation.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/fokker-planck/fokker-planck-relaxation.toml b/examples/fokker-planck/fokker-planck-relaxation.toml
index f2d5651a1..0f633c12b 100644
--- a/examples/fokker-planck/fokker-planck-relaxation.toml
+++ b/examples/fokker-planck/fokker-planck-relaxation.toml
@@ -13,7 +13,7 @@ initial_temperature1 = 1.0
 initial_density2 = 0.5
 initial_temperature2 = 1.0
 z_IC_option1 = "sinusoid"
-z_IC_density_amplitude1 = 0.001
+z_IC_density_amplitude1 = 0.0
 z_IC_density_phase1 = 0.0
 z_IC_upar_amplitude1 = 0.0
 z_IC_upar_phase1 = 0.0

From f2def768aa9d2237229a6ba61325dfae92e528a7 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 14:55:41 +0100
Subject: [PATCH 331/428] Make sure both vperp_bc options are precompiled.

---
 util/precompile_run.jl | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/util/precompile_run.jl b/util/precompile_run.jl
index f5a7fb512..e927edd36 100644
--- a/util/precompile_run.jl
+++ b/util/precompile_run.jl
@@ -66,17 +66,24 @@ for input ∈ [base_input, cheb_input, wall_bc_input, wall_bc_cheb_input]
     push!(inputs_list, x)
 end
 
-collisions_input = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
+collisions_input1 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
                                                   "krook_collisions" => Dict{String,Any}("use_krook" => true),
                                                   "fokker_planck_collisions" => Dict{String,Any}("use_fokker_planck" => true, "self_collisions" => true, "slowing_down_test" => true),
                                                   "vperp_discretization" => "gausslegendre_pseudospectral",
                                                   "vpa_discretization" => "gausslegendre_pseudospectral",
                                                  ))
+collisions_input2 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
+                                                  "krook_collisions" => Dict{String,Any}("use_krook" => true),
+                                                  "fokker_planck_collisions" => Dict{String,Any}("use_fokker_planck" => true, "self_collisions" => true, "slowing_down_test" => true),
+                                                  "vperp_discretization" => "gausslegendre_pseudospectral",
+                                                  "vpa_discretization" => "gausslegendre_pseudospectral",
+                                                  "vperp_bc" => "zero-no-regularity",
+                                                 ))
 # add an additional input for every geometry option available in addition to the default
 geo_input1 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
                                             "geometry" => Dict{String,Any}("option" => "1D-mirror", "DeltaB" => 0.5, "pitch" => 0.5, "rhostar" => 1.0))) 
 
-push!(inputs_list, collisions_input, geo_input1)
+push!(inputs_list, collisions_input1, collisions_input2, geo_input1)
 
 for input in inputs_list
     run_moment_kinetics(input)

From 2afbbae96627056960779bc06cad40aa3187fd00 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 28 Jun 2024 14:58:41 +0100
Subject: [PATCH 332/428] Make conserving correction terms an optional feature
 through "use_conserving_corrections" in the Fokker Planck input namelist.

---
 moment_kinetics/src/fokker_planck.jl | 12 +++++++++---
 moment_kinetics/src/input_structs.jl |  2 ++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/fokker_planck.jl b/moment_kinetics/src/fokker_planck.jl
index 1f520e9f6..fb9c6ef19 100644
--- a/moment_kinetics/src/fokker_planck.jl
+++ b/moment_kinetics/src/fokker_planck.jl
@@ -90,6 +90,7 @@ function setup_fkpl_collisions_input(toml_input::Dict, reference_params)
        nuii = -1.0,
        frequency_option = "reference_parameters",
        self_collisions = true,
+       use_conserving_corrections = true,
        slowing_down_test = false,
        sd_density = 1.0,
        sd_temp = 0.01,
@@ -257,6 +258,7 @@ function explicit_fp_collisions_weak_form_Maxwellian_cross_species!(pdf_out,pdf_
     @boundscheck r.n == size(dSdt,2) || throw(BoundsError(dSdt))
     @boundscheck n_ion_species == size(dSdt,3) || throw(BoundsError(dSdt))
     
+    use_conserving_corrections = collisions.fkpl.use_conserving_corrections
     fkin = collisions.fkpl
     # masses charge numbers and collision frequencies
     mref = 1.0 # generalise if multiple ions evolved
@@ -282,8 +284,10 @@ function explicit_fp_collisions_weak_form_Maxwellian_cross_species!(pdf_out,pdf_
         # enforce the boundary conditions on CC before it is used for timestepping
         enforce_vpavperp_BCs!(fkpl_arrays.CC,vpa,vperp,vpa_spectral,vperp_spectral)
         # make sure that the cross-species terms conserve density
-        density_conserving_correction!(fkpl_arrays.CC, pdf_in[:,:,iz,ir,is], vpa, vperp,
+        if use_conserving_corrections
+            density_conserving_correction!(fkpl_arrays.CC, pdf_in[:,:,iz,ir,is], vpa, vperp,
                                 fkpl_arrays.S_dummy)
+        end
         # advance this part of s,r,z with the resulting sum_s' C[Fs,Fs']
         begin_anyv_vperp_vpa_region()
         CC = fkpl_arrays.CC
@@ -326,6 +330,7 @@ function explicit_fokker_planck_collisions_weak_form!(pdf_out,pdf_in,dSdt,compos
     nuref = collisions.fkpl.nuii # generalise!
     Zi = collisions.fkpl.Zi # generalise!
     nussp = nuref*(Zi^4) # include charge number factor for self collisions
+    use_conserving_corrections = collisions.fkpl.use_conserving_corrections
     # N.B. parallelisation using special 'anyv' region
     begin_s_r_z_anyv_region()
     @loop_s_r_z is ir iz begin
@@ -336,9 +341,10 @@ function explicit_fokker_planck_collisions_weak_form!(pdf_out,pdf_in,dSdt,compos
         # enforce the boundary conditions on CC before it is used for timestepping
         enforce_vpavperp_BCs!(fkpl_arrays.CC,vpa,vperp,vpa_spectral,vperp_spectral)
         # make ad-hoc conserving corrections
-        conserving_corrections!(fkpl_arrays.CC, pdf_in[:,:,iz,ir,is], vpa, vperp,
+        if use_conserving_corrections
+            conserving_corrections!(fkpl_arrays.CC, pdf_in[:,:,iz,ir,is], vpa, vperp,
                                 fkpl_arrays.S_dummy)
-        
+        end
         # advance this part of s,r,z with the resulting C[Fs,Fs]
         begin_anyv_vperp_vpa_region()
         CC = fkpl_arrays.CC
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index fe7661e9a..1c3f18f3f 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -365,6 +365,8 @@ Base.@kwdef struct fkpl_collisions_input
     nuii::mk_float
     # option to determine if self collisions are used (for physics test)
     self_collisions::Bool
+    # option to determine if ad-hoc moment_kinetics-style conserving corrections are used
+    use_conserving_corrections::Bool
     # option to determine if cross-collisions against fixed Maxwellians are used
     slowing_down_test::Bool
     # Setting to switch between different options for Fokker-Planck collision frequency input

From 23e7b696f7921f6d9ee4d41ef6cff749cc4d2cbd Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Sun, 23 Jun 2024 17:59:15 +0100
Subject: [PATCH 333/428] stage maxwell diffusion operator first draft. Works
 for neutrals and ions.

---
 moment_kinetics/src/input_structs.jl         |  18 ++
 moment_kinetics/src/maxwell_diffusion.jl     | 259 +++++++++++++++++++
 moment_kinetics/src/moment_kinetics.jl       |   1 +
 moment_kinetics/src/moment_kinetics_input.jl |  12 +-
 moment_kinetics/src/time_advance.jl          |  22 ++
 5 files changed, 310 insertions(+), 2 deletions(-)
 create mode 100644 moment_kinetics/src/maxwell_diffusion.jl

diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index f81d28a30..6ba339298 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -105,6 +105,8 @@ mutable struct advance_info
     neutral_ionization_collisions_1V::Bool
     ionization_source::Bool
     krook_collisions_ii::Bool
+    mxwl_diff_collisions_ii::Bool
+    mxwl_diff_collisions_nn::Bool
     explicit_weakform_fp_collisions::Bool
     external_source::Bool
     ion_numerical_dissipation::Bool
@@ -359,7 +361,19 @@ struct drive_input
 end
 
 """
+Structs set up for the collision operators so far in use. These will each
+be contained in the main collisions_input struct below, as substructs. 
 """
+Base.@kwdef struct mxwl_diff_collisions_input
+    use_maxwell_diffusion::Bool
+    # different diffusion coefficients for each species, has units of 
+    # frequency * velocity^2. Diffusion coefficients usually denoted D
+    D_ii::mk_float
+    D_nn::mk_float
+    # Setting to switch between different options for Krook collision operator
+    diffusion_coefficient_option::String # "reference_parameters" # "manual", 
+end
+
 Base.@kwdef struct krook_collisions_input
     use_krook::Bool
     # Ion-ion Coulomb collision rate at the reference density and temperature
@@ -400,6 +414,8 @@ Base.@kwdef struct fkpl_collisions_input
 end
 
 """
+Collisions input struct to contain all the different collisions substructs and overall 
+collision input parameters.
 """
 struct collisions_input
     # charge exchange collision frequency
@@ -412,6 +428,8 @@ struct collisions_input
     krook::krook_collisions_input
     # struct of parameters for the Fokker-Planck operator
     fkpl::fkpl_collisions_input
+    # struct of parameters for the Maxwellian Diffusion operator
+    mxwl_diff::mxwl_diff_collisions_input
 end
 
 """
diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
new file mode 100644
index 000000000..7cf57d6be
--- /dev/null
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -0,0 +1,259 @@
+"""
+Model Diffusion operator - this will act to push the distribution towards
+Maxwellian, but purely based on velocity gradients. So need the second derivative of 
+the distribution function. Integration constants should not be a problem, i.e. pushing 
+the second derivatives of a Maxwellian and our function together should also move their 
+values towards each other, as the boundary conditions of both would be the same. 
+
+This operator will mostly come into effect in places where there is ringing in the 
+distribution, which can be expected in the grid points near the walls. Fortunately this
+region is also where the plasma is most collisional, so having such an operator is also
+most valid here.
+"""
+
+module maxwell_diffusion
+
+export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
+
+using ..looping
+using ..input_structs: mxwl_diff_collisions_input, set_defaults_and_check_section!
+using ..calculus: second_derivative!
+using ..reference_parameters: get_reference_collision_frequency_ii
+
+"""
+Function for reading Maxwell diffusion operator input parameters. 
+Structure the namelist as follows.
+
+[maxwell_diffusion_collisions]
+use_maxwell_diffusion = true
+D_ii = 1.0
+diffusion_coefficient_option = "manual"
+"""
+function setup_mxwl_diff_collisions_input(toml_input::Dict, reference_params)
+    # get reference diffusion coefficient, made up of collision frequency and 
+    # thermal speed for now. NOTE THAT THIS CONSTANT PRODUCES ERRORS. DO NOT USE
+    D_ii_mxwl_diff_default = get_reference_collision_frequency_ii(reference_params) *
+                             2 * reference_params.Tref/reference_params.mref
+    D_nn_mxwl_diff_default = D_ii_mxwl_diff_default
+    # read the input toml and specify a sensible default    
+    input_section = set_defaults_and_check_section!(toml_input, "maxwell_diffusion_collisions",
+       # begin default inputs (as kwargs)
+       use_maxwell_diffusion = false,
+       D_ii = -1.0,
+       D_nn = -1.0,
+       diffusion_coefficient_option = "reference_parameters")
+       
+    # ensure that the diffusion coefficient is consistent with the input option
+    diffusion_coefficient_option = input_section["diffusion_coefficient_option"]
+    if diffusion_coefficient_option == "reference_parameters"
+        input_section["D_ii"] = D_ii_mxwl_diff_default
+        input_section["D_nn"] = D_nn_mxwl_diff_default
+    elseif diffusion_coefficient_option == "manual" 
+        # use the diffusion coefficient from the input file
+        # do nothing
+    else
+        error("Invalid option [maxwell_diffusion_collisions] "
+              * "diffusion_coefficient_option=$(diffusion_coefficient_option) passed")
+    end
+    # finally, ensure prefactor < 0 if use_maxwell_diffusion is false
+    # so that prefactor > 0 is the only check required in the rest of the code
+    if !input_section["use_maxwell_diffusion"]
+        input_section["D_ii"] = -1.0
+        input_section["D_nn"] = -1.0
+    end
+    input = Dict(Symbol(k)=>v for (k,v) in input_section)
+
+    return mxwl_diff_collisions_input(; input...)
+end
+
+"""
+Calculate the Maxwellian associated with the current ion pdf moments, and then 
+subtract this from current pdf. Then take second derivative of this function
+to act as the diffusion operator. 
+"""
+function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::T_spectral, 
+                                       dt, diffusion_coefficient) where T_spectral
+    
+    # If negative input (should be -1.0), then none of this diffusion will happen. 
+    # This number can be put in as some parameter in the input file called something
+    # like 'maxwellian_diffusion_coefficient'
+    if diffusion_coefficient <= 0.0 || vpa.n == 1
+        return nothing
+    end
+
+    if vperp.n > 1 && (moments.evolve_density || moments.evolve_upar || moments.evolve_ppar)
+        error("Maxwell diffusion not implemented for 2V moment-kinetic cases yet")
+    end
+
+    # Otherwise, build the maxwellian function (which is going to be subtracted from 
+    # the current distribution) using the moments of the distribution (so that the 
+    # operator itself conserves the moments), and then this result will be the one 
+    # whose second derivative will be added to the RHS (i.e. subtracted from the current)
+    begin_s_r_z_vperp_region()
+
+    # In what follows, there are eight combinations of booleans (though not all have been
+    # fully implemented yet). In line with moment kinetics, the Maxwellian is normalised
+    # in the relevant ways:
+    # - density: normalise by n 
+    # - upar: working in peculiar velocity space, so no upar subtraction from vpa 
+    # - ppar: normalisation by vth, in 1D is 1/vth prefactor, and grid is normalised by vth,
+    # hence no 1/vth^2 term in the exponent.
+    if moments.evolve_density && moments.evolve_upar && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_density && moments.evolve_upar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            1.0 / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_density && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2))
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_upar && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_density
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            1.0 / vth * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_upar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            n = f_in.density[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    elseif moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    else
+        # Drift kinetic version is the only one that currently can support 2V. 
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
+            if vperp.n == 1
+                vth_prefactor = 1.0 / vth
+            else
+                vth_prefactor = 1.0 / vth^3
+            end
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
+                            exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+    end
+    return nothing
+end
+
+"""
+Calculate the Maxwellian associated with the current neutral pdf moments, and then 
+subtract this from current pdf. Then take second derivative of this function
+to act as the diffusion operator. 
+"""
+function neutral_vz_maxwell_diffusion!(f_out, f_in, moments, vzeta, vr, vz, spectral::T_spectral, 
+                                       dt, diffusion_coefficient) where T_spectral
+    
+    # If negative input (should be -1.0), then none of this diffusion will happen. 
+    # This number can be put in as some parameter in the input file called something
+    # like 'maxwellian_diffusion_coefficient'
+    if diffusion_coefficient <= 0.0 || vz.n == 1
+        return nothing
+    end
+
+    if vr.n > 1 && (moments.evolve_density || moments.evolve_upar || moments.evolve_ppar)
+        error("Maxwell diffusion not implemented for 2V moment-kinetic cases yet")
+    end
+
+    # Otherwise, build the maxwellian function (which is going to be subtracted from 
+    # the current distribution) using the moments of the distribution (so that the 
+    # operator itself conserves the moments), and then this result will be the one 
+    # whose second derivative will be added to the RHS (i.e. subtracted from the current pdf)
+    begin_sn_r_z_vzeta_vr_region()
+
+
+    if moments.evolve_ppar && moments.evolve_upar
+        # See similar comments in krook_collisions! function. 
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            @views @. vz.scratch = f_in.pdf_neutral[:,ivr,ivzeta,iz,ir,isn] - 
+                            exp(-((vz.grid[:])^2 + (vr.grid[ivr])^2 + (vzeta.grid[ivzeta])^2) )
+            second_derivative!(vz.scratch2, vz.scratch, vz, spectral)
+            @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch2
+        end
+    elseif moments.evolve_ppar
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            vth = moments.neutral.vth[iz,ir,isn]
+            uz = f_in.uz_neutral[iz,ir,isn]
+            @views @. vz.scratch = f_in.pdf_neutral[:,ivr,ivzeta,iz,ir,isn] - 
+                            exp(- ((vz.grid[:] - uz)^2 + (vr.grid[ivr])^2 + (vzeta.grid[ivzeta])^2)/(vth^2) )
+            second_derivative!(vz.scratch2, vz.scratch, vz, spectral)
+            @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch2
+        end
+    elseif moments.evolve_upar
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            vth = moments.neutral.vth[iz,ir,isn]
+            @views @. vz.scratch = f_in.pdf_neutral[:,ivr,ivzeta,iz,ir,isn] - 
+                            1.0 / vth * exp(- ((vz.grid[:])^2 + (vr.grid[ivr])^2 + (vzeta.grid[ivzeta])^2)/(vth^2) )
+            second_derivative!(vz.scratch2, vz.scratch, vz, spectral)
+            @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch2
+        end
+    elseif moments.evolve_density
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            vth = moments.neutral.vth[iz,ir,isn]
+            uz = f_in.uz_neutral[iz,ir,isn]
+            @views @. vz.scratch = f_in.pdf_neutral[:,ivr,ivzeta,iz,ir,isn] - 
+                            1.0 / vth * exp(- ((vz.grid[:] - uz)^2 + 
+                                            (vr.grid[ivr])^2 + (vzeta.grid[ivzeta])^2)/(vth^2) )
+            second_derivative!(vz.scratch2, vz.scratch, vz, spectral)
+            @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch2
+        end
+    else 
+        @loop_sn_r_z_vzeta_vr isn ir iz ivzeta ivr begin
+            n = f_in.density_neutral[iz,ir,isn]
+            uz = f_in.uz_neutral[iz,ir,isn]
+            vth = moments.neutral.vth[iz,ir,isn]
+            if vr.n == 1 && vzeta.n == 1
+                vth_prefactor = 1.0 / vth
+            else
+                vth_prefactor = 1.0 / vth^3
+            end
+            @views @. vz.scratch = f_in.pdf_neutral[:,ivr,ivzeta,iz,ir,isn] - n * vth_prefactor * 
+                            exp(-( (vz.grid[:] - uz)^2 + (vr.grid[ivr])^2 + (vzeta.grid[ivzeta])^2)/(vth^2) )
+            second_derivative!(vz.scratch2, vz.scratch, vz, spectral)
+            @views @. f_out[:,ivr,ivzeta,iz,ir,isn] += dt * diffusion_coefficient * vz.scratch2
+        end
+    end
+    return nothing
+end
+
+end # maxwell_diffusion
\ No newline at end of file
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 90fd7f1bd..6e6ac2bef 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -59,6 +59,7 @@ include("neutral_vz_advection.jl")
 include("charge_exchange.jl")
 include("ionization.jl")
 include("krook_collisions.jl")
+include("maxwell_diffusion.jl")
 include("continuity.jl")
 include("energy_equation.jl")
 include("force_balance.jl")
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 82b8b00fb..45675996f 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -15,6 +15,7 @@ using ..coordinates: define_coordinate
 using ..external_sources
 using ..file_io: io_has_parallel, input_option_error, open_ascii_output_file
 using ..krook_collisions: setup_krook_collisions_input
+using ..maxwell_diffusion: setup_mxwl_diff_collisions_input
 using ..fokker_planck: setup_fkpl_collisions_input
 using ..finite_differences: fd_check_option
 using ..input_structs
@@ -180,14 +181,21 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     end
     #################### end specification of species inputs #####################
 
+    # Build the main collisions struct using scan_input and reference parameters
     charge_exchange = get(scan_input, "charge_exchange_frequency", 2.0*sqrt(species.ion[1].initial_temperature))
     ionization = get(scan_input, "ionization_frequency", charge_exchange)
     constant_ionization_rate = get(scan_input, "constant_ionization_rate", false)
     # set up krook collision inputs
     krook_input = setup_krook_collisions_input(scan_input, reference_params)
-    # set up krook collision inputs
+    # set up Fokker-Planck collision inputs
     fkpl_input = setup_fkpl_collisions_input(scan_input, reference_params)
-    collisions = collisions_input(charge_exchange, ionization, constant_ionization_rate, krook_input, fkpl_input)
+    # set up maxwell diffusion collision inputs
+    mxwl_diff_input = setup_mxwl_diff_collisions_input(scan_input, reference_params)
+    # write total collision struct using the structs above, as each setup function 
+    # for the collisions outputs itself a struct of the type of collision, which
+    # is a substruct of the overall collisions_input struct.
+    collisions = collisions_input(charge_exchange, ionization, constant_ionization_rate, 
+                                  krook_input, fkpl_input, mxwl_diff_input)
 
     # parameters related to the time stepping
     timestepping_section = set_defaults_and_check_section!(
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 2460dba15..a4fd01a03 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -46,6 +46,7 @@ using ..ionization: ion_ionization_collisions_1V!, neutral_ionization_collisions
                     ion_ionization_collisions_3V!, neutral_ionization_collisions_3V!,
                     constant_ionization_source!
 using ..krook_collisions: krook_collisions!
+using ..maxwell_diffusion: ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
 using ..external_sources
 using ..nonlinear_solvers
 using ..numerical_dissipation: vpa_boundary_buffer_decay!,
@@ -766,6 +767,8 @@ function setup_advance_flags(moments, composition, t_params, collisions,
     advance_neutral_ionization_1V = false
     advance_ionization_source = false
     advance_krook_collisions_ii = false
+    advance_maxwell_diffusion_ii = false
+    advance_maxwell_diffusion_nn = false
     advance_external_source = false
     advance_ion_numerical_dissipation = false
     advance_neutral_numerical_dissipation = false
@@ -850,9 +853,17 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         if collisions.ionization > 0.0 && collisions.constant_ionization_rate && !t_params.implicit_ion_advance
             advance_ionization_source = true
         end
+        # set flags for krook and maxwell diffusion collisions, and negative coefficient
+        # in both cases (as usual) will mean not employing that operator (flag remains false)
         if collisions.krook.nuii0 > 0.0
             advance_krook_collisions_ii = !t_params.implicit_ion_advance
         end
+        if collisions.mxwl_diff.D_ii > 0.0
+            advance_maxwell_diffusion_ii = true
+        end
+        if collisions.mxwl_diff.D_nn > 0.0
+            advance_maxwell_diffusion_nn = true
+        end
         advance_external_source = external_source_settings.ion.active && !t_params.implicit_ion_advance
         advance_neutral_external_source = external_source_settings.neutral.active
         advance_ion_numerical_dissipation = !(t_params.implicit_ion_advance || t_params.implicit_vpa_advection)
@@ -1034,6 +1045,7 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
                         advance_neutral_ionization, advance_ion_ionization_1V,
                         advance_neutral_ionization_1V,
                         advance_ionization_source, advance_krook_collisions_ii,
+                        advance_maxwell_diffusion_ii, advance_maxwell_diffusion_nn,
                         explicit_weakform_fp_collisions,
                         advance_external_source, advance_ion_numerical_dissipation,
                         advance_neutral_numerical_dissipation, advance_sources,
@@ -2573,6 +2585,16 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
         krook_collisions!(fvec_out.pdf, fvec_in, moments, composition, collisions,
                           vperp, vpa, dt)
     end
+    # Add maxwellian diffusion collision operator for ions
+    if advance.mxwl_diff_collisions_ii
+        ion_vpa_maxwell_diffusion!(fvec_out.pdf, fvec_in, moments, vpa, vperp, vpa_spectral, 
+                                   dt, collisions.mxwl_diff.D_ii)
+    end
+    # Add maxwellian diffusion collision operator for neutrals
+    if advance.mxwl_diff_collisions_nn
+        neutral_vz_maxwell_diffusion!(fvec_out.pdf_neutral, fvec_in, moments, vzeta, vr, vz, vz_spectral, 
+                                   dt, collisions.mxwl_diff.D_nn)
+    end
 
     if advance.external_source
         external_ion_source!(fvec_out.pdf, fvec_in, moments, external_source_settings.ion,

From a9d2a2afe87b4bc439d481c9338885c96b23fd1e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 26 Jun 2024 15:52:57 +0100
Subject: [PATCH 334/428] Fix reference collision frequency for Maxwell
 diffusion

Do not need to multiply by 2*Tref/mref=cref^2 as these factors cancel
the cref^(-2) that normalises the velocity space derivatives (or
normalises explicit factors of vth^2 when using a normalised velocity
coordinate in a moment-kinetic mode).
---
 moment_kinetics/src/maxwell_diffusion.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 7cf57d6be..fc93994f2 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -32,8 +32,8 @@ diffusion_coefficient_option = "manual"
 function setup_mxwl_diff_collisions_input(toml_input::Dict, reference_params)
     # get reference diffusion coefficient, made up of collision frequency and 
     # thermal speed for now. NOTE THAT THIS CONSTANT PRODUCES ERRORS. DO NOT USE
-    D_ii_mxwl_diff_default = get_reference_collision_frequency_ii(reference_params) *
-                             2 * reference_params.Tref/reference_params.mref
+    D_ii_mxwl_diff_default = get_reference_collision_frequency_ii(reference_params)# *
+                             #2 * reference_params.Tref/reference_params.mref
     D_nn_mxwl_diff_default = D_ii_mxwl_diff_default
     # read the input toml and specify a sensible default    
     input_section = set_defaults_and_check_section!(toml_input, "maxwell_diffusion_collisions",

From 6700a0d239468085f4b5e526a5ed0d171232d0cc Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 26 Jun 2024 16:18:06 +0100
Subject: [PATCH 335/428] Fix "energy" option for external_sources

---
 moment_kinetics/src/external_sources.jl | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 08b23bd5f..95f40ea7c 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -745,12 +745,15 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
 
     is = 1
     ion_moments = moments.ion
+    density = fvec_in.density
+    upar = fvec_in.upar
+    ppar = fvec_in.ppar
 
     if ion_source_settings.source_type == "Maxwellian"
         if moments.evolve_ppar
             @loop_r_z ir iz begin
                 ion_moments.external_source_pressure_amplitude[iz,ir] =
-                    (0.5 * ion_source_settings.source_T + fvec_in.upar[iz,ir,is]^2) *
+                    (0.5 * ion_source_settings.source_T + upar[iz,ir,is]^2) *
                     ion_moments.external_source_amplitude[iz,ir]
             end
         end
@@ -758,7 +761,7 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
         if moments.evolve_upar
             @loop_r_z ir iz begin
                 ion_moments.external_source_momentum_amplitude[iz,ir] =
-                    - ion_moments.density[iz,ir] * ion_moments.upar[iz,ir] *
+                    - density[iz,ir] * upar[iz,ir] *
                       ion_source_settings.source_strength *
                       ion_source_settings.r_amplitude[ir] *
                       ion_source_settings.z_amplitude[iz]
@@ -767,8 +770,7 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
         if moments.evolve_ppar
             @loop_r_z ir iz begin
                 ion_moments.external_source_pressure_amplitude[iz,ir] =
-                    (0.5 * ion_source_settings.source_T + ion_moments.upar[iz,ir]^2 -
-                     ion_moments.ppar[iz,ir]) *
+                    (0.5 * ion_source_settings.source_T + upar[iz,ir]^2 - ppar[iz,ir]) *
                     ion_source_settings.source_strength *
                     ion_source_settings.r_amplitude[ir] *
                     ion_source_settings.z_amplitude[iz]
@@ -784,8 +786,8 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
                     ion_source_settings.PI_density_target_iz !== nothing
                 # This process has the target point
 
-                n_mid = fvec_in.density[ion_source_settings.PI_density_target_iz,
-                                        ion_source_settings.PI_density_target_ir, is]
+                n_mid = density[ion_source_settings.PI_density_target_iz,
+                                ion_source_settings.PI_density_target_ir, is]
                 n_error = ion_source_settings.PI_density_target - n_mid
 
                 ion_moments.external_source_controller_integral[1,1] +=
@@ -820,14 +822,13 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
         if moments.evolve_ppar
             @loop_r_z ir iz begin
                 ion_moments.external_source_pressure_amplitude[iz,ir] =
-                    (0.5 * ion_source_settings.source_T + fvec_in.upar[iz,ir,is]^2) *
+                    (0.5 * ion_source_settings.source_T + upar[iz,ir,is]^2) *
                     amplitude * ion_source_settings.controller_source_profile[iz,ir]
             end
         end
     elseif ion_source_settings.source_type == "density_profile_control"
         begin_r_z_region()
 
-        density = fvec_in.density
         target = ion_source_settings.PI_density_target
         P = ion_source_settings.PI_density_controller_P
         I = ion_source_settings.PI_density_controller_I
@@ -847,7 +848,7 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
         if moments.evolve_ppar
             @loop_r_z ir iz begin
                 ion_moments.external_source_pressure_amplitude[iz,ir] =
-                    (0.5 * ion_source_settings.source_T + fvec_in.upar[iz,ir,is]^2) *
+                    (0.5 * ion_source_settings.source_T + upar[iz,ir,is]^2) *
                     amplitude[iz,ir]
             end
         end

From 4f5f581f3169da2f73893c68ea38be67e61df9d1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 26 Jun 2024 16:18:24 +0100
Subject: [PATCH 336/428] Disable neutral 'Maxwell diffusion' when using
 "reference_parameters"

Neutral collisions would have a different parameter dependence. Unsure
what that would be, so just disable for now.
---
 moment_kinetics/src/maxwell_diffusion.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index fc93994f2..03b4daa96 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -47,7 +47,7 @@ function setup_mxwl_diff_collisions_input(toml_input::Dict, reference_params)
     diffusion_coefficient_option = input_section["diffusion_coefficient_option"]
     if diffusion_coefficient_option == "reference_parameters"
         input_section["D_ii"] = D_ii_mxwl_diff_default
-        input_section["D_nn"] = D_nn_mxwl_diff_default
+        input_section["D_nn"] = -1.0 #D_nn_mxwl_diff_default
     elseif diffusion_coefficient_option == "manual" 
         # use the diffusion coefficient from the input file
         # do nothing

From 009e7382e0c60d8907e1bb6e679d0ad8715f7185 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 16:57:28 +0100
Subject: [PATCH 337/428] Fix v-space bc for Maxwell-diffusion

Need to set vpa_diffusion and vz_diffusion advance-flags when using
'Maxwell diffusion'.
---
 moment_kinetics/src/time_advance.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index a4fd01a03..997b110da 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -911,9 +911,9 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         # flag to determine if a d^2/dvpa^2 operator is present
         # When using implicit_vpa_advection, the vpa diffusion is included in the implicit
         # step
-        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
+        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1) || advance_maxwell_diffusion_ii)
         vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
-        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0)
+        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0 || advance_maxwell_diffusion_nn)
     end
 
     manufactured_solns_test = manufactured_solns_input.use_for_advance

From c7eed6e7e974e1872879fffa55401fa3de125ba5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 16:58:52 +0100
Subject: [PATCH 338/428] Attempted fix - use scratch3 instead of scratch in
 second_derivative!()

...implementation for weak_discretization_info. Makes the choice of
'scratch*' buffer consistent with the chebyshev_pseudospectral version
of second_derivative!(), so easier to use the right 'scratch*' buffer in
calling functions.
---
 moment_kinetics/src/calculus.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl
index d2c16a478..0065ca196 100644
--- a/moment_kinetics/src/calculus.jl
+++ b/moment_kinetics/src/calculus.jl
@@ -148,7 +148,7 @@ function second_derivative!(d2f, f, coord, spectral::weak_discretization_info; h
     # g = d^2 f / d coord^2, which is 
     # M * g = K * f, with M the mass matrix and K an appropriate stiffness matrix
     # by multiplying by basis functions and integrating by parts    
-    mul!(coord.scratch, spectral.K_matrix, f)
+    mul!(coord.scratch3, spectral.K_matrix, f)
 
     if handle_periodic && coord.bc == "periodic"
         if coord.nrank > 1
@@ -156,8 +156,8 @@ function second_derivative!(d2f, f, coord, spectral::weak_discretization_info; h
                   * "distributed coordinate")
         end
 
-        coord.scratch[1] = 0.5 * (coord.scratch[1] + coord.scratch[end])
-        coord.scratch[end] = coord.scratch[1]
+        coord.scratch3[1] = 0.5 * (coord.scratch3[1] + coord.scratch3[end])
+        coord.scratch3[end] = coord.scratch3[1]
     end
 
     # solve weak form matrix problem M * g = K * f to obtain g = d^2 f / d coord^2
@@ -165,7 +165,7 @@ function second_derivative!(d2f, f, coord, spectral::weak_discretization_info; h
         error("mass_matrix_solve!() does not support a "
               * "distributed coordinate")
     end
-    mass_matrix_solve!(d2f, coord.scratch, spectral)
+    mass_matrix_solve!(d2f, coord.scratch3, spectral)
 end
 
 function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info)
@@ -174,7 +174,7 @@ function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info
     # M * g = K * f, with M the mass matrix, and K an appropriate stiffness matrix,
     # by multiplying by basis functions and integrating by parts.
     # for all other coord.name, do exactly the same as second_derivative! above.
-    mul!(coord.scratch, spectral.L_matrix, f)
+    mul!(coord.scratch3, spectral.L_matrix, f)
 
     if handle_periodic && coord.bc == "periodic"
         if coord.nrank > 1
@@ -182,8 +182,8 @@ function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info
                   * "distributed coordinate")
         end
 
-        coord.scratch[1] = 0.5 * (coord.scratch[1] + coord.scratch[end])
-        coord.scratch[end] = coord.scratch[1]
+        coord.scratch3[1] = 0.5 * (coord.scratch3[1] + coord.scratch3[end])
+        coord.scratch3[end] = coord.scratch3[1]
     end
 
     # solve weak form matrix problem M * g = K * f to obtain g = d^2 f / d coord^2
@@ -191,7 +191,7 @@ function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info
         error("mass_matrix_solve!() does not support a "
               * "distributed coordinate")
     end
-    mass_matrix_solve!(d2f, coord.scratch, spectral)
+    mass_matrix_solve!(d2f, coord.scratch3, spectral)
 end
 
 """

From 664c390c2bc75a45ba422c2f2cde00099167fb37 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 17:00:45 +0100
Subject: [PATCH 339/428] Fix nspecies when loading ion or neutral distribution
 functions

---
 moment_kinetics/src/load_data.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 430edb5a9..cff04edfb 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -2976,7 +2976,6 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
             !isa(iz, mk_int) && push!(dims, nz)
             !isa(ir, mk_int) && push!(dims, nr)
             if is === (:)
-                nspecies = size(variable[1], 3)
                 push!(dims, nspecies)
             elseif !isa(is, mk_int)
                 push!(dims, nspecies)
@@ -2993,7 +2992,6 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
             !isa(iz, mk_int) && push!(dims, nz)
             !isa(ir, mk_int) && push!(dims, nr)
             if is === (:)
-                nspecies = size(variable[1], 3)
                 push!(dims, nspecies)
             elseif !isa(is, mk_int)
                 push!(dims, nspecies)

From 8aa7e1e986d94fee81d520644c5a68cac41f8010 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 17:01:58 +0100
Subject: [PATCH 340/428] Hacky conversion to Mousseau's 'collision operator'

---
 moment_kinetics/src/maxwell_diffusion.jl | 44 ++++++++++++++++--------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 03b4daa96..69ede0588 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -17,7 +17,7 @@ export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_
 
 using ..looping
 using ..input_structs: mxwl_diff_collisions_input, set_defaults_and_check_section!
-using ..calculus: second_derivative!
+using ..calculus: derivative!, second_derivative!
 using ..reference_parameters: get_reference_collision_frequency_ii
 
 """
@@ -100,10 +100,17 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
     # hence no 1/vth^2 term in the exponent.
     if moments.evolve_density && moments.evolve_upar && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+            #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
+            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
+            @views @. vpa.scratch = (vpa.grid - upar / vth) * f_in.pdf[:,ivperp,iz,ir,is]
+            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
         end
     elseif moments.evolve_density && moments.evolve_upar
         @loop_s_r_z_vperp is ir iz ivperp begin
@@ -113,6 +120,7 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     elseif moments.evolve_density && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             upar = f_in.upar[iz,ir,is]
@@ -121,6 +129,7 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     elseif moments.evolve_upar && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
@@ -129,6 +138,7 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     elseif moments.evolve_density
         @loop_s_r_z_vperp is ir iz ivperp begin
             vth = moments.ion.vth[iz,ir,is]
@@ -138,6 +148,7 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     elseif moments.evolve_upar
         @loop_s_r_z_vperp is ir iz ivperp begin
             vth = moments.ion.vth[iz,ir,is]
@@ -147,6 +158,7 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     elseif moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
@@ -156,21 +168,25 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
+        error("hack not implemented for this case")
     else
         # Drift kinetic version is the only one that currently can support 2V. 
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
             upar = f_in.upar[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            if vperp.n == 1
-                vth_prefactor = 1.0 / vth
-            else
-                vth_prefactor = 1.0 / vth^3
-            end
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
-                            exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+            #if vperp.n == 1
+            #    vth_prefactor = 1.0 / vth
+            #else
+            #    vth_prefactor = 1.0 / vth^3
+            #end
+            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
+            #                exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
+            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
+            @views @. vpa.scratch = (vpa.grid - upar) * f_in.pdf[:,ivperp,iz,ir,is]
+            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
         end
     end
     return nothing

From aefc2787142a70311a2fd3aa3facb2a4c2957da6 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 21:11:43 +0100
Subject: [PATCH 341/428] Implicit solve option for Maxwell diffusion

---
 moment_kinetics/src/coordinates.jl           |  13 +-
 moment_kinetics/src/input_structs.jl         |   1 +
 moment_kinetics/src/maxwell_diffusion.jl     | 343 ++++++++++++++-----
 moment_kinetics/src/moment_kinetics_input.jl |   1 +
 moment_kinetics/src/time_advance.jl          |  60 +++-
 5 files changed, 311 insertions(+), 107 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index d034a2dd0..70f5096ae 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -85,6 +85,10 @@ struct coordinate{T <: AbstractVector{mk_float}}
     scratch6::Array{mk_float,1}
     # scratch7 is an array used for intermediate calculations requiring n entries
     scratch7::Array{mk_float,1}
+    # scratch8 is an array used for intermediate calculations requiring n entries
+    scratch8::Array{mk_float,1}
+    # scratch9 is an array used for intermediate calculations requiring n entries
+    scratch9::Array{mk_float,1}
     # scratch_shared is a shared-memory array used for intermediate calculations requiring
     # n entries
     scratch_shared::T
@@ -231,10 +235,11 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch),
         copy(scratch), copy(scratch), copy(scratch), copy(scratch), copy(scratch),
-        scratch_shared, scratch_shared2, scratch_2d, copy(scratch_2d), advection,
-        send_buffer, receive_buffer, input.comm, local_io_range, global_io_range,
-        element_scale, element_shift, input.element_spacing_option, element_boundaries,
-        radau_first_element, other_nodes, one_over_denominator)
+        copy(scratch), copy(scratch), scratch_shared, scratch_shared2, scratch_2d,
+        copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
+        local_io_range, global_io_range, element_scale, element_shift,
+        input.element_spacing_option, element_boundaries, radau_first_element,
+        other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 6ba339298..b982bd4e4 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -76,6 +76,7 @@ struct time_info{Terrorsum <: Real, Trkimp, Timpzero}
     maximum_dt::mk_float
     implicit_ion_advance::Bool
     implicit_vpa_advection::Bool
+    implicit_ion_maxwell_diffusion::Bool
     write_after_fixed_step_count::Bool
     error_sum_zero::Terrorsum
     split_operators::Bool
diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 69ede0588..0b67c142c 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -13,12 +13,17 @@ most valid here.
 
 module maxwell_diffusion
 
-export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
+export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!, implicit_ion_maxwell_diffusion!
 
 using ..looping
 using ..input_structs: mxwl_diff_collisions_input, set_defaults_and_check_section!
+using ..boundary_conditions: enforce_v_boundary_condition_local!, vpagrid_to_dzdt
 using ..calculus: derivative!, second_derivative!
+using ..moment_constraints: moment_constraints_on_residual!
+using ..moment_kinetics_structs: scratch_pdf
+using ..nonlinear_solvers: newton_solve!
 using ..reference_parameters: get_reference_collision_frequency_ii
+using ..velocity_moments: update_derived_moments!, calculate_ion_moment_derivatives!
 
 """
 Function for reading Maxwell diffusion operator input parameters. 
@@ -66,13 +71,31 @@ function setup_mxwl_diff_collisions_input(toml_input::Dict, reference_params)
     return mxwl_diff_collisions_input(; input...)
 end
 
+function ion_vpa_maxwell_diffusion_inner!(f_out, f_in, n, upar, vth, vpa, spectral,
+                                          diffusion_coefficient, dt, ::Val{true},
+                                          ::Val{true}, ::Val{true})
+    second_derivative!(vpa.scratch2, f_in, vpa, spectral)
+    @. vpa.scratch = vpa.grid * f_in
+    derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+    @. f_out += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
+end
+
+function ion_vpa_maxwell_diffusion_inner!(f_out, f_in, n, upar, vth, vpa, spectral,
+                                          diffusion_coefficient, dt, ::Val{false},
+                                          ::Val{false}, ::Val{false})
+    second_derivative!(vpa.scratch2, f_in, vpa, spectral)
+    @. vpa.scratch = (vpa.grid - upar) * f_in
+    derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+    @. f_out += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
+end
+
 """
 Calculate the Maxwellian associated with the current ion pdf moments, and then 
 subtract this from current pdf. Then take second derivative of this function
 to act as the diffusion operator. 
 """
-function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::T_spectral, 
-                                       dt, diffusion_coefficient) where T_spectral
+function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral, dt,
+                                    diffusion_coefficient)
     
     # If negative input (should be -1.0), then none of this diffusion will happen. 
     # This number can be put in as some parameter in the input file called something
@@ -85,6 +108,10 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
         error("Maxwell diffusion not implemented for 2V moment-kinetic cases yet")
     end
 
+    evolve_density = Val(moments.evolve_density)
+    evolve_upar = Val(moments.evolve_upar)
+    evolve_ppar = Val(moments.evolve_ppar)
+
     # Otherwise, build the maxwellian function (which is going to be subtracted from 
     # the current distribution) using the moments of the distribution (so that the 
     # operator itself conserves the moments), and then this result will be the one 
@@ -98,97 +125,76 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
     # - upar: working in peculiar velocity space, so no upar subtraction from vpa 
     # - ppar: normalisation by vth, in 1D is 1/vth prefactor, and grid is normalised by vth,
     # hence no 1/vth^2 term in the exponent.
-    if moments.evolve_density && moments.evolve_upar && moments.evolve_ppar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-            #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-            n = f_in.density[iz,ir,is]
-            upar = f_in.upar[iz,ir,is]
-            vth = moments.ion.vth[iz,ir,is]
-            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
-            @views @. vpa.scratch = (vpa.grid - upar / vth) * f_in.pdf[:,ivperp,iz,ir,is]
-            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
-        end
-    elseif moments.evolve_density && moments.evolve_upar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            vth = moments.ion.vth[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            1.0 / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    elseif moments.evolve_density && moments.evolve_ppar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            upar = f_in.upar[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2))
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    elseif moments.evolve_upar && moments.evolve_ppar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            n = f_in.density[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            n * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    elseif moments.evolve_density
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            vth = moments.ion.vth[iz,ir,is]
-            upar = f_in.upar[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            1.0 / vth * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    elseif moments.evolve_upar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            vth = moments.ion.vth[iz,ir,is]
-            n = f_in.density[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            n / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    elseif moments.evolve_ppar
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            n = f_in.density[iz,ir,is]
-            upar = f_in.upar[iz,ir,is]
-            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-                            n * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2) )
-            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        end
-        error("hack not implemented for this case")
-    else
-        # Drift kinetic version is the only one that currently can support 2V. 
-        @loop_s_r_z_vperp is ir iz ivperp begin
-            n = f_in.density[iz,ir,is]
-            upar = f_in.upar[iz,ir,is]
-            vth = moments.ion.vth[iz,ir,is]
-            #if vperp.n == 1
-            #    vth_prefactor = 1.0 / vth
-            #else
-            #    vth_prefactor = 1.0 / vth^3
-            #end
-            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
-            #                exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
-            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
-            @views @. vpa.scratch = (vpa.grid - upar) * f_in.pdf[:,ivperp,iz,ir,is]
-            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
-        end
+    @loop_s_r_z_vperp is ir iz ivperp begin
+        #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+        #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+        #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+        #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        @views ion_vpa_maxwell_diffusion_inner!(f_out[:,ivperp,iz,ir,is],
+                                                f_in.pdf[:,ivperp,iz,ir,is],
+                                                f_in.density[iz,ir,is],
+                                                f_in.upar[iz,ir,is],
+                                                moments.ion.vth[iz,ir,is], vpa, spectral,
+                                                diffusion_coefficient, dt, evolve_density,
+                                                evolve_upar, evolve_ppar)
     end
+    #elseif moments.evolve_density && moments.evolve_upar
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        vth = moments.ion.vth[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        1.0 / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
+    #elseif moments.evolve_density && moments.evolve_ppar
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        upar = f_in.upar[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2))
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
+    #elseif moments.evolve_upar && moments.evolve_ppar
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        n = f_in.density[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        n * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
+    #elseif moments.evolve_density
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        vth = moments.ion.vth[iz,ir,is]
+    #        upar = f_in.upar[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        1.0 / vth * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
+    #elseif moments.evolve_upar
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        vth = moments.ion.vth[iz,ir,is]
+    #        n = f_in.density[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        n / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
+    #elseif moments.evolve_ppar
+    #    @loop_s_r_z_vperp is ir iz ivperp begin
+    #        n = f_in.density[iz,ir,is]
+    #        upar = f_in.upar[iz,ir,is]
+    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+    #                        n * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2) )
+    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+    #    end
+    #    error("hack not implemented for this case")
     return nothing
 end
 
@@ -272,4 +278,153 @@ function neutral_vz_maxwell_diffusion!(f_out, f_in, moments, vzeta, vr, vz, spec
     return nothing
 end
 
-end # maxwell_diffusion
\ No newline at end of file
+"""
+"""
+function implicit_ion_maxwell_diffusion!(f_out, fvec_in, moments, z_advect, vpa, vperp, z,
+                                         r, dt, r_spectral, vpa_spectral, composition,
+                                         collisions, geometry, nl_solver_params, gyroavs,
+                                         scratch_dummy)
+    if vperp.n > 1 && (moments.evolve_density || moments.evolve_upar || moments.evolve_ppar)
+        error("Moment constraints in implicit_maxwell_diffusion!() do not support 2V runs yet")
+    end
+
+    # Ensure moments are consistent with f_new
+    new_scratch = scratch_pdf(f_out, fvec_in.density, fvec_in.upar, fvec_in.ppar,
+                              fvec_in.pperp, fvec_in.temp_z_s, fvec_in.pdf_neutral,
+                              fvec_in.density_neutral, fvec_in.uz_neutral,
+                              fvec_in.pz_neutral)
+    update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
+                            r_spectral, geometry, gyroavs, scratch_dummy, z_advect, false)
+
+    begin_s_r_z_vperp_region()
+
+    evolve_density = Val(moments.evolve_density)
+    evolve_upar = Val(moments.evolve_upar)
+    evolve_ppar = Val(moments.evolve_ppar)
+    coords = (vpa=vpa,)
+    vpa_bc = vpa.bc
+    D_ii = collisions.mxwl_diff.D_ii
+    zero = 1.0e-14
+    @loop_s is begin
+        @loop_r_z_vperp ir iz ivperp begin
+            f_old_no_bc = @view fvec_in.pdf[:,ivperp,iz,ir,is]
+            this_f_out = @view f_out[:,ivperp,iz,ir,is]
+            n = fvec_in.density[iz,ir,is]
+            upar = fvec_in.upar[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
+
+            if z.irank == 0 && iz == 1
+                @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, moments.ion.vth[iz,ir,is],
+                                                 fvec_in.upar[iz,ir,is],
+                                                 moments.evolve_ppar,
+                                                 moments.evolve_upar)
+                icut_lower_z = vpa.n
+                for ivpa ∈ vpa.n:-1:1
+                    # for left boundary in zed (z = -Lz/2), want
+                    # f(z=-Lz/2, v_parallel > 0) = 0
+                    if vpa.scratch[ivpa] ≤ zero
+                        icut_lower_z = ivpa + 1
+                        break
+                    end
+                end
+            end
+            if z.irank == z.nrank - 1 && iz == z.n
+                @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, moments.ion.vth[iz,ir,is],
+                                                 fvec_in.upar[iz,ir,is],
+                                                 moments.evolve_ppar,
+                                                 moments.evolve_upar)
+                icut_upper_z = 0
+                for ivpa ∈ 1:vpa.n
+                    # for right boundary in zed (z = Lz/2), want
+                    # f(z=Lz/2, v_parallel < 0) = 0
+                    if vpa.scratch[ivpa] ≥ -zero
+                        icut_upper_z = ivpa - 1
+                        break
+                    end
+                end
+            end
+
+            function apply_bc!(x)
+                # Boundary condition
+                enforce_v_boundary_condition_local!(x, vpa_bc, nothing, true,
+                                                    vpa, vpa_spectral)
+
+                if z.bc == "wall"
+                    # Wall boundary conditions. Note that as density, upar, ppar do not
+                    # change in this implicit step, f_new, f_old, and residual should all
+                    # be zero at exactly the same set of grid points, so it is reasonable
+                    # to zero-out `residual` to impose the boundary condition. We impose
+                    # this after subtracting f_old in case rounding errors, etc. mean that
+                    # at some point f_old had a different boundary condition cut-off
+                    # index.
+                    if z.irank == 0 && iz == 1
+                        x[icut_lower_z:end] .= 0.0
+                    end
+                    # absolute velocity at right boundary
+                    if z.irank == z.nrank - 1 && iz == z.n
+                        x[1:icut_upper_z] .= 0.0
+                    end
+                end
+            end
+
+            # Need to apply 'new' boundary conditions to `f_old`, so that by imposing them
+            # on `residual`, they are automatically imposed on `f_new`.
+            f_old = vpa.scratch9 .= f_old_no_bc
+            apply_bc!(f_old)
+
+            left_preconditioner = identity
+            right_preconditioner = identity
+
+            # Define a function whose input is `f_new`, so that when it's output
+            # `residual` is zero, f_new is the result of a backward-Euler timestep:
+            #   (f_new - f_old) / dt = RHS(f_new)
+            # ⇒ f_new - f_old - dt*RHS(f_new) = 0
+            function residual_func!(residual, f_new)
+                apply_bc!(f_new)
+                residual .= f_old
+                ion_vpa_maxwell_diffusion_inner!(residual, f_new, n, upar, vth, vpa,
+                                                 vpa_spectral, D_ii, dt, evolve_density,
+                                                 evolve_upar, evolve_ppar)
+
+                # Now
+                #   residual = f_old + dt*RHS(f_new)
+                # so update to desired residual
+                @. residual = f_new - residual
+
+                apply_bc!(residual)
+                moment_constraints_on_residual!(residual, f_new, moments, vpa)
+            end
+
+            # Buffers
+            # Note vpa.scratch, vpa.scratch2 and vpa.scratch3 are used by advance_f!, so
+            # we cannot use it here.
+            residual = vpa.scratch4
+            delta_x = vpa.scratch5
+            rhs_delta = vpa.scratch6
+            v = vpa.scratch7
+            w = vpa.scratch8
+
+            # Use forward-Euler step for initial guess
+            # By passing this_f_out, which is equal to f_old at this point, the 'residual'
+            # is
+            #   f_new - f_old - dt*RHS(f_old) = -dt*RHS(f_old)
+            # so to get a forward-Euler step we have to subtract this 'residual'
+            residual_func!(residual, this_f_out)
+            this_f_out .-= residual
+
+            success = newton_solve!(this_f_out, residual_func!, residual, delta_x,
+                                    rhs_delta, v, w, nl_solver_params, coords=coords,
+                                    left_preconditioner=left_preconditioner,
+                                    right_preconditioner=right_preconditioner)
+            if !success
+                return success
+            end
+        end
+    end
+
+    nl_solver_params.stage_counter[] += 1
+
+    return true
+end
+
+end # maxwell_diffusion
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 45675996f..5222f8c6f 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -221,6 +221,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         maximum_dt=Inf,
         implicit_ion_advance=true,
         implicit_vpa_advection=false,
+        implicit_ion_maxwell_diffusion=false,
         write_after_fixed_step_count=false,
         high_precision_error_sum=false,
        )
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 997b110da..6054222e5 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -46,7 +46,8 @@ using ..ionization: ion_ionization_collisions_1V!, neutral_ionization_collisions
                     ion_ionization_collisions_3V!, neutral_ionization_collisions_3V!,
                     constant_ionization_source!
 using ..krook_collisions: krook_collisions!
-using ..maxwell_diffusion: ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
+using ..maxwell_diffusion: ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!,
+                           implicit_ion_maxwell_diffusion!
 using ..external_sources
 using ..nonlinear_solvers
 using ..numerical_dissipation: vpa_boundary_buffer_decay!,
@@ -310,6 +311,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         # Not an IMEX scheme, so cannot have any implicit terms
         t_input["implicit_ion_advance"] = false
         t_input["implicit_vpa_advection"] = false
+        t_input["implicit_ion_maxwell_diffusion"] = false
     end
 
     if t_input["high_precision_error_sum"]
@@ -329,6 +331,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      t_input["last_fail_proximity_factor"], t_input["minimum_dt"],
                      t_input["maximum_dt"], t_input["implicit_ion_advance"],
                      t_input["implicit_vpa_advection"],
+                     t_input["implicit_ion_maxwell_diffusion"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
                      t_input["split_operators"], t_input["steady_state_residual"],
                      t_input["converged_residual_value"],
@@ -498,13 +501,36 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     else
         nl_solver_vpa_advection_params = nothing
     end
+    if t_params.implicit_ion_maxwell_diffusion
+        # Implicit solve for ion maxwell diffusion term should be done in serial, as it
+        # will be called within a parallelised s_r_z_vperp loop.
+        nl_solver_maxwell_diffusion_params =
+            setup_nonlinear_solve(input_dict, (vpa=vpa,),
+                                  (composition.n_ion_species, r, z, vperp);
+                                  default_rtol=t_params.rtol / 10.0,
+                                  default_atol=t_params.atol / 10.0,
+                                  serial_solve=true, preconditioner_type="lu")
+    else
+        nl_solver_maxwell_diffusion_params = nothing
+    end
     if nl_solver_ion_advance_params !== nothing &&
             nl_solver_vpa_advection_params !== nothing
         error("Cannot use implicit_ion_advance and implicit_vpa_advection at the same "
               * "time")
     end
+    if nl_solver_ion_advance_params !== nothing &&
+            nl_solver_maxwell_diffusion_params !== nothing
+        error("Cannot use implicit_ion_advance and implicit_ion_maxwell_diffusion at the same "
+              * "time")
+    end
+    if nl_solver_vpa_advection_params !== nothing &&
+            nl_solver_maxwell_diffusion_params !== nothing
+        error("Cannot use implicit_vpa_advection and implicit_ion_maxwell_diffusion at the same "
+              * "time")
+    end
     nl_solver_params = (ion_advance=nl_solver_ion_advance_params,
-                        vpa_advection=nl_solver_vpa_advection_params,)
+                        vpa_advection=nl_solver_vpa_advection_params,
+                        maxwell_diffusion=nl_solver_maxwell_diffusion_params,)
 
     begin_serial_region()
 
@@ -858,7 +884,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         if collisions.krook.nuii0 > 0.0
             advance_krook_collisions_ii = !t_params.implicit_ion_advance
         end
-        if collisions.mxwl_diff.D_ii > 0.0
+        if collisions.mxwl_diff.D_ii > 0.0 && !t_params.implicit_ion_maxwell_diffusion
             advance_maxwell_diffusion_ii = true
         end
         if collisions.mxwl_diff.D_nn > 0.0
@@ -911,9 +937,9 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         # flag to determine if a d^2/dvpa^2 operator is present
         # When using implicit_vpa_advection, the vpa diffusion is included in the implicit
         # step
-        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1) || advance_maxwell_diffusion_ii)
+        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1) || collisions.mxwl_diff.D_ii > 0.0)
         vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
-        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0 || advance_maxwell_diffusion_nn)
+        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0 || collisions.mxwl_diff.D_nn > 0.0)
     end
 
     manufactured_solns_test = manufactured_solns_input.use_for_advance
@@ -923,8 +949,9 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                         advance_neutral_vz_advection, advance_ion_cx, advance_neutral_cx,
                         advance_ion_cx_1V, advance_neutral_cx_1V, advance_ion_ionization,
                         advance_neutral_ionization, advance_ion_ionization_1V,
-                        advance_neutral_ionization_1V, advance_ionization_source,
-                        advance_krook_collisions_ii,
+                        advance_neutral_ionization_1V,
+                        advance_ionization_source, advance_krook_collisions_ii,
+                        advance_maxwell_diffusion_ii, advance_maxwell_diffusion_nn,
                         explicit_weakform_fp_collisions,
                         advance_external_source, advance_ion_numerical_dissipation,
                         advance_neutral_numerical_dissipation, advance_sources,
@@ -959,6 +986,8 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     advance_neutral_ionization_1V = false
     advance_ionization_source = false
     advance_krook_collisions_ii = false
+    advance_maxwell_diffusion_ii = false
+    advance_maxwell_diffusion_nn = false
     advance_external_source = false
     advance_ion_numerical_dissipation = false
     advance_neutral_numerical_dissipation = false
@@ -1022,6 +1051,8 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     elseif t_params.implicit_vpa_advection
         advance_vpa_advection = true
         advance_ion_numerical_dissipation = true
+    elseif t_params.implicit_ion_maxwell_diffusion && collisions.mxwl_diff.D_ii > 0.0
+        advance_maxwell_diffusion_ii = true
     end
     # *_diffusion flags are set regardless of whether diffusion is included in explicit or
     # implicit part of timestep, because they are used for boundary conditions, not to
@@ -1032,9 +1063,9 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     # flag to determine if a d^2/dvpa^2 operator is present
     # When using implicit_vpa_advection, the vpa diffusion is included in the implicit
     # step
-    vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
+    vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1)) || (collisions.mxwl_diff.D_ii > 0.0)
     vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
-    vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0)
+    vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0) || (collisions.mxwl_diff.D_nn > 0.0)
 
     manufactured_solns_test = manufactured_solns_input.use_for_advance
 
@@ -2371,6 +2402,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, t, t_params, vz, vr, vzeta, vpa
 
     reset_nonlinear_per_stage_counters(nl_solver_params.ion_advance)
     reset_nonlinear_per_stage_counters(nl_solver_params.vpa_advection)
+    reset_nonlinear_per_stage_counters(nl_solver_params.maxwell_diffusion)
 
     istage = n_rk_stages+1
 
@@ -2732,6 +2764,16 @@ function backward_euler!(fvec_out, fvec_in, pdf, fields, moments, advect_objects
         if !success
             return success
         end
+    elseif advance.mxwl_diff_collisions_ii
+        success = implicit_ion_maxwell_diffusion!(fvec_out.pdf, fvec_in, moments,
+                                                  z_advect, vpa, vperp, z, r, dt,
+                                                  r_spectral, vpa_spectral, composition,
+                                                  collisions, geometry,
+                                                  nl_solver_params.maxwell_diffusion,
+                                                  gyroavs, scratch_dummy)
+        if !success
+            return success
+        end
     end
 
     return true

From 6b581a62484b554fcea38d32b26c196b47094c31 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 29 Jun 2024 21:48:55 +0100
Subject: [PATCH 342/428] In nonlinear solve, always do at least one Newton
 iteration

Ensures implicitly advanced quantities always get updated even when
timestep is very small. Do allow continuing without doing an iteration
if the residual is very small, to avoid creating NaNs.
---
 moment_kinetics/src/nonlinear_solvers.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl
index 7219c1530..7b20cdd8b 100644
--- a/moment_kinetics/src/nonlinear_solvers.jl
+++ b/moment_kinetics/src/nonlinear_solvers.jl
@@ -254,7 +254,7 @@ function newton_solve!(x, residual_func!, residual, delta_x, rhs_delta, v, w,
     close_linear_counter = -1
     success = true
     previous_residual_norm = residual_norm
-    while residual_norm > 1.0
+    while (counter < 1 && residual_norm > 1.0e-8) || residual_norm > 1.0
         counter += 1
         #println("\nNewton ", counter)
 

From 6253195929822c3129bd00e8f27aa4615ae4b318 Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.co.uk>
Date: Mon, 1 Jul 2024 15:48:25 +0000
Subject: [PATCH 343/428] Input file for slowing-down Fokker-Planck relaxation,
 without enforcing dF/dvperp = 0. Simulation shows reasonable behaviour with
 density preserved and thermal speed hitting a steady-state value near 0.15,
 which is held for multiple collision times. Ti/Talpha = 1/40 was chosen to
 reduce strain on resolution.

---
 ...planck-relaxation-slowing-down-alphas.toml | 92 +++++++++++++++++++
 1 file changed, 92 insertions(+)
 create mode 100644 examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml

diff --git a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
new file mode 100644
index 000000000..24aaf6e13
--- /dev/null
+++ b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
@@ -0,0 +1,92 @@
+# cheap input file for a 0D2V relaxation to a collisional Maxwellian distribution with self-ion collisions and collisions with fixed Maxwellian background of cold ions and electrons.
+n_ion_species = 1
+n_neutral_species = 0
+electron_physics = "boltzmann_electron_response"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.0
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 0.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "isotropic-beam"
+vpa_IC_v01 = 1.0
+vpa_IC_vth01 = 0.1
+#vpa_IC_option1 = "directed-beam"
+#vpa_IC_vpa01 = -1.5
+#vpa_IC_vperp01 = 0.0
+charge_exchange_frequency = 0.0
+ionization_frequency = 0.0
+constant_ionization_rate = false
+
+z_ngrid = 1
+z_nelement = 1
+z_nelement_local = 1
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+r_ngrid = 1
+r_nelement = 1
+r_nelement_local = 1
+r_bc = "periodic"
+r_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 5
+vpa_nelement = 32
+vpa_L = 3.0
+vpa_bc = "zero"
+vpa_discretization = "gausslegendre_pseudospectral"
+vperp_ngrid = 5
+vperp_nelement = 16
+vperp_L = 1.5
+vperp_discretization = "gausslegendre_pseudospectral"
+vperp_bc = "zero-no-regularity"
+# Fokker-Planck operator requires the "gausslegendre_pseudospectral
+# options for the vpa and vperp grids
+
+[fokker_planck_collisions]
+# nuii sets the normalised input C[F,F] Fokker-Planck collision frequency
+nuii = 1.876334222e-3 #(1/nu_alphae, as computed from input diagnostic)
+Zi = 2.0
+self_collisions = true
+slowing_down_test = true
+frequency_option = "manual"
+use_fokker_planck = true
+use_conserving_corrections = true
+sd_density = 1.0
+sd_temp = 0.025 # TD/Ealpha
+sd_mi = 0.5 # mD/malpha
+sd_me = 0.000013616 # 0.25/1836.0 me/malpha
+sd_q = 1.0
+
+[ion_source]
+active=false
+source_strength=0.0
+source_T=0.005
+source_n=1.0
+r_profile="constant"
+r_width=1.0
+r_relative_minimum=0.0
+z_profile="gaussian"
+z_width=0.1
+z_relative_minimum=0.0
+source_v0 = 1.0
+#source_type="alphas"
+source_type="alphas-with-losses"
+#source_type="beam-with-losses"
+#source_vpa0 = 1.0
+#source_vperp0 = 1.0
+sink_strength = 0.0
+sink_vth=0.07071067811865475                     
+
+[timestepping]
+nstep = 50000
+dt = 5.0e-4
+nwrite = 100
+nwrite_dfns = 100

From 5715036b45796d85e95d984ed7c6dfa461ab38b9 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 3 Jul 2024 12:57:05 +0000
Subject: [PATCH 344/428] Correct typo in input for vpa_IC_v0 parameter, no
 underscore assumed for other parameters and no underscore being used in the
 examples files, so remove it in the source for consistency.

---
 moment_kinetics/src/moment_kinetics_input.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 8e0772442..cddadfec0 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -548,7 +548,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
             species.ion[is].vpa_IC.upar_amplitude, species.ion[is].vpa_IC.upar_phase,
             species.ion[is].vpa_IC.temperature_amplitude,
             species.ion[is].vpa_IC.temperature_phase, species.ion[is].vpa_IC.monomial_degree,
-            get(scan_input, "vpa_IC_v0_$is", 0.5*sqrt(vperp.L^2 + (0.5*vpa.L)^2)),
+            get(scan_input, "vpa_IC_v0$is", 0.5*sqrt(vperp.L^2 + (0.5*vpa.L)^2)),
             get(scan_input, "vpa_IC_vth0$is", 0.1*sqrt(vperp.L^2 + (0.5*vpa.L)^2)),
             get(scan_input, "vpa_IC_vpa0$is", 0.25*0.5*abs(vpa.L)),
             get(scan_input, "vpa_IC_vperp0$is", 0.5*abs(vperp.L)))

From 7efa515554d09cfbacccfc9178e984c8b7053434 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 3 Jul 2024 14:14:45 +0000
Subject: [PATCH 345/428] Prevent animations being made if one of the
 dimensions has size 1.

---
 .../plots_post_processing/src/plots_post_processing.jl | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
index 972b92620..353df4961 100644
--- a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
+++ b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
@@ -2858,7 +2858,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
         spec_string = [string("_", spec_type)]
     end
     # make a gif animation of f(vpa,z,t) at a given (vperp,r) location
-    if pp.animate_f_vs_vpa_z
+    if pp.animate_f_vs_vpa_z && (vpa.n > 1 && z_local.n > 1)
         pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa;
@@ -2876,7 +2876,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
         end
     end
     # make a gif animation of f(vpa,r,t) at a given (vperp,z) location
-    if pp.animate_f_vs_vpa_r
+    if pp.animate_f_vs_vpa_r && (vpa.n > 1 && r_local.n > 1)
         pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                              itime_min:iskip:itime_max, n_species,
                                              r_local, z_local, vperp, vpa;
@@ -2894,7 +2894,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
         end
     end
     # make a gif animation of f(vperp,z,t) at a given (vpa,r) location
-    if pp.animate_f_vs_vperp_z
+    if pp.animate_f_vs_vperp_z && (vperp.n > 1 && z_local.n > 1)
         pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa; ivpa=ivpa0,
@@ -2908,7 +2908,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
         end
     end
     # make a gif animation of f(vperp,r,t) at a given (vpa,z) location
-    if pp.animate_f_vs_vperp_r
+    if pp.animate_f_vs_vperp_r && (vperp.n > 1 && r_local.n > 1)
         pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                                  itime_min:iskip:itime_max, n_species,
                                                  r_local, z_local, vperp, vpa; ivpa=ivpa0,
@@ -2970,7 +2970,7 @@ function plot_ion_pdf(run_name, run_name_label, vpa, vperp, z, r, z_local, r_loc
         end
     end
     # make a gif animation of f(z,r,t) at a given (vpa,vperp) location
-    if pp.animate_f_vs_r_z
+    if pp.animate_f_vs_r_z && (z_local.n > 1 && r_local.n > 1)
         pdf = load_distributed_ion_pdf_slice(run_name, nblocks,
                                              itime_min:iskip:itime_max, n_species,
                                              r_local, z_local, vperp, vpa; ivpa=ivpa0,

From 41e04fdedb8c6da537c8801b2cca5ef99f545417 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 3 Jul 2024 14:26:37 +0000
Subject: [PATCH 346/428] Reduce time step size of example input file to avoid
 numerical instability.

---
 .../fokker-planck-relaxation-slowing-down-alphas.toml           | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
index 24aaf6e13..dfab272e3 100644
--- a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
+++ b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
@@ -87,6 +87,6 @@ sink_vth=0.07071067811865475
 
 [timestepping]
 nstep = 50000
-dt = 5.0e-4
+dt = 2.5e-4
 nwrite = 100
 nwrite_dfns = 100

From 797952d5eb632183ebf60e3e0f0f38c10d7127ba Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Thu, 4 Jul 2024 14:53:16 +0000
Subject: [PATCH 347/428] Prevent pointless plots being made when r and z
 dimension size is 1.

---
 .../src/plots_post_processing.jl              | 63 ++++++++++---------
 1 file changed, 33 insertions(+), 30 deletions(-)

diff --git a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
index 353df4961..a29b57dd7 100644
--- a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
+++ b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
@@ -2806,13 +2806,13 @@ function plot_fields_rt(phi, delta_phi, time, itime_min, itime_max, nwrite_movie
         outfile = string(run_name, "_delta_phi(r0,z0)_vs_t.pdf")
         trysavefig(outfile)
     end
-    if pp.plot_phi_vs_z_t
+    if pp.plot_phi_vs_z_t && r.n > 1
         # make a heatmap plot of ϕ(r,t)
         heatmap(time, r.grid, phi, xlabel="time", ylabel="r", title="ϕ", c = :deep)
         outfile = string(run_name, "_phi_vs_r_t.pdf")
         trysavefig(outfile)
     end
-    if pp.animate_phi_vs_z
+    if pp.animate_phi_vs_z && r.n > 1
         # make a gif animation of ϕ(r) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views plot(r.grid, phi[:,i], xlabel="r", ylabel="ϕ", ylims = (phimin,phimax))
@@ -2826,10 +2826,11 @@ function plot_fields_rt(phi, delta_phi, time, itime_min, itime_max, nwrite_movie
     # 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")
     # trysavefig(outfile)
-    plot(r.grid, phi[:,end], xlabel="r/Lr", ylabel="eϕ/Te", label="", linewidth=2)
-    outfile = string(run_name, "_phi(r)_final.pdf")
-    trysavefig(outfile)
-
+    if r.n > 1
+        plot(r.grid, phi[:,end], xlabel="r/Lr", ylabel="eϕ/Te", label="", linewidth=2)
+        outfile = string(run_name, "_phi(r)_final.pdf")
+        trysavefig(outfile)
+    end
     println("done.")
 end
 
@@ -3067,15 +3068,16 @@ end
 function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
     itime_min, itime_max, nwrite_movie, run_name, pp, description)
     nr = size(r,1)
+    nz = size(z,1)
     print("Plotting fields data...")
     phimin = minimum(phi)
     phimax = maximum(phi)
-    if pp.plot_phi_vs_r0_z # plot last timestep phi[z,ir0]
+    if pp.plot_phi_vs_r0_z && nz > 1 # plot last timestep phi[z,ir0]
         @views plot(z, phi[:,ir0,end], xlabel=L"z/L_z", ylabel=L"\phi")
         outfile = string(run_name, "_phi"*description*"(r0,z)_vs_z.pdf")
         trysavefig(outfile)
     end
-    if pp.animate_phi_vs_r_z && nr > 1
+    if pp.animate_phi_vs_r_z && nr > 1 && nz > 1
         # make a gif animation of ϕ(z) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views heatmap(r, z, phi[:,:,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3087,7 +3089,7 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
         end
         outfile = string(run_name, "_phi(zwall-)_vs_r.gif")
         trygif(anim, outfile, fps=5)
-    elseif pp.animate_phi_vs_r_z && nr == 1 # make a gif animation of ϕ(z) at different times
+    elseif pp.animate_phi_vs_r_z && nr == 1 && nz > 1 # make a gif animation of ϕ(z) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views plot(z, phi[:,1,i], xlabel="z", ylabel=L"\widetilde{\phi}", ylims = (phimin,phimax))
         end
@@ -3096,17 +3098,17 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
     end
     Ezmin = minimum(Ez)
     Ezmax = maximum(Ez)
-    if pp.plot_Ez_vs_r0_z # plot last timestep Ez[z,ir0]
+    if pp.plot_Ez_vs_r0_z && nz > 1 # plot last timestep Ez[z,ir0]
         @views plot(z, Ez[:,ir0,end], xlabel=L"z/L_z", ylabel=L"E_z")
         outfile = string(run_name, "_Ez"*description*"(r0,z)_vs_z.pdf")
         trysavefig(outfile)
     end
-    if pp.plot_wall_Ez_vs_r && nr > 1 # plot last timestep Ez[z_wall,r]
+    if pp.plot_wall_Ez_vs_r && nr > 1 && nz > 1 # plot last timestep Ez[z_wall,r]
         @views plot(r, Ez[1,:,end], xlabel=L"r/L_r", ylabel=L"E_z")
         outfile = string(run_name, "_Ez"*description*"(r,z_wall-)_vs_r.pdf")
         trysavefig(outfile)
     end
-    if pp.animate_Ez_vs_r_z && nr > 1
+    if pp.animate_Ez_vs_r_z && nr > 1 && nz > 1
         # make a gif animation of ϕ(z) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views heatmap(r, z, Ez[:,:,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3118,7 +3120,7 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
         end
         outfile = string(run_name, "_Ez(zwall-)_vs_r.gif")
         trygif(anim, outfile, fps=5)
-    elseif pp.animate_Ez_vs_r_z && nr == 1
+    elseif pp.animate_Ez_vs_r_z && nr == 1 && nz > 1
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views plot(z, Ez[:,1,i], xlabel="z", ylabel=L"\widetilde{E}_z", ylims = (Ezmin,Ezmax))
         end
@@ -3127,12 +3129,12 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
     end
     Ermin = minimum(Er)
     Ermax = maximum(Er)
-    if pp.plot_Er_vs_r0_z # plot last timestep Er[z,ir0]
+    if pp.plot_Er_vs_r0_z && nz > 1 # plot last timestep Er[z,ir0]
         @views plot(z, Er[:,ir0,end], xlabel=L"z/L_z", ylabel=L"E_r")
         outfile = string(run_name, "_Er"*description*"(r0,z)_vs_z.pdf")
         trysavefig(outfile)
     end
-    if pp.plot_wall_Er_vs_r && nr > 1 # plot last timestep Er[z_wall,r]
+    if pp.plot_wall_Er_vs_r && nr > 1 && nz > 1 # plot last timestep Er[z_wall,r]
         @views plot(r, Er[1,:,end], xlabel=L"r/L_r", ylabel=L"E_r")
         outfile = string(run_name, "_Er"*description*"(r,z_wall-)_vs_r.pdf")
         trysavefig(outfile)
@@ -3142,7 +3144,7 @@ function plot_fields_2D(phi, Ez, Er, time, z, r, iz0, ir0,
         outfile = string(run_name, "_Er(zwall-)_vs_r.gif")
         trygif(anim, outfile, fps=5)
     end
-    if pp.animate_Er_vs_r_z && nr > 1
+    if pp.animate_Er_vs_r_z && nr > 1 && nz > 1
         # make a gif animation of ϕ(z) at different times
         anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
             @views heatmap(r, z, Er[:,:,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3159,6 +3161,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
     time, z, r, iz0, ir0, n_ion_species,
     itime_min, itime_max, nwrite_movie, run_name, pp)
     nr = size(r,1)
+    nz = size(z,1)
     ntime = size(time,1)
     print("Plotting ion moments data...")
     for is in 1:n_ion_species
@@ -3166,7 +3169,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 		# the density
 		densitymin = minimum(density[:,:,is,:])
 		densitymax = maximum(density)
-		if pp.plot_density_vs_r0_z # plot last timestep density[z,ir0]
+		if pp.plot_density_vs_r0_z && nz > 1 # plot last timestep density[z,ir0]
 			@views plot(z, density[:,ir0,is,end], xlabel=L"z/L_z", ylabel=L"n_i")
 			outfile = string(run_name, "_density"*description*"(r0,z)_vs_z.pdf")
 			trysavefig(outfile)
@@ -3176,7 +3179,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_density"*description*"(r,z_wall)_vs_r.pdf")
 			trysavefig(outfile)
 		end
-		if pp.animate_density_vs_r_z && nr > 1
+		if pp.animate_density_vs_r_z && nr > 1 && nz > 1
 			# make a gif animation of ϕ(z) at different times
 			anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
 				@views heatmap(r, z, density[:,:,is,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3184,7 +3187,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_density"*description*"_vs_r_z.gif")
 			trygif(anim, outfile, fps=5)
 		end
-		if pp.plot_density_vs_r_z && nr > 1
+		if pp.plot_density_vs_r_z && nr > 1 && nz > 1
 			@views heatmap(r, z, density[:,:,is,end], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
 			windowsize = (360,240), margin = 15pt)
 			outfile = string(run_name, "_density"*description*"_vs_r_z.pdf")
@@ -3201,17 +3204,17 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 		# the parallel flow
 		parallel_flowmin = minimum(parallel_flow[:,:,is,:])
 		parallel_flowmax = maximum(parallel_flow)
-		if pp.plot_parallel_flow_vs_r0_z # plot last timestep parallel_flow[z,ir0]
+		if pp.plot_parallel_flow_vs_r0_z && nz > 1 # plot last timestep parallel_flow[z,ir0]
 			@views plot(z, parallel_flow[:,ir0,is,end], xlabel=L"z/L_z", ylabel=L"u_{i\|\|}")
 			outfile = string(run_name, "_parallel_flow"*description*"(r0,z)_vs_z.pdf")
 			trysavefig(outfile)
 		end
-		if pp.plot_wall_parallel_flow_vs_r && nr > 1 # plot last timestep parallel_flow[z_wall,r]
+		if pp.plot_wall_parallel_flow_vs_r && nr > 1 && nz > 1# plot last timestep parallel_flow[z_wall,r]
 			@views plot(r, parallel_flow[end,:,is,end], xlabel=L"r/L_r", ylabel=L"u_{i\|\|}")
 			outfile = string(run_name, "_parallel_flow"*description*"(r,z_wall)_vs_r.pdf")
 			trysavefig(outfile)
 		end
-		if pp.animate_parallel_flow_vs_r_z && nr > 1
+		if pp.animate_parallel_flow_vs_r_z && nr > 1 && nz > 1
 			# make a gif animation of ϕ(z) at different times
 			anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
 				@views heatmap(r, z, parallel_flow[:,:,is,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3219,7 +3222,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_parallel_flow"*description*"_vs_r_z.gif")
 			trygif(anim, outfile, fps=5)
 		end
-		if pp.plot_parallel_flow_vs_r_z && nr > 1
+		if pp.plot_parallel_flow_vs_r_z && nr > 1 && nz > 1
 			@views heatmap(r, z, parallel_flow[:,:,is,end], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
 			windowsize = (360,240), margin = 15pt)
 			outfile = string(run_name, "_parallel_flow"*description*"_vs_r_z.pdf")
@@ -3236,7 +3239,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 		# the parallel pressure
 		parallel_pressuremin = minimum(parallel_pressure[:,:,is,:])
 		parallel_pressuremax = maximum(parallel_pressure)
-		if pp.plot_parallel_pressure_vs_r0_z # plot last timestep parallel_pressure[z,ir0]
+		if pp.plot_parallel_pressure_vs_r0_z && nz > 1 # plot last timestep parallel_pressure[z,ir0]
 			@views plot(z, parallel_pressure[:,ir0,is,end], xlabel=L"z/L_z", ylabel=L"p_{i\|\|}")
 			outfile = string(run_name, "_parallel_pressure"*description*"(r0,z)_vs_z.pdf")
 			trysavefig(outfile)
@@ -3246,7 +3249,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_parallel_pressure"*description*"(r,z_wall)_vs_r.pdf")
 			trysavefig(outfile)
 		end
-		if pp.animate_parallel_pressure_vs_r_z && nr > 1
+		if pp.animate_parallel_pressure_vs_r_z && nr > 1 && nz > 1
 			# make a gif animation of ϕ(z) at different times
 			anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
 				@views heatmap(r, z, parallel_pressure[:,:,is,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3254,7 +3257,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_parallel_pressure"*description*"_vs_r_z.gif")
 			trygif(anim, outfile, fps=5)
 		end
-		if pp.plot_parallel_pressure_vs_r_z && nr > 1
+		if pp.plot_parallel_pressure_vs_r_z && nr > 1 && nz > 1
 			@views heatmap(r, z, parallel_pressure[:,:,is,end], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
 			windowsize = (360,240), margin = 15pt)
 			outfile = string(run_name, "_parallel_pressure"*description*"_vs_r_z.pdf")
@@ -3264,17 +3267,17 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
                 # Note factor of 2 here because currently temperatures are normalised by
                 # Tref, while pressures are normalised by m*nref*c_ref^2=2*nref*Tref
                 temperature = 2.0 * parallel_pressure ./ density
-                if pp.plot_parallel_temperature_vs_r0_z # plot last timestep parallel_temperature[z,ir0]
+                if pp.plot_parallel_temperature_vs_r0_z  && nz > 1# plot last timestep parallel_temperature[z,ir0]
                     @views plot(z, temperature[:,ir0,is,end], xlabel=L"z/L_z", ylabel=L"T_i")
                     outfile = string(run_name, "_temperature"*description*"(r0,z)_vs_z.pdf")
                     trysavefig(outfile)
                 end
-                if pp.plot_wall_parallel_temperature_vs_r && nr > 1 # plot last timestep parallel_temperature[z_wall,r]
+                if pp.plot_wall_parallel_temperature_vs_r && nr > 1 && nz > 1 # plot last timestep parallel_temperature[z_wall,r]
                     @views plot(r, temperature[end,:,is,end], xlabel=L"r/L_r", ylabel=L"T_i")
                     outfile = string(run_name, "_temperature"*description*"(r,z_wall)_vs_r.pdf")
                     trysavefig(outfile)
                 end
-                if pp.animate_parallel_temperature_vs_r_z && nr > 1
+                if pp.animate_parallel_temperature_vs_r_z && nr > 1 && nz > 1
                     # make a gif animation of T_||(z) at different times
                     anim = @animate for i ∈ itime_min:nwrite_movie:itime_max
                         @views heatmap(r, z, temperature[:,:,is,i], xlabel="r", ylabel="z", c = :deep, interpolation = :cubic)
@@ -3282,7 +3285,7 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
                     outfile = string(run_name, "_temperature"*description*"_vs_r_z.gif")
                     trygif(anim, outfile, fps=5)
                 end
-                if pp.plot_parallel_temperature_vs_r_z && nr > 1
+                if pp.plot_parallel_temperature_vs_r_z && nr > 1 && nz > 1
                     @views heatmap(r, z, temperature[:,:,is,end], xlabel=L"r", ylabel=L"z", c = :deep, interpolation = :cubic,
                                   windowsize = (360,240), margin = 15pt)
                     outfile = string(run_name, "_temperature"*description*"_vs_r_z.pdf")

From e6bdb90023ea44ed8c5537d8a635448837c79f89 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Thu, 4 Jul 2024 16:07:30 +0000
Subject: [PATCH 348/428] Rename and add example file for slowing-down
 distribution test. Testing without self collisions is much faster.

---
 ...lowing-down-alphas-no-self-collisions.toml | 91 +++++++++++++++++++
 ...lanck-source-sink-slowing-down-alphas.toml | 91 +++++++++++++++++++
 2 files changed, 182 insertions(+)
 create mode 100644 examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
 create mode 100644 examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml

diff --git a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
new file mode 100644
index 000000000..e3070a6f7
--- /dev/null
+++ b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
@@ -0,0 +1,91 @@
+# cheap input file for a 0D2V relaxation to a collisional Maxwellian distribution with self-ion collisions and collisions with fixed Maxwellian background of cold ions and electrons.
+n_ion_species = 1
+n_neutral_species = 0
+electron_physics = "boltzmann_electron_response"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.0
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 0.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "isotropic-beam"
+vpa_IC_v01 = 1.0
+vpa_IC_vth01 = 0.1
+#vpa_IC_option1 = "directed-beam"
+#vpa_IC_vpa01 = -1.5
+#vpa_IC_vperp01 = 0.0
+charge_exchange_frequency = 0.0
+ionization_frequency = 0.0
+constant_ionization_rate = false
+
+z_ngrid = 1
+z_nelement = 1
+z_nelement_local = 1
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+r_ngrid = 1
+r_nelement = 1
+r_nelement_local = 1
+r_bc = "periodic"
+r_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 5
+vpa_nelement = 32
+vpa_L = 3.0
+vpa_bc = "zero"
+vpa_discretization = "gausslegendre_pseudospectral"
+vperp_ngrid = 5
+vperp_nelement = 16
+vperp_L = 1.5
+vperp_discretization = "gausslegendre_pseudospectral"
+vperp_bc = "zero-no-regularity"
+# Fokker-Planck operator requires the "gausslegendre_pseudospectral
+# options for the vpa and vperp grids
+
+[fokker_planck_collisions]
+# nuii sets the normalised input C[F,F] Fokker-Planck collision frequency
+nuii = 1.876334222e-3 #(1/nu_alphae, as computed from input diagnostic)
+Zi = 2.0
+self_collisions = false
+slowing_down_test = true
+frequency_option = "manual"
+use_fokker_planck = true
+sd_density = 1.0
+sd_temp = 0.0025 # TD/Ealpha
+sd_mi = 0.5 # mD/malpha
+sd_me = 0.000013616 # 0.25/1836.0 me/malpha
+sd_q = 1.0
+
+[ion_source]
+active=true
+source_strength=1.0
+source_T=0.005
+source_n=1.0
+r_profile="constant"
+r_width=1.0
+r_relative_minimum=0.0
+z_profile="gaussian"
+z_width=0.1
+z_relative_minimum=0.0
+source_v0 = 1.0
+#source_type="alphas"
+source_type="alphas-with-losses"
+#source_type="beam-with-losses"
+#source_vpa0 = 1.0
+#source_vperp0 = 1.0
+sink_strength = 1.0
+sink_vth=0.07071067811865475                     
+
+[timestepping]
+nstep = 250000
+dt = 1.0e-4
+nwrite = 500
+nwrite_dfns = 500
diff --git a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml
new file mode 100644
index 000000000..841af3a1e
--- /dev/null
+++ b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml
@@ -0,0 +1,91 @@
+# cheap input file for a 0D2V relaxation to a collisional Maxwellian distribution with self-ion collisions and collisions with fixed Maxwellian background of cold ions and electrons.
+n_ion_species = 1
+n_neutral_species = 0
+electron_physics = "boltzmann_electron_response"
+evolve_moments_density = false
+evolve_moments_parallel_flow = false
+evolve_moments_parallel_pressure = false
+evolve_moments_conservation = false
+T_e = 1.0
+T_wall = 1.0
+initial_density1 = 1.0
+initial_temperature1 = 1.0
+z_IC_option1 = "sinusoid"
+z_IC_density_amplitude1 = 0.0
+z_IC_density_phase1 = 0.0
+z_IC_upar_amplitude1 = 0.0
+z_IC_upar_phase1 = 0.0
+z_IC_temperature_amplitude1 = 0.0
+z_IC_temperature_phase1 = 0.0
+vpa_IC_option1 = "isotropic-beam"
+vpa_IC_v01 = 1.0
+vpa_IC_vth01 = 0.1
+#vpa_IC_option1 = "directed-beam"
+#vpa_IC_vpa01 = -1.5
+#vpa_IC_vperp01 = 0.0
+charge_exchange_frequency = 0.0
+ionization_frequency = 0.0
+constant_ionization_rate = false
+
+z_ngrid = 1
+z_nelement = 1
+z_nelement_local = 1
+z_bc = "wall"
+z_discretization = "chebyshev_pseudospectral"
+r_ngrid = 1
+r_nelement = 1
+r_nelement_local = 1
+r_bc = "periodic"
+r_discretization = "chebyshev_pseudospectral"
+vpa_ngrid = 5
+vpa_nelement = 32
+vpa_L = 3.0
+vpa_bc = "zero"
+vpa_discretization = "gausslegendre_pseudospectral"
+vperp_ngrid = 5
+vperp_nelement = 16
+vperp_L = 1.5
+vperp_discretization = "gausslegendre_pseudospectral"
+vperp_bc = "zero-no-regularity"
+# Fokker-Planck operator requires the "gausslegendre_pseudospectral
+# options for the vpa and vperp grids
+
+[fokker_planck_collisions]
+# nuii sets the normalised input C[F,F] Fokker-Planck collision frequency
+nuii = 1.876334222e-3 #(1/nu_alphae, as computed from input diagnostic)
+Zi = 2.0
+self_collisions = true
+slowing_down_test = true
+frequency_option = "manual"
+use_fokker_planck = true
+sd_density = 1.0
+sd_temp = 0.0025 # TD/Ealpha
+sd_mi = 0.5 # mD/malpha
+sd_me = 0.000013616 # 0.25/1836.0 me/malpha
+sd_q = 1.0
+
+[ion_source]
+active=true
+source_strength=1.0
+source_T=0.005
+source_n=1.0
+r_profile="constant"
+r_width=1.0
+r_relative_minimum=0.0
+z_profile="gaussian"
+z_width=0.1
+z_relative_minimum=0.0
+source_v0 = 1.0
+#source_type="alphas"
+source_type="alphas-with-losses"
+#source_type="beam-with-losses"
+#source_vpa0 = 1.0
+#source_vperp0 = 1.0
+sink_strength = 1.0
+sink_vth=0.07071067811865475                     
+
+[timestepping]
+nstep = 250000
+dt = 1.0e-4
+nwrite = 500
+nwrite_dfns = 500

From 8c84a1e8af5436678cba99843614e76a5965822d Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Fri, 5 Jul 2024 18:51:04 +0100
Subject: [PATCH 349/428] Change vperp numerical dissipation to use
 laplacian_derivative!(), support laplacian_derivative!() for chebyshev points
 with direct differentiation calculation.

---
 moment_kinetics/src/calculus.jl              | 70 +++++++++++++++++++-
 moment_kinetics/src/numerical_dissipation.jl |  4 +-
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/calculus.jl b/moment_kinetics/src/calculus.jl
index fbff04a2f..f4a73fb83 100644
--- a/moment_kinetics/src/calculus.jl
+++ b/moment_kinetics/src/calculus.jl
@@ -131,6 +131,70 @@ function second_derivative!(d2f, f, coord, spectral; handle_periodic=true)
     return nothing
 end
 
+function laplacian_derivative!(d2f, f, coord, spectral; handle_periodic=true)
+    # computes (1/coord) d / coord ( coord d f / d(coord)) for vperp coordinate
+    # For spectral element methods, calculate second derivative by applying first
+    # derivative twice, with special treatment for element boundaries
+
+    # First derivative
+    elementwise_derivative!(coord, f, spectral)
+    derivative_elements_to_full_grid!(coord.scratch3, coord.scratch_2d, coord)
+    if coord.name == "vperp"
+        # include the Jacobian
+        @. coord.scratch3 *= coord.grid
+    end
+    # MPI reconcile code here if used with z or r coords
+
+    # Save elementwise first derivative result
+    coord.scratch2_2d .= coord.scratch_2d
+
+    # Second derivative for element interiors
+    elementwise_derivative!(coord, coord.scratch3, spectral)
+    derivative_elements_to_full_grid!(d2f, coord.scratch_2d, coord)
+    if coord.name == "vperp"
+        # include the Jacobian
+        @. d2f /= coord.grid
+    end
+    # MPI reconcile code here if used with z or r coords
+
+    # Add contribution to penalise discontinuous first derivatives at element
+    # boundaries. For smooth functions this would do nothing so should not affect
+    # convergence of the second derivative. Aims to stabilise numerical instability when
+    # spike develops at an element boundary. The coefficient is an arbitrary choice, it
+    # should probably be large enough for stability but as small as possible.
+    #
+    # Arbitrary numerical coefficient
+    C = 1.0
+    function penalise_discontinuous_first_derivative!(d2f, imin, imax, df)
+        # Left element boundary
+        d2f[imin] += C * df[1]
+
+        # Right element boundary
+        d2f[imax] -= C * df[end]
+
+        return nothing
+    end
+    @views penalise_discontinuous_first_derivative!(d2f, 1, coord.imax[1],
+                                                    coord.scratch2_2d[:,1])
+    for ielement ∈ 2:coord.nelement_local
+        @views penalise_discontinuous_first_derivative!(d2f, coord.imin[ielement]-1,
+                                                        coord.imax[ielement],
+                                                        coord.scratch2_2d[:,ielement])
+    end
+
+    if coord.bc ∈ ("zero", "both_zero", "zero-no-regularity")
+        # For stability don't contribute to evolution at boundaries, in case these
+        # points are not set by a boundary condition.
+        # Full grid may be across processes and bc only applied to extreme ends of the
+        # domain.
+        if coord.irank == coord.nrank - 1
+            d2f[end] = 0.0
+        end
+    else
+        error("Unsupported bc '$(coord.bc)'")
+    end
+    return nothing
+end
 """
     mass_matrix_solve!(f, b, spectral::weak_discretization_info)
 
@@ -176,9 +240,11 @@ function laplacian_derivative!(d2f, f, coord, spectral::weak_discretization_info
     # for all other coord.name, do exactly the same as second_derivative! above.
     mul!(coord.scratch, spectral.L_matrix, f)
 
-    if handle_periodic && coord.bc == "periodic"
+    if coord.bc == "periodic" && coord.name == "vperp"
+        error("laplacian_derivative!() cannot handle periodic boundaries for vperp")
+    elseif coord.bc == "periodic"
         if coord.nrank > 1
-            error("second_derivative!() cannot handle periodic boundaries for a "
+            error("laplacian_derivative!() cannot handle periodic boundaries for a "
                   * "distributed coordinate")
         end
 
diff --git a/moment_kinetics/src/numerical_dissipation.jl b/moment_kinetics/src/numerical_dissipation.jl
index 86e517229..f48a158c6 100644
--- a/moment_kinetics/src/numerical_dissipation.jl
+++ b/moment_kinetics/src/numerical_dissipation.jl
@@ -12,7 +12,7 @@ export setup_numerical_dissipation, vpa_boundary_buffer_decay!,
 using Base.Iterators: flatten
 
 using ..looping
-using ..calculus: derivative!, second_derivative!
+using ..calculus: derivative!, second_derivative!, laplacian_derivative!
 using ..derivatives: derivative_r!, derivative_z!, second_derivative_r!,
                      second_derivative_z!
 using ..type_definitions: mk_float
@@ -369,7 +369,7 @@ function vperp_dissipation!(f_out, f_in, vperp, spectral::T_spectral, dt,
     end
     
     @loop_s_r_z_vpa is ir iz ivpa begin
-        @views second_derivative!(vperp.scratch, f_in[ivpa,:,iz,ir,is], vperp, spectral)
+        @views laplacian_derivative!(vperp.scratch, f_in[ivpa,:,iz,ir,is], vperp, spectral)
         @views @. f_out[ivpa,:,iz,ir,is] += dt * diffusion_coefficient * vperp.scratch
     end
 

From 4269c9ed3de61d5003458e2c84555d8311d28c91 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 11:28:34 +0100
Subject: [PATCH 350/428] Addition of tests of laplacian_derivative!(). Note
 the large errors, especially associated with the Chebyshev pseudospectral
 method, which does not show good convergence even at large ngrid. The
 finite-element GaussLegendre method shows much better behaviour.

---
 moment_kinetics/test/calculus_tests.jl | 201 +++++++++++++++++++++++++
 1 file changed, 201 insertions(+)

diff --git a/moment_kinetics/test/calculus_tests.jl b/moment_kinetics/test/calculus_tests.jl
index e439ac4cf..ef13b75bb 100644
--- a/moment_kinetics/test/calculus_tests.jl
+++ b/moment_kinetics/test/calculus_tests.jl
@@ -5,6 +5,7 @@ include("setup.jl")
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.calculus: derivative!, second_derivative!, integral
+using moment_kinetics.calculus: laplacian_derivative!
 
 using MPI
 using Random
@@ -1355,6 +1356,101 @@ function runtests()
             end
         end
         
+        @testset "Chebyshev pseudospectral cylindrical laplacian derivatives (4 argument), zero" verbose=false begin
+            @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈
+                    (
+                     (4, 7, 2.e-1),
+                     (4, 8, 2.e-1),
+                     (4, 9, 4.e-2),
+                     (4, 10, 4.e-2),
+                     (4, 11, 5.e-3),
+                     (4, 12, 5.e-3),
+                     (4, 13, 5.e-3),
+                     (4, 14, 5.e-3),
+                     (4, 15, 5.e-3),
+                     (4, 16, 5.e-3),
+                     (4, 17, 5.e-3),
+                     (4, 18, 5.e-3),
+                     (4, 19, 5.e-3),
+                     (4, 20, 5.e-3),
+                     (4, 21, 5.e-3),
+                     (4, 22, 5.e-3),
+                     (4, 23, 5.e-3),
+                     (4, 24, 4.e-3),
+                     (4, 25, 4.e-3),
+                     (4, 26, 4.e-3),
+                     (4, 27, 4.e-3),
+                     (4, 28, 4.e-3),
+                     (4, 29, 4.e-3),
+                     (4, 30, 4.e-3),
+                     (4, 31, 4.e-3),
+                     (4, 32, 4.e-3),
+                     (4, 33, 4.e-3),
+
+                     (5, 7, 2.e-1),
+                     (5, 8, 8.e-2),
+                     (5, 9, 5.e-2),
+                     (5, 10, 8.e-3),
+                     (5, 11, 8.e-3),
+                     (5, 12, 8.e-3),
+                     (5, 13, 8.e-3),
+                     (5, 14, 8.e-3),
+                     (5, 15, 8.e-3),
+                     (5, 16, 2.e-3),
+                     (5, 17, 2.e-3),
+                     (5, 18, 2.e-3),
+                     (5, 19, 2.e-3),
+                     (5, 20, 2.e-3),
+                     (5, 21, 2.e-3),
+                     (5, 22, 2.e-3),
+                     (5, 23, 2.e-3),
+                     (5, 24, 2.e-3),
+                     (5, 25, 2.e-3),
+                     (5, 26, 2.e-3),
+                     (5, 27, 2.e-3),
+                     (5, 28, 2.e-3),
+                     (5, 29, 2.e-3),
+                     (5, 30, 2.e-3),
+                     (5, 31, 2.e-3),
+                     (5, 32, 2.e-3),
+                     (5, 33, 2.e-3),
+                    ), cheb_option in ("FFT","matrix")
+
+                # define inputs needed for the test
+                L = 6.0
+                bc = "zero"
+                # fd_option and adv_input not actually used so given values unimportant
+                fd_option = ""
+                adv_input = advection_input("default", 1.0, 0.0, 0.0)
+                # create the 'input' struct containing input info needed to create a
+                # coordinate
+                nelement_local = nelement
+				nrank_per_block = 0 # dummy value
+				irank = 0 # dummy value
+				comm = MPI.COMM_NULL # dummy value
+                element_spacing_option = "uniform"
+				input = grid_input("vperp", ngrid, nelement,
+                    nelement_local, nrank_per_block, irank, L,
+                    "chebyshev_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
+                    element_spacing_option)
+                # create the coordinate struct 'x' and info for derivatives, etc.
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true)
+
+                f = @. exp(-x.grid^2)
+                expected_d2f = @. 4.0*(x.grid^2 - 1.0)*exp(-x.grid^2)
+                # create array for the derivative d2f/dx2
+                d2f = similar(f)
+
+                # differentiate f
+                laplacian_derivative!(d2f, f, x, spectral)
+
+                @test isapprox(d2f, expected_d2f, rtol=rtol, atol=1.e-10,
+                               norm=maxabs_norm)
+            end
+        end
+        
         @testset "GaussLegendre pseudospectral second derivatives (4 argument), periodic" verbose=false begin
             @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈
                     (
@@ -1460,6 +1556,111 @@ function runtests()
                 # differentiate f
                 second_derivative!(d2f, f, x, spectral)
 
+                @test isapprox(d2f, expected_d2f, rtol=rtol, atol=1.e-10,
+                               norm=maxabs_norm)
+            end
+        end
+        
+    @testset "GaussLegendre pseudospectral cylindrical laplacian derivatives (4 argument), zero" verbose=false begin
+            @testset "$nelement $ngrid" for (nelement, ngrid, rtol) ∈
+                    (
+                     (1, 8, 5.e-2),
+                     (1, 9, 1.e-1),
+                     (1, 10, 2.e-1),
+                     (1, 11, 5.e-2),
+                     (1, 12, 5.e-2),
+                     (1, 13, 5.e-2),
+                     (1, 14, 5.e-2),
+                     (1, 15, 5.e-3),
+                     (1, 16, 5.e-2),
+                     (1, 17, 5.e-3),
+                     
+                     (2, 6, 8.e-2),
+                     (2, 7, 8.e-2),
+                     (2, 8, 5.e-2),
+                     (2, 9, 5.e-2),
+                     (2, 10, 5.e-2),
+                     (2, 11, 5.e-3),
+                     (2, 12, 5.e-3),
+                     (2, 13, 5.e-4),
+                     (2, 14, 5.e-4),
+                     (2, 15, 5.e-4),
+                     (2, 16, 5.e-4),
+                     (2, 17, 5.e-4),
+                     
+                     (3, 6, 5.e-2),
+                     (3, 7, 5.e-3),
+                     (3, 8, 5.e-2),
+                     (3, 9, 5.e-4),
+                     (3, 10, 5.e-3),
+                     (3, 11, 5.e-4),
+                     (3, 12, 5.e-5),
+                     (3, 13, 5.e-5),
+                     (3, 14, 5.e-6),
+                     (3, 15, 5.e-6),
+                     (3, 16, 5.e-6),
+                     (3, 17, 5.e-8),
+                     
+                     (4, 5, 5.e-2),
+                     (4, 6, 2.e-2),
+                     (4, 7, 2.e-2),
+                     (4, 8, 2.e-3),
+                     (4, 9, 1.e-3),
+                     (4, 10, 1.e-4),
+                     (4, 11, 8.e-5),
+                     (4, 12, 8.e-6),
+                     (4, 13, 8.e-6),
+                     (4, 14, 8.e-7),
+                     (4, 15, 8.e-7),
+                     (4, 16, 8.e-8),
+                     (4, 17, 8.e-8),
+                     
+                     (5, 5, 4.e-2),
+                     (5, 6, 8.e-3),
+                     (5, 7, 5.e-3),
+                     (5, 8, 5.e-4),
+                     (5, 9, 8.e-5),
+                     (5, 10, 5.e-6),
+                     (5, 11, 8.e-6),
+                     (5, 12, 4.e-7),
+                     (5, 13, 2.e-7),
+                     (5, 14, 2.e-7),
+                     (5, 15, 8.e-7),
+                     (5, 16, 8.e-10),
+                     (5, 17, 8.e-10),
+                     )
+
+                # define inputs needed for the test
+                L = 6.0
+                bc = "zero"
+                # fd_option and adv_input not actually used so given values unimportant
+                fd_option = ""
+                adv_input = advection_input("default", 1.0, 0.0, 0.0)
+                cheb_option = ""
+                # create the 'input' struct containing input info needed to create a
+                # coordinate
+                nelement_local = nelement
+				nrank_per_block = 0 # dummy value
+				irank = 0 # dummy value
+				comm = MPI.COMM_NULL # dummy value
+                element_spacing_option = "uniform"
+				input = grid_input("vperp", ngrid, nelement,
+                    nelement_local, nrank_per_block, irank, L,
+                    "gausslegendre_pseudospectral", fd_option, cheb_option, bc, adv_input, comm,
+                    element_spacing_option)
+                # create the coordinate struct 'x' and info for derivatives, etc.
+                # This test runs effectively in serial, so use `ignore_MPI=true` to avoid
+                # errors due to communicators not being fully set up.
+                x, spectral = define_coordinate(input; ignore_MPI=true, collision_operator_dim=false)
+
+                f = @. exp(-x.grid^2)
+                expected_d2f = @. 4.0*(x.grid^2 - 1.0)*exp(-x.grid^2)
+                # create array for the derivative d2f/dx2
+                d2f = similar(f)
+
+                # differentiate f
+                laplacian_derivative!(d2f, f, x, spectral)
+
                 @test isapprox(d2f, expected_d2f, rtol=rtol, atol=1.e-10,
                                norm=maxabs_norm)
             end

From 99e9525a5d388b1d06056ea6bd2ee03eaa03bc42 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 12:31:25 +0100
Subject: [PATCH 351/428] Initial refactor moving Er_constant to geo struct.

---
 moment_kinetics/ext/manufactured_solns_ext.jl | 19 +++++++++++------
 moment_kinetics/src/em_fields.jl              |  9 ++++----
 moment_kinetics/src/geo.jl                    | 21 ++++++++++++++-----
 moment_kinetics/src/input_structs.jl          |  6 +++++-
 moment_kinetics/src/moment_kinetics_input.jl  |  6 +++---
 moment_kinetics/src/time_advance.jl           |  8 +++----
 6 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/moment_kinetics/ext/manufactured_solns_ext.jl b/moment_kinetics/ext/manufactured_solns_ext.jl
index de6d4e44b..b2bb21290 100644
--- a/moment_kinetics/ext/manufactured_solns_ext.jl
+++ b/moment_kinetics/ext/manufactured_solns_ext.jl
@@ -43,6 +43,8 @@ using IfElse
     # does not appear in a particular geometric coefficient, because 
     # that coefficient is a constant. 
     struct geometric_coefficients_sym{}
+        Er_constant::mk_float
+        Ez_constant::mk_float
         rhostar::mk_float
         Bzed::Union{mk_float,Num}
         Bzeta::Union{mk_float,Num}
@@ -62,6 +64,8 @@ using IfElse
         option = geometry_input_data.option
         rhostar = geometry_input_data.rhostar
         pitch = geometry_input_data.pitch
+        Er_constant = geometry_input_data.Er_constant
+        Ez_constant = geometry_input_data.Ez_constant
         if option == "constant-helical" || option == "default"
             bzed = pitch
             bzeta = sqrt(1 - bzed^2)
@@ -93,7 +97,8 @@ using IfElse
         else
             input_option_error("$option", option)
         end
-        geo_sym = geometric_coefficients_sym(rhostar,Bzed,Bzeta,Bmag,bzed,bzeta,dBdz,dBdr,jacobian)
+        geo_sym = geometric_coefficients_sym(Er_constant,Ez_constant,
+          rhostar,Bzed,Bzeta,Bmag,bzed,bzeta,dBdz,dBdr,jacobian)
         return geo_sym
     end
 
@@ -272,7 +277,7 @@ using IfElse
             upari = 0.0
         elseif z_bc == "wall"
             densi = densi_sym(Lr,Lz,r_bc,z_bc,composition,manufactured_solns_input,species)
-            Er, Ez, phi = electric_fields(Lr,Lz,r_bc,z_bc,composition,nr,manufactured_solns_input,species)
+            Er, Ez, phi = electric_fields(Lr,Lz,r_bc,z_bc,composition,geometry,nr,manufactured_solns_input,species)
             Bzeta = geometry.Bzeta
             Bmag = geometry.Bmag
             rhostar = geometry.rhostar
@@ -402,7 +407,7 @@ using IfElse
         return dfni
     end
 
-    function electric_fields(Lr, Lz, r_bc, z_bc, composition, nr,
+    function electric_fields(Lr, Lz, r_bc, z_bc, composition, geometry, nr,
                              manufactured_solns_input, species)
        
         # define derivative operators
@@ -500,11 +505,13 @@ using IfElse
         return manufactured_solns_list
     end 
     
-    function manufactured_electric_fields(Lr, Lz, r_bc, z_bc, composition, nr,
+    function manufactured_electric_fields(Lr, Lz, r_bc, z_bc, composition, geometry_input_data::geometry_input, nr,
                                           manufactured_solns_input, species)
         
+        # calculate the geometry symbolically
+        geometry = geometry_sym(geometry_input_data,Lz,Lr,nr)
         # calculate the electric fields and the potential
-        Er, Ez, phi = electric_fields(Lr, Lz, r_bc, z_bc, composition, nr,
+        Er, Ez, phi = electric_fields(Lr, Lz, r_bc, z_bc, composition, geometry, nr,
                                       manufactured_solns_input, species)
         
         Er_func = build_function(Er, z, r, t, expression=Val{false})
@@ -603,7 +610,7 @@ using IfElse
         
         # calculate the electric fields and the potential
         Er, Ez, phi = electric_fields(r_coord.L, z_coord.L, r_coord.bc, z_coord.bc,
-                                      composition, r_coord.n, manufactured_solns_input,
+                                      composition, geometry, r_coord.n, manufactured_solns_input,
                                       ion_species)
 
         # the adiabatic invariant (for compactness)
diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 6591e9062..5c021099b 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -31,7 +31,7 @@ end
 """
 update_phi updates the electrostatic potential, phi
 """
-function update_phi!(fields, fvec, vperp, z, r, composition, z_spectral, r_spectral, scratch_dummy, gyroavs::gyro_operators)
+function update_phi!(fields, fvec, vperp, z, r, composition, geometry, z_spectral, r_spectral, scratch_dummy, gyroavs::gyro_operators)
     n_ion_species = composition.n_ion_species
     @boundscheck size(fields.phi,1) == z.n || throw(BoundsError(fields.phi))
     @boundscheck size(fields.phi,2) == r.n || throw(BoundsError(fields.phi))
@@ -125,8 +125,8 @@ function update_phi!(fields, fvec, vperp, z, r, composition, z_spectral, r_spect
         end
     else
         @loop_r_z ir iz begin
-            fields.Er[iz,ir] = composition.Er_constant
-            # Er_constant defaults to 0.0 in moment_kinetics_input.jl
+            fields.Er[iz,ir] = geometry.input.Er_constant
+            # Er_constant defaults to 0.0 in geo.jl
         end
     end
     #Ez = - d phi / dz
@@ -137,7 +137,8 @@ function update_phi!(fields, fvec, vperp, z, r, composition, z_spectral, r_spect
                 z_spectral,z)
     else
         @serial_region begin
-            fields.Ez[:,:] .= 0.0
+            fields.Ez[:,:] .= = geometry.input.Ez_constant
+            # Ez_constant defaults to 0.0 in geo.jl
         end
     end
     
diff --git a/moment_kinetics/src/geo.jl b/moment_kinetics/src/geo.jl
index 9ad6d4885..e9c188bb8 100644
--- a/moment_kinetics/src/geo.jl
+++ b/moment_kinetics/src/geo.jl
@@ -73,11 +73,22 @@ option = ""
 
 """
 function setup_geometry_input(toml_input::Dict, reference_rhostar)
-    input_section = get(toml_input, "geometry", Dict{String,Any}())
-    if !("rhostar" ∈ keys(input_section))
-        # Set default rhostar with reference value
-        input_section["rhostar"] = get(input_section, "rhostar", reference_rhostar)
-    end
+    # read the input toml and specify a sensible default
+    input_section = set_defaults_and_check_section!(toml_input, "geometry",
+        # begin default inputs (as kwargs)
+        # rhostar ion (ref)
+        rhostar = reference_rhostar #used to premultiply ExB drift terms
+        # magnetic geometry option
+        option = "constant-helical" # "1D-mirror"
+        # pitch ( = Bzed/Bmag if geometry_option == "constant-helical")
+        pitch = 1.0,
+        # DeltaB ( = (Bzed(z=L/2) - Bzed(0))/Bref if geometry_option == "1D-mirror")
+        DeltaB = 0.0,
+        # constant for testing nonzero Er when nr = 1
+        Er_constant = 0.0,
+        # constant for testing nonzero Ez when nz = 1
+        Ez_constant = 0.0)
+    
     input = Dict(Symbol(k)=>v for (k,v) in input_section)
     #println(input)
     return geometry_input(; input...)
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 1c3f18f3f..634b2fadd 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -305,7 +305,7 @@ mutable struct species_composition
     # wall potential used if electron_physics=boltzmann_electron_response_with_simple_sheath
     phi_wall::mk_float
     # constant for testing nonzero Er
-    Er_constant::mk_float
+    #Er_constant::mk_float
     # ratio of the neutral particle mass to the ion mass
     mn_over_mi::mk_float
     # ratio of the electron particle mass to the ion mass
@@ -414,6 +414,10 @@ Base.@kwdef struct geometry_input
     pitch::mk_float = 1.0
     # DeltaB ( = (Bzed(z=L/2) - Bzed(0))/Bref if geometry_option == "1D-mirror")
     DeltaB::mk_float = 0.0
+    # constant for testing nonzero Er when nr = 1
+    Er_constant::mk_float
+    # constant for testing nonzero Ez when nz = 1
+    Ez_constant::mk_float  
 end
 
 @enum binary_format_type hdf5 netcdf
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index cddadfec0..393d5e55d 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -103,7 +103,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
     # if false use true Knudsen cosine for neutral wall bc
     composition.use_test_neutral_wall_pdf = get(scan_input, "use_test_neutral_wall_pdf", false)
     # constant to be used to test nonzero Er in wall boundary condition
-    composition.Er_constant = get(scan_input, "Er_constant", 0.0)
+    #composition.Er_constant = get(scan_input, "Er_constant", 0.0)
     # The ion flux reaching the wall that is recycled as neutrals is reduced by
     # `recycling_fraction` to account for ions absorbed by the wall.
     composition.recycling_fraction = get(scan_input, "recycling_fraction", 1.0)
@@ -1015,7 +1015,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     # wall potential at z = 0
     phi_wall = 0.0
     # constant to test nonzero Er
-    Er_constant = 0.0
+    #Er_constant = 0.0
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0
     # ratio of the electron particle mass to the ion particle mass
@@ -1025,7 +1025,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     recycling_fraction = 1.0
     gyrokinetic_ions = false
     composition = species_composition(n_species, n_ion_species, n_neutral_species,
-        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, Er_constant,
+        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, #Er_constant,
         mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
     
     species_ion = Array{species_parameters_mutable,1}(undef,n_ion_species)
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 27eec4a0f..0fbc1d1ac 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -394,7 +394,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
 
     # initialize the electrostatic potential
     begin_serial_region()
-    update_phi!(fields, scratch[1], vperp, z, r, composition, z_spectral, r_spectral, scratch_dummy, gyroavs)
+    update_phi!(fields, scratch[1], vperp, z, r, composition, geometry, z_spectral, r_spectral, scratch_dummy, gyroavs)
     @serial_region begin
         # save the initial phi(z) for possible use later (e.g., if forcing phi)
         fields.phi0 .= fields.phi
@@ -592,7 +592,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         end
     end
 
-    update_phi!(fields, scratch[1], vperp, z, r, composition, z_spectral, r_spectral,
+    update_phi!(fields, scratch[1], vperp, z, r, composition, geometry, z_spectral, r_spectral,
                 scratch_dummy, gyroavs)
     calculate_ion_moment_derivatives!(moments, scratch[1], scratch_dummy, z, z_spectral, 
                                       ion_mom_diss_coeff)
@@ -1572,7 +1572,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
     end
 
     # update the electrostatic potential phi
-    update_phi!(fields, scratch[istage+1], vperp, z, r, composition, z_spectral,
+    update_phi!(fields, scratch[istage+1], vperp, z, r, composition, geometry, z_spectral,
                 r_spectral, scratch_dummy, gyroavs)
     # _block_synchronize() here because phi needs to be read on different ranks than
     # it was written on, even though the loop-type does not change here. However,
@@ -1603,7 +1603,7 @@ function rk_update!(scratch, pdf, moments, fields, boundary_distributions, vz, v
             end
 
             # update the electrostatic potential phi
-            update_phi!(fields, scratch[istage+1], vperp, z, r, composition, z_spectral,
+            update_phi!(fields, scratch[istage+1], vperp, z, r, composition, geometry, z_spectral,
                         r_spectral, scratch_dummy, gyroavs)
             if !(( moments.evolve_upar || moments.evolve_ppar) &&
                       istage == length(scratch)-1)

From 4cde46c91d85bedf618e2fe47d6aba823b789009 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 12:35:21 +0100
Subject: [PATCH 352/428] Remove Er_constant from comments.

---
 .../fokker-planck-1D2V-even_nz-shorttest-nstep200.toml         | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml b/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
index eb7d60b18..7422ef0b0 100644
--- a/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
+++ b/examples/fokker-planck-1D2V/fokker-planck-1D2V-even_nz-shorttest-nstep200.toml
@@ -9,7 +9,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 #force_Er_zero_at_wall = false #true
-#Er_constant = 0.0
 #epsilon_offset = 0.1
 #use_vpabar_in_mms_dfni = true
 T_e = 1.0
@@ -95,4 +94,4 @@ vperp_discretization = "gausslegendre_pseudospectral"
 [fokker_planck_collisions]
 use_fokker_planck = true
 nuii = 1.0
-frequency_option = "manual"
\ No newline at end of file
+frequency_option = "manual"

From 705f4a3a0c66a646c62e86558188815722faf7eb Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 12:37:03 +0100
Subject: [PATCH 353/428] Delete out of date input files.

---
 ..._new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml | 98 ------------------
 ...MS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml | 99 -------------------
 2 files changed, 197 deletions(-)
 delete mode 100644 runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
 delete mode 100644 runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml

diff --git a/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml b/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
deleted file mode 100644
index 9a475346b..000000000
--- a/runs/1D-mirror_MMS_new_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
+++ /dev/null
@@ -1,98 +0,0 @@
-n_ion_species = 1
-n_neutral_species = 0
-electron_physics = "boltzmann_electron_response"
-#electron_physics = "boltzmann_electron_response_with_simple_sheath"
-evolve_moments_density = false
-evolve_moments_parallel_flow = false
-evolve_moments_parallel_pressure = false
-evolve_moments_conservation = false
-force_Er_zero_at_wall = false #true
-geometry_option="1D-mirror"
-DeltaB=10.0
-#geometry_option="constant-helical"
-#pitch=1.0
-Er_constant = 0.0
-T_e = 1.0
-T_wall = 1.0
-rhostar = 1.0
-initial_density1 = 0.5
-initial_temperature1 = 1.0
-initial_density2 = 0.5
-initial_temperature2 = 1.0
-z_IC_option1 = "sinusoid"
-z_IC_density_amplitude1 = 0.001
-z_IC_density_phase1 = 0.0
-z_IC_upar_amplitude1 = 0.0
-z_IC_upar_phase1 = 0.0
-z_IC_temperature_amplitude1 = 0.0
-z_IC_temperature_phase1 = 0.0
-z_IC_option2 = "sinusoid"
-z_IC_density_amplitude2 = 0.001
-z_IC_density_phase2 = 0.0
-z_IC_upar_amplitude2 = 0.0
-z_IC_upar_phase2 = 0.0
-z_IC_temperature_amplitude2 = 0.0
-z_IC_temperature_phase2 = 0.0
-charge_exchange_frequency = 0.0
-ionization_frequency = 0.0
-nstep = 2000
-dt = 0.0005
-nwrite = 200
-nwrite_dfns = 200
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
-z_ngrid = 17
-z_nelement = 16
-z_nelement_local = 16
-z_bc = "wall"
-z_discretization = "chebyshev_pseudospectral"
-r_ngrid = 1
-r_nelement = 1
-r_nelement_local = 1
-r_bc = "periodic"
-r_discretization = "chebyshev_pseudospectral"
-vpa_ngrid = 17
-vpa_nelement = 16
-vpa_L = 12.0
-vpa_bc = "zero"
-vpa_discretization = "chebyshev_pseudospectral"
-vperp_ngrid = 17
-vperp_nelement = 8
-vperp_L = 6.0
-#vperp_discretization = "finite_difference"
-vperp_discretization = "chebyshev_pseudospectral"
-vperp_bc = "zero"
-
-vz_ngrid = 17
-vz_nelement = 4
-vz_L = 12.0
-vz_bc = "periodic"
-vz_discretization = "chebyshev_pseudospectral"
-
-vr_ngrid = 17
-vr_nelement = 4
-vr_L = 12.0
-vr_bc = "periodic"
-vr_discretization = "chebyshev_pseudospectral"
-
-vzeta_ngrid = 17
-vzeta_nelement = 4
-vzeta_L = 12.0
-vzeta_bc = "periodic"
-vzeta_discretization = "chebyshev_pseudospectral"
-
-[manufactured_solns]
- use_for_advance=true
- use_for_init=true
- # constant to be used to control Ez divergence in MMS tests
- epsilon_offset=0.1
- # bool to control if dfni is a function of vpa or vpabar in MMS test
- use_vpabar_in_mms_dfni=true
- alpha_switch=1.0
- type="default"
-[ion_numerical_dissipation]
-vpa_dissipation_coefficient = 0.001
-vperp_dissipation_coefficient = 0.001
-#z_dissipation_coefficient = 0.1
-r_dissipation_coefficient = 0.0
diff --git a/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml b/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml
deleted file mode 100644
index 347dd2b04..000000000
--- a/runs/1D-mirror_MMS_new_nel_r_1_z_6_vpa_6_vperp_3_diss.toml
+++ /dev/null
@@ -1,99 +0,0 @@
-n_ion_species = 1
-n_neutral_species = 0
-electron_physics = "boltzmann_electron_response"
-#electron_physics = "boltzmann_electron_response_with_simple_sheath"
-evolve_moments_density = false
-evolve_moments_parallel_flow = false
-evolve_moments_parallel_pressure = false
-evolve_moments_conservation = false
-force_Er_zero_at_wall = false #true
-geometry_option="1D-mirror"
-DeltaB=0.5
-#geometry_option="constant-helical"
-#pitch=1.0
-Er_constant = 0.0
-T_e = 1.0
-T_wall = 1.0
-rhostar = 1.0
-initial_density1 = 0.5
-initial_temperature1 = 1.0
-initial_density2 = 0.5
-initial_temperature2 = 1.0
-z_IC_option1 = "sinusoid"
-z_IC_density_amplitude1 = 0.001
-z_IC_density_phase1 = 0.0
-z_IC_upar_amplitude1 = 0.0
-z_IC_upar_phase1 = 0.0
-z_IC_temperature_amplitude1 = 0.0
-z_IC_temperature_phase1 = 0.0
-z_IC_option2 = "sinusoid"
-z_IC_density_amplitude2 = 0.001
-z_IC_density_phase2 = 0.0
-z_IC_upar_amplitude2 = 0.0
-z_IC_upar_phase2 = 0.0
-z_IC_temperature_amplitude2 = 0.0
-z_IC_temperature_phase2 = 0.0
-charge_exchange_frequency = 0.0
-ionization_frequency = 0.0
-nstep = 2000
-dt = 0.0005
-nwrite = 200
-nwrite_dfns = 200
-use_semi_lagrange = false
-n_rk_stages = 4
-split_operators = false
-z_ngrid = 17
-z_nelement = 6
-z_nelement_local = 6
-z_bc = "wall"
-z_discretization = "chebyshev_pseudospectral"
-r_ngrid = 1
-r_nelement = 1
-r_nelement_local = 1
-r_bc = "periodic"
-r_discretization = "chebyshev_pseudospectral"
-vpa_ngrid = 17
-vpa_nelement = 6
-vpa_L = 12.0
-vpa_bc = "zero"
-vpa_discretization = "chebyshev_pseudospectral"
-vperp_ngrid = 17
-vperp_nelement = 3
-vperp_L = 6.0
-#vperp_discretization = "finite_difference"
-#vperp_discretization = "chebyshev_pseudospectral"
-vperp_discretization = "gausslegendre_pseudospectral"
-vperp_bc = "zero"
-
-vz_ngrid = 17
-vz_nelement = 4
-vz_L = 12.0
-vz_bc = "periodic"
-vz_discretization = "chebyshev_pseudospectral"
-
-vr_ngrid = 17
-vr_nelement = 4
-vr_L = 12.0
-vr_bc = "periodic"
-vr_discretization = "chebyshev_pseudospectral"
-
-vzeta_ngrid = 17
-vzeta_nelement = 4
-vzeta_L = 12.0
-vzeta_bc = "periodic"
-vzeta_discretization = "chebyshev_pseudospectral"
-
-[manufactured_solns]
- use_for_advance=true
- use_for_init=true
- # constant to be used to control Ez divergence in MMS tests
- epsilon_offset=0.1
- # bool to control if dfni is a function of vpa or vpabar in MMS test
- use_vpabar_in_mms_dfni=true
- alpha_switch=1.0
- type="default"
-[ion_numerical_dissipation]
-vpa_dissipation_coefficient = 0.001
-vperp_dissipation_coefficient = 0.001
-#z_dissipation_coefficient = 0.1
-r_dissipation_coefficient = 0.0

From d1d091f266e3abc98ac589b6398b9e4cd8a891a8 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 13:06:21 +0100
Subject: [PATCH 354/428] Move Er_constant to geometry namelist.

---
 ...1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml | 2 +-
 ...D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml | 2 +-
 runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml | 2 +-
 ...D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml | 2 +-
 runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml | 2 +-
 runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml     | 3 ++-
 runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml    | 3 ++-
 runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml     | 3 ++-
 ...D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml | 2 +-
 ...-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml | 2 +-
 runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml | 2 +-
 runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml | 2 +-
 runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml        | 2 +-
 13 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
index 883a73275..f22e52894 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_16_vpa_16_vperp_8_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=1.0
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
index 50e878cb3..fefea7811 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_32_vpa_32_vperp_16_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=1.0
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
index 1624ca129..58051541c 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_4_vpa_4_vperp_2_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=1.0
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
index 74e3af388..0cc6b6ef1 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_64_vpa_64_vperp_32_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=1.0
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
index f4d4956ab..0f3a24c74 100644
--- a/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
+++ b/runs/1D-mirror_MMS_ngrid_9_nel_r_1_z_8_vpa_8_vperp_4_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=1.0
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
index 8bf508ee1..96611183c 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -89,3 +88,5 @@ split_operators = false
 vpa_dissipation_coefficient = 0.01
 #z_dissipation_coefficient = 0.1
 #r_dissipation_coefficient = 0.0
+[geometry]
+Er_constant = 0.0
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
index 1abe96001..346996f69 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_16_vperp_1_krook.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -93,3 +92,5 @@ split_operators = false
 vpa_dissipation_coefficient = -1.0
 z_dissipation_coefficient = -1.0
 r_dissipation_coefficient = -1.0
+[geometry]
+Er_constant = 0.0
diff --git a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
index d9ddc1aeb..3378a93b1 100644
--- a/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
+++ b/runs/1D-wall_MMS_new_nel_r_1_z_16_vpa_8_vperp_8_krook.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -96,3 +95,5 @@ split_operators = false
 vpa_dissipation_coefficient = -1.0
 z_dissipation_coefficient = -1.0
 r_dissipation_coefficient = -1.0
+[geometry]
+Er_constant = 0.0
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
index a4a31ccbb..29542c43f 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_16_z_16_vpa_16_vperp_8_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=0.5
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
index 01dd1b487..3561f1d11 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_32_z_32_vpa_32_vperp_16_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=0.5
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
index b13205d70..a8add1f3e 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_4_z_4_vpa_4_vperp_2_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=0.5
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml b/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
index f19b91e39..2eb57f6ee 100644
--- a/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
+++ b/runs/2D-mirror_MMS_ngrid_5_nel_r_8_z_8_vpa_8_vperp_4_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -99,3 +98,4 @@ DeltaB=0.5
 #option="constant-helical"
 pitch=0.5
 rhostar = 1.0
+Er_constant = 0.0
diff --git a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
index 274fc59e6..fa1f6c220 100644
--- a/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
+++ b/runs/2D-wall_MMS_nel_r_32_z_32_vpa_16_vperp_1_diss.toml
@@ -7,7 +7,6 @@ evolve_moments_parallel_flow = false
 evolve_moments_parallel_pressure = false
 evolve_moments_conservation = false
 force_Er_zero_at_wall = false #true
-Er_constant = 0.0
 T_e = 1.0
 T_wall = 1.0
 initial_density1 = 0.5
@@ -95,3 +94,4 @@ r_dissipation_coefficient = 0.1
 option="constant-helical"
 pitch=0.5
 rhostar = 1.0
+Er_constant = 0.0

From 1ad159cadb03364a4bbabaa9bc7d9aec29a8db6c Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 13:50:36 +0100
Subject: [PATCH 355/428] Fix various bugs in last commit so that refactored
 code runs and MMS tests can be executed.

---
 .../makie_post_processing/src/shared_utils.jl | 19 +++++--------------
 moment_kinetics/ext/manufactured_solns_ext.jl |  8 ++++----
 moment_kinetics/src/em_fields.jl              |  4 ++--
 moment_kinetics/src/geo.jl                    |  6 +++---
 moment_kinetics/test/gyroaverage_tests.jl     | 12 +++++-------
 .../src/plot_MMS_sequence.jl                  |  2 +-
 .../src/plots_post_processing.jl              |  2 +-
 7 files changed, 21 insertions(+), 32 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/shared_utils.jl b/makie_post_processing/makie_post_processing/src/shared_utils.jl
index 031bcda75..f168bd457 100644
--- a/makie_post_processing/makie_post_processing/src/shared_utils.jl
+++ b/makie_post_processing/makie_post_processing/src/shared_utils.jl
@@ -5,7 +5,6 @@ export calculate_and_write_frequencies, get_geometry, get_composition
 using moment_kinetics.analysis: fit_delta_phi_mode
 using moment_kinetics.array_allocation: allocate_float
 using moment_kinetics.coordinates: define_coordinate
-using moment_kinetics.geo: init_magnetic_geometry
 using moment_kinetics.input_structs: boltzmann_electron_response,
                                      boltzmann_electron_response_with_simple_sheath,
                                      grid_input, geometry_input, species_composition
@@ -13,7 +12,7 @@ using moment_kinetics.moment_kinetics_input: get_default_rhostar, setup_referenc
 using moment_kinetics.type_definitions: mk_float, mk_int
 using moment_kinetics.reference_parameters: setup_reference_parameters
 using moment_kinetics.moment_kinetics_input: get_default_rhostar
-using moment_kinetics.geo: init_magnetic_geometry
+using moment_kinetics.geo: init_magnetic_geometry, setup_geometry_input
 using MPI
 
 """
@@ -90,7 +89,7 @@ function get_composition(scan_input)
     use_test_neutral_wall_pdf = get(scan_input, "use_test_neutral_wall_pdf", false)
     gyrokinetic_ions = get(scan_input, "gyrokinetic_ions", false)
     # constant to be used to test nonzero Er in wall boundary condition
-    Er_constant = get(scan_input, "Er_constant", 0.0)
+    #Er_constant = get(scan_input, "Er_constant", 0.0)
     recycling_fraction = get(scan_input, "recycling_fraction", 1.0)
     # constant to be used to control Ez divergences
     epsilon_offset = get(scan_input, "epsilon_offset", 0.001)
@@ -106,7 +105,7 @@ function get_composition(scan_input)
     # ratio of the electron particle mass to the ion particle mass
     me_over_mi = 1.0/1836.0
     composition = species_composition(n_species, n_ion_species, n_neutral_species,
-        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, Er_constant,
+        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, #Er_constant,
         mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
     return composition
 
@@ -114,18 +113,10 @@ end
 
 function get_geometry(scan_input,z,r)
     reference_params = setup_reference_parameters(scan_input)
-    # set geometry_input
-    # MRH need to get this in way that does not duplicate code
-    # MRH from moment_kinetics_input.jl
-    option = get(scan_input, "geometry_option", "constant-helical") #"1D-mirror"
-    pitch = get(scan_input, "pitch", 1.0)
-    rhostar = get(scan_input, "rhostar", get_default_rhostar(reference_params))
-    DeltaB = get(scan_input, "DeltaB", 1.0)
-    geo_in = geometry_input(rhostar,option,pitch,DeltaB)
+    reference_rhostar = get_default_rhostar(reference_params)
+    geo_in = setup_geometry_input(scan_input, reference_rhostar)
     geometry = init_magnetic_geometry(geo_in,z,r)
-    
     return geometry
-
 end
 
 end # shared_utils.jl
diff --git a/moment_kinetics/ext/manufactured_solns_ext.jl b/moment_kinetics/ext/manufactured_solns_ext.jl
index 8f1f42650..314853725 100644
--- a/moment_kinetics/ext/manufactured_solns_ext.jl
+++ b/moment_kinetics/ext/manufactured_solns_ext.jl
@@ -362,7 +362,7 @@ using IfElse
 
         if manufactured_solns_input.type == "default"
             # calculate the electric fields and the potential
-            Er, Ez, phi = electric_fields(Lr, Lz, r_bc, z_bc, composition, nr,
+            Er, Ez, phi = electric_fields(Lr, Lz, r_bc, z_bc, composition, geometry, nr,
                                           manufactured_solns_input, species)
 
             # get geometric/composition data
@@ -417,7 +417,7 @@ using IfElse
         # get N_e factor for boltzmann response
         if composition.electron_physics == boltzmann_electron_response_with_simple_sheath && nr == 1 
             # so 1D MMS test with 3V neutrals where ion current can be calculated prior to knowing Er
-            jpari_into_LHS_wall = jpari_into_LHS_wall_sym(Lr, Lz, r_bc, z_bc, composition,
+            jpari_into_LHS_wall = jpari_into_LHS_wall_sym(Lr, Lz, r_bc, z_bc,
                                                           manufactured_solns_input)
             N_e = -2.0*sqrt(pi*composition.me_over_mi)*exp(-composition.phi_wall/composition.T_e)*jpari_into_LHS_wall
         elseif composition.electron_physics == boltzmann_electron_response_with_simple_sheath && nr > 1 
@@ -442,8 +442,8 @@ using IfElse
         # calculate the electric fields
         dense = densi # get the electron density via quasineutrality with Zi = 1
         phi = composition.T_e*log(dense/N_e) # use the adiabatic response of electrons for me/mi -> 0
-        Er = -Dr(phi)*rfac + composition.Er_constant
-        Ez = -Dz(phi)
+        Er = -Dr(phi)*rfac + geometry.Er_constant
+        Ez = -Dz(phi)      + geometry.Ez_constant
         
         Er_expanded = expand_derivatives(Er)
         Ez_expanded = expand_derivatives(Ez)
diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 5c021099b..67534d3c3 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -136,8 +136,8 @@ function update_phi!(fields, fvec, vperp, z, r, composition, geometry, z_spectra
                 scratch_dummy.buffer_r_3, scratch_dummy.buffer_r_4,
                 z_spectral,z)
     else
-        @serial_region begin
-            fields.Ez[:,:] .= = geometry.input.Ez_constant
+        @loop_r_z ir iz begin
+            fields.Ez[iz,ir] = geometry.input.Ez_constant
             # Ez_constant defaults to 0.0 in geo.jl
         end
     end
diff --git a/moment_kinetics/src/geo.jl b/moment_kinetics/src/geo.jl
index e9c188bb8..75637425a 100644
--- a/moment_kinetics/src/geo.jl
+++ b/moment_kinetics/src/geo.jl
@@ -8,7 +8,7 @@ module geo
 export init_magnetic_geometry
 export setup_geometry_input
 
-using ..input_structs: geometry_input
+using ..input_structs: geometry_input, set_defaults_and_check_section!
 using ..file_io: input_option_error
 using ..array_allocation: allocate_float
 using ..type_definitions: mk_float, mk_int
@@ -77,9 +77,9 @@ function setup_geometry_input(toml_input::Dict, reference_rhostar)
     input_section = set_defaults_and_check_section!(toml_input, "geometry",
         # begin default inputs (as kwargs)
         # rhostar ion (ref)
-        rhostar = reference_rhostar #used to premultiply ExB drift terms
+        rhostar = reference_rhostar, #used to premultiply ExB drift terms
         # magnetic geometry option
-        option = "constant-helical" # "1D-mirror"
+        option = "constant-helical",# "1D-mirror"
         # pitch ( = Bzed/Bmag if geometry_option == "constant-helical")
         pitch = 1.0,
         # DeltaB ( = (Bzed(z=L/2) - Bzed(0))/Bref if geometry_option == "1D-mirror")
diff --git a/moment_kinetics/test/gyroaverage_tests.jl b/moment_kinetics/test/gyroaverage_tests.jl
index 38a6c7085..21c43f0b7 100644
--- a/moment_kinetics/test/gyroaverage_tests.jl
+++ b/moment_kinetics/test/gyroaverage_tests.jl
@@ -10,7 +10,7 @@ using SpecialFunctions: besselj0
 import moment_kinetics
 using moment_kinetics.input_structs
 using moment_kinetics.coordinates: define_coordinate
-using moment_kinetics.geo: init_magnetic_geometry
+using moment_kinetics.geo: init_magnetic_geometry, setup_geometry_input
 using moment_kinetics.communication
 using moment_kinetics.looping
 using moment_kinetics.array_allocation: allocate_float, allocate_shared_float
@@ -142,11 +142,9 @@ function gyroaverage_test(absolute_error; rhostar=0.1, pitch=0.5, ngrid=5, kr=2,
         gyrophase, gyrophase_spectral = define_coordinate(gyrophase_input; collision_operator_dim=false, run_directory=test_output_directory)
         
         # create test geometry
-        #rhostar = 0.1 #rhostar of ions for ExB drift
         option = "constant-helical"
-        #pitch = 1.0
-        DeltaB = 1.0
-        geometry_in = geometry_input(rhostar,option,pitch,DeltaB)
+        inputdict = Dict("geometry" => Dict("option" => option, "rhostar" => rhostar, "pitch" => pitch))
+        geometry_in = setup_geometry_input(inputdict, 0.1)
         geometry = init_magnetic_geometry(geometry_in,z,r)
         
         # create test composition
@@ -243,7 +241,7 @@ function create_test_composition()
     # wall potential at z = 0
     phi_wall = 0.0
     # constant to test nonzero Er
-    Er_constant = 0.0
+    #Er_constant = 0.0
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0
     # ratio of the electron particle mass to the ion particle mass
@@ -253,7 +251,7 @@ function create_test_composition()
     recycling_fraction = 1.0
     gyrokinetic_ions = true
     return composition = species_composition(n_species, n_ion_species, n_neutral_species,
-            electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, Er_constant,
+            electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, #Er_constant,
             mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
 end
 
diff --git a/plots_post_processing/plots_post_processing/src/plot_MMS_sequence.jl b/plots_post_processing/plots_post_processing/src/plot_MMS_sequence.jl
index 84bfcf21b..92d780b27 100644
--- a/plots_post_processing/plots_post_processing/src/plot_MMS_sequence.jl
+++ b/plots_post_processing/plots_post_processing/src/plot_MMS_sequence.jl
@@ -255,7 +255,7 @@ function get_MMS_error_data(path_list,scan_type,scan_name)
         dfnn_func = manufactured_solns_list.dfnn_func
         densn_func = manufactured_solns_list.densn_func
         
-        manufactured_E_fields = manufactured_electric_fields(Lr_in,z.L,r_bc,z_bc,composition,r.n,manufactured_solns_input,species)
+        manufactured_E_fields = manufactured_electric_fields(Lr_in,z.L,r_bc,z_bc,composition,geo_in,r.n,manufactured_solns_input,species)
         Er_func = manufactured_E_fields.Er_func
         Ez_func = manufactured_E_fields.Ez_func
         phi_func = manufactured_E_fields.phi_func
diff --git a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
index a29b57dd7..81d1e3828 100644
--- a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
+++ b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
@@ -1097,7 +1097,7 @@ function analyze_and_plot_data(prefix...; run_index=nothing)
         densn_func = manufactured_solns_list.densn_func
         manufactured_E_fields =
             manufactured_electric_fields(Lr_in, z_global.L, r_global.bc, z_global.bc,
-                                         composition, r_global.n,
+                                         composition, geometry.input, r_global.n,
                                          manufactured_solns_input, species)
         Er_func = manufactured_E_fields.Er_func
         Ez_func = manufactured_E_fields.Ez_func

From 72bd6496dc75cac72d931490f9e2c68d0d7dc325 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Mon, 8 Jul 2024 16:45:07 +0100
Subject: [PATCH 356/428] Add option to simulate a 0D2V plasma with collisions
 and advection terms in vpa and vperp, for testing of the compatibility of the
 pseudospectral and the finite-element methods for advection and collisons,
 respectively. The option is chosen by selecting the [geometry] option =
 "0D-Spitzer-test", and picking fixed constant values of Ez, dBdr, Er, and
 dBdr. The vperp_advection!() routine must update the z_advect and r_advect
 structs if nr == 1 and nz == 1 and the "0D-Spitzer-test" option is being
 used, leading to some duplication of code. This could be avoided if the speed
 updates happened outside the function where the advection terms are computed.

---
 moment_kinetics/ext/manufactured_solns_ext.jl |  9 ++++
 moment_kinetics/src/geo.jl                    | 31 +++++++++++++-
 moment_kinetics/src/input_structs.jl          |  6 ++-
 moment_kinetics/src/time_advance.jl           | 11 ++---
 moment_kinetics/src/vperp_advection.jl        | 42 ++++++++++++++++++-
 5 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/ext/manufactured_solns_ext.jl b/moment_kinetics/ext/manufactured_solns_ext.jl
index 314853725..e65c964db 100644
--- a/moment_kinetics/ext/manufactured_solns_ext.jl
+++ b/moment_kinetics/ext/manufactured_solns_ext.jl
@@ -94,6 +94,15 @@ using IfElse
             end
             dBdz = expand_derivatives(Dz(Bmag))
             jacobian = 1.0
+        elseif option == "0D-Spitzer-test"
+            Bmag = 1.0
+            dBdz = geometry_input_data.dBdz_constant
+            dBdr = geometry_input_data.dBdr_constant
+            bzed = pitch
+            bzeta = sqrt(1 - bzed^2)
+            Bzed = Bmag*bzed
+            Bzeta = Bmag*bzeta
+            jacobian = 1.0
         else
             input_option_error("$option", option)
         end
diff --git a/moment_kinetics/src/geo.jl b/moment_kinetics/src/geo.jl
index 75637425a..3635230dc 100644
--- a/moment_kinetics/src/geo.jl
+++ b/moment_kinetics/src/geo.jl
@@ -87,7 +87,11 @@ function setup_geometry_input(toml_input::Dict, reference_rhostar)
         # constant for testing nonzero Er when nr = 1
         Er_constant = 0.0,
         # constant for testing nonzero Ez when nz = 1
-        Ez_constant = 0.0)
+        Ez_constant = 0.0,
+        # constant for testing nonzero dBdz when nz = 1
+        dBdz_constant = 0.0,
+        # constant for testing nonzero dBdr when nr = 1
+        dBdr_constant = 0.0)
     
     input = Dict(Symbol(k)=>v for (k,v) in input_section)
     #println(input)
@@ -198,6 +202,31 @@ function init_magnetic_geometry(geometry_input_data::geometry_input,z,r)
                 gbdriftz[iz,ir] = cvdriftz[iz,ir]
             end
         end
+    elseif option == "0D-Spitzer-test"
+     # a 0D configuration with certain geometrical factors
+     # set to be constants to enable testing of velocity
+     # space operators such as mirror or vperp advection terms
+     pitch = geometry_input_data.pitch
+     dBdz_constant = geometry_input_data.dBdz_constant
+     dBdr_constant = geometry_input_data.dBdr_constant
+     B0 = 1.0 # chose reference field strength to be Bzeta at r = 1
+     for ir in 1:nr
+            for iz in 1:nz
+                Bmag[iz,ir] = B0
+                bzed[iz,ir] = pitch
+                bzeta[iz,ir] = sqrt(1 - pitch^2)
+                Bzed[iz,ir] = bzed[iz,ir]*Bmag[iz,ir]
+                Bzeta[iz,ir] = bzeta[iz,ir]*Bmag[iz,ir]
+                dBdz[iz,ir] = dBdz_constant
+                dBdr[iz,ir] = dBdr_constant
+                jacobian[iz,ir] = 1.0
+                
+                cvdriftr[iz,ir] = 0.0
+                cvdriftz[iz,ir] = 0.0
+                gbdriftr[iz,ir] = 0.0
+                gbdriftz[iz,ir] = 0.0
+            end
+        end
     else 
         input_option_error("$option", option)
     end
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 634b2fadd..cddb11226 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -417,7 +417,11 @@ Base.@kwdef struct geometry_input
     # constant for testing nonzero Er when nr = 1
     Er_constant::mk_float
     # constant for testing nonzero Ez when nz = 1
-    Ez_constant::mk_float  
+    Ez_constant::mk_float
+    # constant for testing nonzero dBdz when nz = 1
+    dBdz_constant::mk_float
+    # constant for testing nonzero dBdr when nr = 1
+    dBdr_constant::mk_float
 end
 
 @enum binary_format_type hdf5 netcdf
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 0fbc1d1ac..f9df6e910 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -649,8 +649,8 @@ function setup_advance_flags(moments, composition, t_params, collisions,
     # otherwise, check to see if the flags need to be set to true
     if !t_params.split_operators
         # default for non-split operators is to include both vpa and z advection together
-        advance_vpa_advection = vpa.n > 1 && z.n > 1
-        advance_vperp_advection = vperp.n > 1 && z.n > 1
+        advance_vpa_advection = vpa.n > 1 #&& z.n > 1
+        advance_vperp_advection = vperp.n > 1 #&& z.n > 1
         advance_z_advection = z.n > 1
         advance_r_advection = r.n > 1
         if collisions.fkpl.nuii > 0.0 && vperp.n > 1 
@@ -2085,7 +2085,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
 
     # r advection relies on derivatives in z to get ExB
-    if advance.r_advection && r.n > 1
+    if advance.r_advection
         r_advection!(fvec_out.pdf, fvec_in, moments, fields, r_advect, r, z, vperp, vpa,
                      dt, r_spectral, composition, geometry, scratch_dummy)
     end
@@ -2093,7 +2093,8 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     # so call vperp_advection! only after z and r advection routines
     if advance.vperp_advection
         vperp_advection!(fvec_out.pdf, fvec_in, vperp_advect, r, z, vperp, vpa,
-                      dt, vperp_spectral, composition, z_advect, r_advect, geometry)
+                      dt, vperp_spectral, composition, z_advect, r_advect, geometry,
+                      moments, fields, t)
     end
 
     if advance.source_terms
@@ -2106,7 +2107,7 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
             r, z, vzeta, vr, vz, dt, t, z_spectral, composition, scratch_dummy)
     end
 
-    if advance.neutral_r_advection && r.n > 1
+    if advance.neutral_r_advection
         neutral_advection_r!(fvec_out.pdf_neutral, fvec_in, neutral_r_advect,
             r, z, vzeta, vr, vz, dt, r_spectral, composition, geometry, scratch_dummy)
     end
diff --git a/moment_kinetics/src/vperp_advection.jl b/moment_kinetics/src/vperp_advection.jl
index 16c2697df..2dc6bd5c8 100644
--- a/moment_kinetics/src/vperp_advection.jl
+++ b/moment_kinetics/src/vperp_advection.jl
@@ -6,13 +6,21 @@ export update_speed_vperp!
 using ..advection: advance_f_local!
 using ..chebyshev: chebyshev_info
 using ..looping
+using ..z_advection: update_speed_z!
+using ..r_advection: update_speed_r!
 
 # do a single stage time advance (potentially as part of a multi-stage RK scheme)
 function vperp_advection!(f_out, fvec_in, vperp_advect, r, z, vperp, vpa,
-                      dt, vperp_spectral, composition, z_advect, r_advect, geometry)
+                      dt, vperp_spectral, composition, z_advect, r_advect, geometry,
+                      moments, fields, t)
+    
+    # if appropriate, update z and r speeds
+    update_z_r_speeds!(z_advect, r_advect, fvec_in, moments, fields,
+                            geometry, vpa, vperp, z, r, t)
+    
     begin_s_r_z_vpa_region()
     @loop_s is begin
-        # get the updated speed along the r direction using the current f
+        # get the updated speed along the vperp direction using the current f
         @views update_speed_vperp!(vperp_advect[is], vpa, vperp, z, r, z_advect[is], r_advect[is], geometry)
         @loop_r_z_vpa ir iz ivpa begin
             @views advance_f_local!(f_out[ivpa,:,iz,ir,is], fvec_in.pdf[ivpa,:,iz,ir,is],
@@ -58,4 +66,34 @@ function update_speed_vperp!(vperp_advect, vpa, vperp, z, r, z_advect, r_advect,
     return nothing
 end
 
+function update_z_r_speeds!(z_advect, r_advect, fvec_in, moments, fields,
+                            geometry, vpa, vperp, z, r, t)
+    update_z_speed = (z.n == 1 && geometry.input.option == "0D-Spitzer-test")
+    if update_z_speed
+    # make sure z speed is physical despite
+    # 0D nature of simulation
+        begin_s_r_vperp_vpa_region()
+        @loop_s is begin
+            # get the updated speed along the z direction using the current f
+            @views update_speed_z!(z_advect[is], fvec_in.upar[:,:,is],
+                                   moments.ion.vth[:,:,is], moments.evolve_upar,
+                                   moments.evolve_ppar, fields, vpa, vperp, z, r, t, geometry, is)
+        end
+    end
+    
+    update_r_speed = (r.n == 1 && geometry.input.option == "0D-Spitzer-test")
+    if update_r_speed
+    # make sure r speed is physical despite
+    # 0D nature of simulation
+        begin_s_z_vperp_vpa_region()
+        @loop_s is begin
+            # get the updated speed along the r direction using the current f
+            update_speed_r!(r_advect[is], fvec_in.upar[:,:,is],
+                               moments.ion.vth[:,:,is], fields, moments.evolve_upar,
+                               moments.evolve_ppar, vpa, vperp, z, r, geometry, is)
+        end
+    end
+    return nothing
+end    
+
 end

From 8bff6a5658e8e3012fb8c0bef81c01c02760fb22 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 10 Jul 2024 11:56:02 +0100
Subject: [PATCH 357/428] Changes to permit testing of "none" boundary
 condition option for Fokker-Planck time evolution. In this mode of operation,
 boundary conditions on the fluxes are imposed through the assembly of the
 finite-element method, and no additional zero-ing of the distribution
 function is carried out. Initial tests suggest that this mode of operation is
 at least as good as other tested options.

---
 moment_kinetics/src/boundary_conditions.jl    |  2 +
 moment_kinetics/src/fokker_planck_calculus.jl | 16 ++++--
 .../fokker_planck_time_evolution_tests.jl     | 56 +++++++++++++++++++
 3 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index 7082c1383..27282d65e 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -954,6 +954,8 @@ function enforce_v_boundary_condition_local!(f, bc, speed, v_diffusion, v, v_spe
     elseif bc == "periodic"
         f[1] = 0.5*(f[1]+f[end])
         f[end] = f[1]
+    elseif bc == "none"
+        # Do nothing
     else
         error("Unsupported boundary condition option '$bc' for $(v.name)")
     end
diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl
index 7838608fe..4007e6356 100644
--- a/moment_kinetics/src/fokker_planck_calculus.jl
+++ b/moment_kinetics/src/fokker_planck_calculus.jl
@@ -2287,18 +2287,22 @@ function enforce_vpavperp_BCs!(pdf,vpa,vperp,vpa_spectral,vperp_spectral)
     D0 = vperp_spectral.radau.D0
     # vpa boundary conditions
     # zero at infinity
-    begin_anyv_vperp_region()
-    @loop_vperp ivperp begin
-        pdf[1,ivperp] = 0.0
-        pdf[nvpa,ivperp] = 0.0
+    if vpa.bc == "zero"
+        begin_anyv_vperp_region()
+        @loop_vperp ivperp begin
+            pdf[1,ivperp] = 0.0
+            pdf[nvpa,ivperp] = 0.0
+        end
     end
     # vperp boundary conditions
     # zero boundary condition at infinity
     # set regularity condition d F / d vperp = 0 at vperp = 0
     # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
     begin_anyv_vpa_region()
-    @loop_vpa ivpa begin
-        pdf[ivpa,nvperp] = 0.0
+    if vperp.bc in ("zero", "zero-no-regularity")
+        @loop_vpa ivpa begin
+            pdf[ivpa,nvperp] = 0.0
+        end
     end
     if vperp.bc == "zero"
         buffer = @view vperp.scratch[1:ngrid_vperp-1]
diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index 669af32de..b5d5086d6 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -128,6 +128,54 @@ expected_data(
     0.007980077110978 0.005701898246888 0.003327438535529 0.000925955494251 -0.000037963012435 -0.000174809935222 0.000000000000000 ;
     0.000473874379555 0.000190726186343 -0.000067413540342 -0.000219129923463 -0.000165609965457 -0.000069940808382 0.000000000000000 ;
     0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000])
+
+const expected_none_bc = 
+expected_data(
+   [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
+   [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
+   # Expected phi:
+   [-1.254104813718096, -1.254104813718360],
+   # Expected n_ion:
+   [0.285331158471152, 0.285331158471077],
+   # Expected upar_ion:
+   [0.0, 0.0],
+   # Expected ppar_ion:
+   [0.182321263834348, 0.154595996906667],
+   # Expected pperp_ion
+   [0.143470130069393, 0.157332763533171],
+   # Expected qpar_ion
+   [0.0, 0.0],
+   # Expected v_t_ion
+   [1.047097792428007, 1.047097792428005],
+   # Expected dSdt
+   [0.000000000000000, 0.000000019115425],
+   # Expected f_ion:
+   [0.000045179366280 0.000030530376095 0.000017024973661 0.000004877736620 0.000000847623528 0.000000089338863 0.000000005711242 ;
+    0.000706724195475 0.000477575434536 0.000266315395816 0.000076300638379 0.000013259062819 0.000001397494940 0.000000089338863 ;
+    0.006705212475828 0.004531109564814 0.002526730124657 0.000723920301085 0.000125798485463 0.000013259062819 0.000000847623528 ;
+    0.038585833650058 0.026074735222541 0.014540327934436 0.004165873701136 0.000723920300963 0.000076300638366 0.000004877736619 ;
+    0.134671678398729 0.091005636629981 0.050748427134097 0.014539667807029 0.002526615411287 0.000266303305116 0.000017024200728 ;
+    0.261709398369233 0.176852554997681 0.098620144126568 0.028255144359305 0.004910007858078 0.000517511020834 0.000033083372713 ;
+    0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 0.000000000000000 ;
+    0.261709398369233 0.176852554997682 0.098620144126568 0.028255144359305 0.004910007858078 0.000517511020834 0.000033083372713 ;
+    0.134671678398729 0.091005636629981 0.050748427134097 0.014539667807029 0.002526615411287 0.000266303305116 0.000017024200728 ;
+    0.038585833650058 0.026074735222541 0.014540327934436 0.004165873701136 0.000723920300963 0.000076300638366 0.000004877736619 ;
+    0.006705212475828 0.004531109564814 0.002526730124657 0.000723920301085 0.000125798485463 0.000013259062819 0.000000847623528 ;
+    0.000706724195475 0.000477575434536 0.000266315395816 0.000076300638379 0.000013259062819 0.000001397494940 0.000000089338863 ;
+    0.000045179366280 0.000030530376095 0.000017024973661 0.000004877736620 0.000000847623528 0.000000089338863 0.000000005711242 ;;;
+    0.000447615535468 0.000364801746561 0.000270093752229 0.000138454732861 0.000052917434226 0.000014331973588 -0.000000022222325 ;
+    0.000530676740860 0.000337644325418 0.000161654162384 0.000044642657051 0.000030848860156 0.000021605350900 0.000014496785601 ;
+    0.006652401806692 0.004725322877613 0.002773765135978 0.000908975538425 0.000210924069594 0.000037798409136 0.000057517483301 ;
+    0.030674916797523 0.021423584760802 0.012211671337997 0.003838328661082 0.000895825759878 0.000045643481270 0.000159762122747 ;
+    0.097751022951371 0.068620792366922 0.039365625553724 0.012397445282748 0.002818297719268 0.000149127033057 0.000311439525732 ;
+    0.194564361178625 0.137570544970062 0.079900546066037 0.025549350260855 0.005673491663727 0.000408759661293 0.000445141805523 ;
+    0.242684098382075 0.171867841736931 0.100113917461427 0.032148650920148 0.007091277425658 0.000554833974109 0.000500018986081 ;
+    0.194564361178625 0.137570544970062 0.079900546066037 0.025549350260855 0.005673491663727 0.000408759661293 0.000445141805523 ;
+    0.097751022951371 0.068620792366922 0.039365625553724 0.012397445282748 0.002818297719268 0.000149127033057 0.000311439525732 ;
+    0.030674916797523 0.021423584760802 0.012211671337997 0.003838328661082 0.000895825759878 0.000045643481270 0.000159762122747 ;
+    0.006652401806692 0.004725322877613 0.002773765135978 0.000908975538425 0.000210924069594 0.000037798409136 0.000057517483301 ;
+    0.000530676740860 0.000337644325418 0.000161654162384 0.000044642657051 0.000030848860156 0.000021605350900 0.000014496785601 ;
+    0.000447615535468 0.000364801746561 0.000270093752229 0.000138454732861 0.000052917434226 0.000014331973588 -0.000000022222325])
 ###########################################################################################
 # to modify the test, with a new expected f, print the new f using the following commands
 # in an interative Julia REPL. The path is the path to the .dfns file. 
@@ -399,6 +447,14 @@ function runtests()
             expected_no_regularity,
              1.0e-14, 1.0e-14; vperp_bc="zero-no-regularity")
         end
+        @testset "Gauss Legendre no (explicitly) enforced boundary conditions" begin
+            run_name = "gausslegendre_pseudospectral_none_bc"
+            vperp_bc = "none"
+            vpa_bc = "none"
+            run_test(test_input_gauss_legendre,
+            expected_none_bc,
+             1.0e-14, 1.0e-14; vperp_bc=vperp_bc, vpa_bc=vpa_bc)
+        end
     end
 end
 

From 7b2c84a003742964b007f0256bbe09bd7f4358e8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 1 Jul 2024 09:53:16 +0100
Subject: [PATCH 358/428] Fix merge of Maxwell diffusion and kinetic electrons

---
 moment_kinetics/src/initial_conditions.jl | 15 ++++++++++-----
 moment_kinetics/src/maxwell_diffusion.jl  |  8 +++++---
 moment_kinetics/src/time_advance.jl       |  7 ++++++-
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index c449d1895..495d796f0 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -444,11 +444,16 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         scratch[1].electron_ppar .= moments.electron.ppar
         scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
         scratch[1].electron_temp .= moments.electron.temp
-        scratch[t_params.electron.n_rk_stages+1].electron_density .= moments.electron.dens
-        scratch[t_params.electron.n_rk_stages+1].electron_upar .= moments.electron.upar
-        scratch[t_params.electron.n_rk_stages+1].electron_ppar .= moments.electron.ppar
-        scratch[t_params.electron.n_rk_stages+1].electron_pperp .= 0.0 #moments.electron.pperp
-        scratch[t_params.electron.n_rk_stages+1].electron_temp .= moments.electron.temp
+        if t_params.electron === nothing
+            n_rk_stages = length(scratch)
+        else
+            n_rk_stages = t_params.electron.n_rk_stages
+        end
+        scratch[n_rk_stages+1].electron_density .= moments.electron.dens
+        scratch[n_rk_stages+1].electron_upar .= moments.electron.upar
+        scratch[n_rk_stages+1].electron_ppar .= moments.electron.ppar
+        scratch[n_rk_stages+1].electron_pperp .= 0.0 #moments.electron.pperp
+        scratch[n_rk_stages+1].electron_temp .= moments.electron.temp
     end
     if scratch_electron !== nothing
         begin_serial_region()
diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 0b67c142c..431bec193 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -290,9 +290,11 @@ function implicit_ion_maxwell_diffusion!(f_out, fvec_in, moments, z_advect, vpa,
 
     # Ensure moments are consistent with f_new
     new_scratch = scratch_pdf(f_out, fvec_in.density, fvec_in.upar, fvec_in.ppar,
-                              fvec_in.pperp, fvec_in.temp_z_s, fvec_in.pdf_neutral,
-                              fvec_in.density_neutral, fvec_in.uz_neutral,
-                              fvec_in.pz_neutral)
+                              fvec_in.pperp, fvec_in.temp_z_s, fvec_in.electron_density,
+                              fvec_in.electron_upar, fvec_in.electron_ppar,
+                              fvec_in.electron_pperp, fvec_in.electron_temp,
+                              fvec_in.pdf_neutral, fvec_in.density_neutral,
+                              fvec_in.uz_neutral, fvec_in.pz_neutral)
     update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
                             r_spectral, geometry, gyroavs, scratch_dummy, z_advect, false)
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 0957f2d56..b6770f1bd 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2866,8 +2866,13 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # Note the timestep for this solve is rk_coefs_implict[istage,istage]*dt.
                 # The diagonal elements are equal to the Butcher 'a' coefficients
                 # rk_coefs_implicit[istage,istage]=a[istage,istage].
+                if scratch_electron === nothing
+                    this_scratch_electron = nothing
+                else
+                    this_scratch_electron = scratch_electron[t_params.electron.n_rk_stages+1]
+                end
                 nl_success = backward_euler!(scratch_implicit[istage], scratch[istage],
-                                             scratch_electron[t_params.electron.n_rk_stages+1],
+                                             this_scratch_electron,
                                              pdf, fields, moments, advect_objects, vz, vr,
                                              vzeta, vpa, vperp, gyrophase, z, r, t,
                                              t_params.dt[] *

From 6c228d7e20ea2d4d2c0a2f59e769bc77b37bb502 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 17:34:14 +0100
Subject: [PATCH 359/428] Simplify ion sheath entrance boundary condition

Create separate function that calculates the indices nearest vpa=0,
which can also be used elsewhere (e.g. in implicit solves).
---
 moment_kinetics/src/boundary_conditions.jl | 43 +++++++++++++---------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index 61dff1a08..cc1aa436c 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -416,6 +416,26 @@ function enforce_zero_incoming_bc!(pdf, speed, z, zero)
         end
     end
 end
+function get_ion_z_boundary_cutoff_indices(density, upar, ppar, evolve_upar, evolve_ppar,
+                                           z, vpa, zero)
+    if z.irank == 0
+        vth = sqrt(2.0*(ppar[1]/density[1]))
+        @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, vth,
+                                         upar[1], evolve_ppar, evolve_upar)
+        last_negative_vpa_ind = searchsortedlast(vpa.scratch, -zero)
+    else
+        last_negative_vpa_ind = nothing
+    end
+    if z.irank == z.nrank - 1
+        vth = sqrt(2.0*(ppar[end]/density[end]))
+        @. vpa.scratch2 = vpagrid_to_dzdt(vpa.grid, vth,
+                                          upar[end], evolve_ppar, evolve_upar)
+        first_positive_vpa_ind = searchsortedfirst(vpa.scratch2, zero)
+    else
+        first_positive_vpa_ind = nothing
+    end
+    return last_negative_vpa_ind, first_positive_vpa_ind
+end
 function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density, upar,
                                    ppar, evolve_upar, evolve_ppar, zero)
     if z.irank != 0 && z.irank != z.nrank - 1
@@ -429,28 +449,15 @@ function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density,
     # so use advection speed below instead of vpa
 
     # absolute velocity at left boundary
+    last_negative_vpa_ind, first_positive_vpa_ind =
+        get_ion_z_boundary_cutoff_indices(density, upar, ppar, evolve_upar, evolve_ppar,
+                                          z, vpa, zero)
     if z.irank == 0
-        @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[1]/density[1])),
-                                         upar[1], evolve_ppar, evolve_upar)
-        @loop_vpa ivpa begin
-            # for left boundary in zed (z = -Lz/2), want
-            # f(z=-Lz/2, v_parallel > 0) = 0
-            if vpa.scratch[ivpa] > zero
-                pdf[ivpa,:,1] .= 0.0
-            end
-        end
+        pdf[last_negative_vpa_ind+1:end, :, 1] .= 0.0
     end
     # absolute velocity at right boundary
     if z.irank == z.nrank - 1
-        @. vpa.scratch2 = vpagrid_to_dzdt(vpa.grid, sqrt(2.0*(ppar[end]/density[end])),
-                                          upar[end], evolve_ppar, evolve_upar)
-        @loop_vpa ivpa begin
-            # for right boundary in zed (z = Lz/2), want
-            # f(z=Lz/2, v_parallel < 0) = 0
-            if vpa.scratch2[ivpa] < -zero
-                pdf[ivpa,:,end] .= 0.0
-            end
-        end
+        pdf[1:first_positive_vpa_ind-1, :, end] .= 0.0
     end
 
     # Special constraint-forcing code that tries to keep the modifications smooth at

From a2bb5d88d79809741c1b6652ce545343ac5d9113 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 19:32:06 +0100
Subject: [PATCH 360/428] Modify constraint enforcement in ion sheath entrance
 bc

Reduces corrections to distribution function at large w_parallel, seems
to improve numerical stability at the point where bc cutoff index jumps
from one point to the next.
---
 moment_kinetics/src/boundary_conditions.jl | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index cc1aa436c..1a929bdbf 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -485,9 +485,12 @@ function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density,
             # Store v_parallel with upar shift removed in vpa.scratch
             vth = sqrt(2.0*ppar[iz]/density[iz])
             @. vpa.scratch = vpa.grid + upar[iz]/vth
-            # Introduce factor to ensure corrections go smoothly to zero near
-            # v_parallel=0
-            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
+            # Introduce factors to ensure corrections go smoothly to zero near
+            # v_parallel=0, and that there are no large corrections aw large w_parallel as
+            # those can have a strong effect on the parallel heat flux and make
+            # timestepping unstable when the cut-off point jumps from one grid point to
+            # another.
+            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch)) / (1.0 + (4.0 * vpa.scratch / vpa.L)^4)
             J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
             J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
             J3 = integrate_over_vspace(vpa.scratch2, vpa.grid, 3, vpa.wgts)
@@ -505,9 +508,12 @@ function enforce_zero_incoming_bc!(pdf, z::coordinate, vpa::coordinate, density,
 
             # Store v_parallel with upar shift removed in vpa.scratch
             @. vpa.scratch = vpa.grid + upar[iz]
-            # Introduce factor to ensure corrections go smoothly to zero near
-            # v_parallel=0
-            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch))
+            # Introduce factors to ensure corrections go smoothly to zero near
+            # v_parallel=0, and that there are no large corrections aw large w_parallel as
+            # those can have a strong effect on the parallel heat flux and make
+            # timestepping unstable when the cut-off point jumps from one grid point to
+            # another.
+            @. vpa.scratch2 = f * abs(vpa.scratch) / (1.0 + abs(vpa.scratch)) / (1.0 + (4.0 * vpa.scratch / vpa.L)^4)
             J1 = integrate_over_vspace(vpa.scratch2, vpa.grid, vpa.wgts)
             J2 = integrate_over_vspace(vpa.scratch2, vpa.grid, 2, vpa.wgts)
 

From fb09eb5e29e95aed2c1e5fa409bc3c45b13afa45 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 19:43:48 +0100
Subject: [PATCH 361/428] Fix external source, source_type="energy"

---
 moment_kinetics/src/external_sources.jl | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 793d89550..be2206f94 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -608,10 +608,18 @@ function external_ion_source!(pdf, fvec, moments, ion_source_settings, vperp, vp
         end
 
         if source_type == "energy"
-            # Take particles out of pdf so source does not change density
-            @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-                pdf[ivpa,ivperp,iz,ir,is] -= dt * source_amplitude[iz,ir] *
-                    fvec.pdf[ivpa,ivperp,iz,ir,is]
+            if moments.evolve_density
+                # Take particles out of pdf so source does not change density
+                @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+                    pdf[ivpa,ivperp,iz,ir,is] -= dt * source_amplitude[iz,ir] *
+                        fvec.pdf[ivpa,ivperp,iz,ir,is]
+                end
+            else
+                # Take particles out of pdf so source does not change density
+                @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+                    pdf[ivpa,ivperp,iz,ir,is] -= dt * source_amplitude[iz,ir] *
+                        fvec.pdf[ivpa,ivperp,iz,ir,is] / fvec.density[iz,ir,is]
+                end
             end
         end
     elseif source_type == "alphas" || source_type == "alphas-with-losses"

From 4f001326e503db81138a23e8ec2e55ce4bbc012f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 20:18:26 +0100
Subject: [PATCH 362/428] Fix initialize_electrons!() when not evolving kinetic
 electrons

---
 moment_kinetics/src/initial_conditions.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 495d796f0..e1025217b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -445,7 +445,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         scratch[1].electron_pperp .= 0.0 #moments.electron.pperp
         scratch[1].electron_temp .= moments.electron.temp
         if t_params.electron === nothing
-            n_rk_stages = length(scratch)
+            n_rk_stages = length(scratch) - 1
         else
             n_rk_stages = t_params.electron.n_rk_stages
         end

From b81e522515c33f46c9bf6902bc254057f94c9f6f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 20:26:30 +0100
Subject: [PATCH 363/428] Move collision_frequency_ei, collision_frequency_ee
 to electron_moment_variables

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

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index c02c89398..23d4f9dae 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -54,11 +54,11 @@ const timestep_diagnostic_variables = ("time_for_run", "step_counter", "dt",
 const em_variables = ("phi", "Er", "Ez")
 const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "thermal_speed", "temperature", "parallel_heat_flux",
-                              "collision_frequency_ii", "collision_frequency_ee",
-                              "collision_frequency_ei", "sound_speed", "mach_number")
+                              "collision_frequency_ii", "sound_speed", "mach_number")
 const electron_moment_variables = ("electron_density", "electron_parallel_flow",
                                    "electron_parallel_pressure", "electron_thermal_speed",
                                    "electron_temperature", "electron_parallel_heat_flux",
+                                   "collision_frequency_ee", "collision_frequency_ei",
                                    "electron_dudz", "electron_dpdz", "electron_dqdz")
 const neutral_moment_variables = ("density_neutral", "uz_neutral", "pz_neutral",
                                   "thermal_speed_neutral", "temperature_neutral",

From d06880211425b326d744f412860b2059a0362d36 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 20:31:28 +0100
Subject: [PATCH 364/428] Diagnostics for total energy and total energy flux

---
 moment_kinetics/src/load_data.jl | 58 ++++++++++++++++++++++++++++++--
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 23d4f9dae..2e89d25e0 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -54,7 +54,8 @@ const timestep_diagnostic_variables = ("time_for_run", "step_counter", "dt",
 const em_variables = ("phi", "Er", "Ez")
 const ion_moment_variables = ("density", "parallel_flow", "parallel_pressure",
                               "thermal_speed", "temperature", "parallel_heat_flux",
-                              "collision_frequency_ii", "sound_speed", "mach_number")
+                              "collision_frequency_ii", "sound_speed", "mach_number",
+                              "total_energy", "total_energy_flux")
 const electron_moment_variables = ("electron_density", "electron_parallel_flow",
                                    "electron_parallel_pressure", "electron_thermal_speed",
                                    "electron_temperature", "electron_parallel_heat_flux",
@@ -62,7 +63,8 @@ const electron_moment_variables = ("electron_density", "electron_parallel_flow",
                                    "electron_dudz", "electron_dpdz", "electron_dqdz")
 const neutral_moment_variables = ("density_neutral", "uz_neutral", "pz_neutral",
                                   "thermal_speed_neutral", "temperature_neutral",
-                                  "qz_neutral")
+                                  "qz_neutral", "total_energy_neutral",
+                                  "total_energy_flux_neutral")
 const all_moment_variables = tuple(em_variables..., ion_moment_variables...,
                                    electron_moment_variables...,
                                    neutral_moment_variables...)
@@ -4173,6 +4175,58 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         upar = get_variable(run_info, "parallel_flow"; kwargs...)
         cs = get_variable(run_info, "sound_speed"; kwargs...)
         variable = upar ./ cs
+    elseif variable_name == "total_energy"
+        if run_info.vperp.n > 1
+            error("total_energy is so far only implemented for 1D1V case")
+        else
+            ppar = get_variable(run_info, "parallel_pressure"; kwargs...)
+            upar = get_variable(run_info, "parallel_flow"; kwargs...)
+            n = get_variable(run_info, "density"; kwargs...)
+
+            variable = @. ppar + 0.5*n*upar^2
+        end
+    elseif variable_name == "total_energy_neutral"
+        if run_info.vzeta.n > 1 || run_info.vr.n > 1
+            error("total_energy_neutral is so far only implemented for 1D1V case")
+        else
+            ppar = get_variable(run_info, "pz_neutral"; kwargs...)
+            upar = get_variable(run_info, "uz_neutral"; kwargs...)
+            n = get_variable(run_info, "density_neutral"; kwargs...)
+
+            # Factor of 3/2 in front of 1/2*n*vth^2*upar because this in 1V - would be 5/2
+            # for 2V/3V cases.
+            variable = @. ppar + 0.5*n*upar^2
+        end
+    elseif variable_name == "total_energy_flux"
+        if run_info.vperp.n > 1
+            error("total_energy_flux is so far only implemented for 1D1V case")
+        else
+            qpar = get_variable(run_info, "parallel_heat_flux"; kwargs...)
+            vth = get_variable(run_info, "thermal_speed"; kwargs...)
+            upar = get_variable(run_info, "parallel_flow"; kwargs...)
+            n = get_variable(run_info, "density"; kwargs...)
+
+            # Note factor of 0.5 in front of qpar because the definition of qpar (see e.g.
+            # `update_qpar_species!()`) is unconventional (i.e. missing a factor of 0.5).
+            # Factor of 3/2 in front of 1/2*n*vth^2*upar because this in 1V - would be 5/2
+            # for 2V/3V cases.
+            variable = @. 0.5*qpar + 0.75*n*vth^2*upar + 0.5*n*upar^3
+        end
+    elseif variable_name == "total_energy_flux_neutral"
+        if run_info.vzeta.n > 1 || run_info.vr.n > 1
+            error("total_energy_flux_neutral is so far only implemented for 1D1V case")
+        else
+            qpar = get_variable(run_info, "qz_neutral"; kwargs...)
+            vth = get_variable(run_info, "thermal_speed_neutral"; kwargs...)
+            upar = get_variable(run_info, "uz_neutral"; kwargs...)
+            n = get_variable(run_info, "density_neutral"; kwargs...)
+
+            # Note factor of 0.5 in front of qpar because the definition of qpar (see e.g.
+            # `update_qpar_species!()`) is unconventional (i.e. missing a factor of 0.5).
+            # Factor of 3/2 in front of 1/2*n*vth^2*upar because this in 1V - would be 5/2
+            # for 2V/3V cases.
+            variable = @. 0.5*qpar + 0.75*n*vth^2*upar + 0.5*n*upar^3
+        end
     elseif variable_name == "z_advect_speed"
         # update_speed_z!() requires all dimensions to be present, so do *not* pass kwargs
         # to get_variable() in this case. Instead select a slice of the result.

From 37777b0cd8382843fc95be4df3ee7bb02059b1f4 Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Tue, 30 Jul 2024 16:56:01 +0000
Subject: [PATCH 365/428] Remove old example file.

---
 ...ation-slowing-down-alphas-source-sink.toml | 91 -------------------
 1 file changed, 91 deletions(-)
 delete mode 100644 examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas-source-sink.toml

diff --git a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas-source-sink.toml b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas-source-sink.toml
deleted file mode 100644
index 098dc787a..000000000
--- a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas-source-sink.toml
+++ /dev/null
@@ -1,91 +0,0 @@
-# cheap input file for a 0D2V relaxation to a collisional Maxwellian distribution with self-ion collisions and collisions with fixed Maxwellian background of cold ions and electrons.
-n_ion_species = 1
-n_neutral_species = 0
-electron_physics = "boltzmann_electron_response"
-evolve_moments_density = false
-evolve_moments_parallel_flow = false
-evolve_moments_parallel_pressure = false
-evolve_moments_conservation = false
-T_e = 1.0
-T_wall = 1.0
-initial_density1 = 1.0
-initial_temperature1 = 1.0
-z_IC_option1 = "sinusoid"
-z_IC_density_amplitude1 = 0.0
-z_IC_density_phase1 = 0.0
-z_IC_upar_amplitude1 = 0.0
-z_IC_upar_phase1 = 0.0
-z_IC_temperature_amplitude1 = 0.0
-z_IC_temperature_phase1 = 0.0
-vpa_IC_option1 = "isotropic-beam"
-vpa_IC_v01 = 1.0
-vpa_IC_vth01 = 0.1
-#vpa_IC_option1 = "directed-beam"
-#vpa_IC_vpa01 = -1.5
-#vpa_IC_vperp01 = 0.0
-charge_exchange_frequency = 0.0
-ionization_frequency = 0.0
-constant_ionization_rate = false
-
-z_ngrid = 1
-z_nelement = 1
-z_nelement_local = 1
-z_bc = "wall"
-z_discretization = "chebyshev_pseudospectral"
-r_ngrid = 1
-r_nelement = 1
-r_nelement_local = 1
-r_bc = "periodic"
-r_discretization = "chebyshev_pseudospectral"
-vpa_ngrid = 5
-vpa_nelement = 32
-vpa_L = 3.0
-vpa_bc = "zero"
-vpa_discretization = "gausslegendre_pseudospectral"
-vperp_ngrid = 5
-vperp_nelement = 16
-vperp_L = 1.5
-vperp_discretization = "gausslegendre_pseudospectral"
-vperp_bc = "zero"
-# Fokker-Planck operator requires the "gausslegendre_pseudospectral
-# options for the vpa and vperp grids
-
-[fokker_planck_collisions]
-# nuii sets the normalised input C[F,F] Fokker-Planck collision frequency
-nuii = 1.876334222e-3 #(1/nu_alphae, as computed from input diagnostic)
-Zi = 2.0
-self_collisions = true
-slowing_down_test = true
-frequency_option = "manual"
-use_fokker_planck = true
-sd_density = 1.0
-sd_temp = 0.0025 # TD/Ealpha
-sd_mi = 0.5 # mD/malpha
-sd_me = 0.000013616 # 0.25/1836.0 me/malpha
-sd_q = 1.0
-
-[ion_source]
-active=true
-source_strength=1.0
-source_T=0.005
-source_n=1.0
-r_profile="constant"
-r_width=1.0
-r_relative_minimum=0.0
-z_profile="gaussian"
-z_width=0.1
-z_relative_minimum=0.0
-source_v0 = 1.0
-#source_type="alphas"
-source_type="alphas-with-losses"
-#source_type="beam-with-losses"
-#source_vpa0 = 1.0
-#source_vperp0 = 1.0
-sink_strength = 1.0
-sink_vth=0.07071067811865475                     
-
-[timestepping]
-nstep = 50000
-dt = 5.0e-4
-nwrite = 100
-nwrite_dfns = 100

From 368302276d375f2eb5ad5b4ce697f983921ae3ff Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Tue, 30 Jul 2024 16:57:25 +0000
Subject: [PATCH 366/428] Add plot of pperp(z) at final time.

---
 moment_kinetics/src/input_structs.jl          |  2 ++
 .../src/plots_post_processing.jl              | 30 ++++++++++++++-----
 .../src/post_processing_input.jl              |  4 ++-
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index cddb11226..f59fbb049 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -553,6 +553,8 @@ struct pp_input
 	animate_parallel_flow_vs_r_z::Bool
 	# if plot_parallel_pressure_vs_r0_z = true  plot last timestep parallel_pressure[z,ir0]
 	plot_parallel_pressure_vs_r0_z::Bool
+	# if plot_perpendicular_pressure_vs_r0_z = true  plot last timestep perpendicular_pressure[z,ir0]
+	plot_perpendicular_pressure_vs_r0_z::Bool
 	# if plot_wall_parallel_pressure_vs_r = true  plot last timestep parallel_pressure[z_wall,r]
 	plot_wall_parallel_pressure_vs_r::Bool
     # if plot_parallel_pressure_vs_r_z = true plot parallel_pressure vs r z at last timestep 
diff --git a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
index 81d1e3828..8fdb8fd81 100644
--- a/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
+++ b/plots_post_processing/plots_post_processing/src/plots_post_processing.jl
@@ -3308,33 +3308,47 @@ function plot_ion_moments_2D(density, parallel_flow, parallel_pressure,
 			outfile = string(run_name, "_delta_perpendicular_pressure"*description*"(iz0,ir0)_vs_t.pdf")
 			trysavefig(outfile)
         end
+        if pp.plot_perpendicular_pressure_vs_r0_z && nz > 1 # plot last timestep perpendicular_pressure[z,ir0]
+           @views plot(z, perpendicular_pressure[:,ir0,is,end], xlabel=L"z/L_z", ylabel=L"p_{i\perp}",label="")
+           outfile = string(run_name, "_perpendicular_pressure"*description*"(r0,z)_vs_z.pdf")
+           trysavefig(outfile)
+        end
         # the total pressure
+        total_pressure = copy(parallel_pressure)
+        @. total_pressure = (2.0/3.0)*perpendicular_pressure + (1.0/3.0)*parallel_pressure
         if pp.plot_ppar0_vs_t && pp.plot_pperp0_vs_t
+            
             @views plot([time, time, time] , 
             [parallel_pressure[iz0,ir0,is,:], perpendicular_pressure[iz0,ir0,is,:], 
-            (2.0/3.0).*perpendicular_pressure[iz0,ir0,is,:] .+ (1.0/3.0).*parallel_pressure[iz0,ir0,is,:]],
-            xlabel=L"t/ (L_{ref}/c_{ref})", ylabel="", label = [L"p_{i\|\|}(t)" L"p_{i\perp}(t)" L"p_{i}(t)"])
+            total_pressure[iz0,ir0,is,:]],
+            xlabel=L"t/ (L_{ref}/c_{ref})", ylabel="", label = [L"p_{i\|\|}(t)" L"p_{i\perp}(t)" L"p_{i}(t)"],
+            foreground_color_legend = nothing, background_color_legend = nothing)
 			outfile = string(run_name, "_pressures"*description*"(iz0,ir0)_vs_t.pdf")
 			trysavefig(outfile)
             @views plot([time, time, time] , 
             [parallel_pressure[iz0,ir0,is,:] .- parallel_pressure[iz0,ir0,is,1], perpendicular_pressure[iz0,ir0,is,:] .- perpendicular_pressure[iz0,ir0,is,1], 
-            (2.0/3.0).*(perpendicular_pressure[iz0,ir0,is,:] .- perpendicular_pressure[iz0,ir0,is,1]).+
-            (1.0/3.0).*(parallel_pressure[iz0,ir0,is,:] .- parallel_pressure[iz0,ir0,is,1])],
-            xlabel=L"t/ (L_{ref}/c_{ref})", ylabel="", label = [L"p_{i\|\|}(t) - p_{i\|\|}(0)" L"p_{i\perp}(t) - p_{i\perp}(0)" L"p_{i}(t) - p_{i}(0)"])
+             total_pressure[iz0,ir0,is,:] .- total_pressure[iz0,ir0,is,1]],
+             xlabel=L"t/ (L_{ref}/c_{ref})", ylabel="", label = [L"p_{i\|\|}(t) - p_{i\|\|}(0)" L"p_{i\perp}(t) - p_{i\perp}(0)" L"p_{i}(t) - p_{i}(0)"],
+             foreground_color_legend = nothing, background_color_legend = nothing)
 			outfile = string(run_name, "_delta_pressures"*description*"(iz0,ir0)_vs_t.pdf")
 			trysavefig(outfile)
             @views plot([time] , 
-            [(2.0/3.0).*perpendicular_pressure[iz0,ir0,is,:] .+ (1.0/3.0).*parallel_pressure[iz0,ir0,is,:]],
+            [total_pressure[iz0,ir0,is,:]],
             xlabel=L"t/ (L_{ref}/c_{ref})", ylabel=L"p_{i}(t)", label = "")
 			outfile = string(run_name, "_pressure"*description*"(iz0,ir0)_vs_t.pdf")
 			trysavefig(outfile)
             @views plot([time] , 
-            [(2.0/3.0).*(perpendicular_pressure[iz0,ir0,is,:] .- perpendicular_pressure[iz0,ir0,is,1]).+
-            (1.0/3.0).*(parallel_pressure[iz0,ir0,is,:] .- parallel_pressure[iz0,ir0,is,1])],
+            [total_pressure[iz0,ir0,is,:] .- total_pressure[iz0,ir0,is,1]],
             xlabel=L"t/ (L_{ref}/c_{ref})", ylabel=L"p_{i}(t) - p_{i}(0)", label = "")
 			outfile = string(run_name, "_delta_pressure"*description*"(iz0,ir0)_vs_t.pdf")
 			trysavefig(outfile)
         end
+        if pp.plot_perpendicular_pressure_vs_r0_z && pp.plot_parallel_pressure_vs_r0_z && nz > 1 # plot last timestep perpendicular_pressure[z,ir0]
+           @views plot([z,z,z], [parallel_pressure[:,ir0,is,end],perpendicular_pressure[:,ir0,is,end],total_pressure[:,ir0,is,end]], xlabel=L"z/L_z", ylabel="",
+                       label = [L"p_{i\|\|}" L"p_{i\perp}" L"p_{i}"], foreground_color_legend = nothing, background_color_legend = nothing)
+           outfile = string(run_name, "_all_pressures"*description*"(r0,z)_vs_z.pdf")
+           trysavefig(outfile)
+        end
         # the thermal speed
         if pp.plot_vth0_vs_t
             @views plot(time, thermal_speed[iz0,ir0,is,:], xlabel=L"t / (L_{ref}/c_{ref})", ylabel=L"v_{i,th}(t)", label = "")
diff --git a/plots_post_processing/plots_post_processing/src/post_processing_input.jl b/plots_post_processing/plots_post_processing/src/post_processing_input.jl
index 5c3b6f58d..08cad55f8 100644
--- a/plots_post_processing/plots_post_processing/src/post_processing_input.jl
+++ b/plots_post_processing/plots_post_processing/src/post_processing_input.jl
@@ -102,6 +102,7 @@ const plot_wall_density_vs_r = true # plot last timestep density[z_wall,r]
 const plot_density_vs_r_z = true
 const animate_density_vs_r_z = true
 const plot_parallel_flow_vs_r0_z = true # plot last timestep parallel_flow[z,ir0]
+const plot_perpendicular_flow_vs_r0_z = true # plot last timestep perpendicular_flow[z,ir0]
 const plot_wall_parallel_flow_vs_r = true # plot last timestep parallel_flow[z_wall,r]
 const plot_parallel_flow_vs_r_z = true
 const animate_parallel_flow_vs_r_z = true
@@ -171,7 +172,8 @@ pp = pp_input(calculate_frequencies, plot_phi0_vs_t, plot_phi_vs_z_t, animate_ph
     animate_Er_vs_r_z, animate_Ez_vs_r_z, animate_phi_vs_r_z, plot_phi_vs_r0_z,
     plot_Ez_vs_r0_z, plot_wall_Ez_vs_r, plot_Er_vs_r0_z, plot_wall_Er_vs_r,
     plot_density_vs_r0_z, plot_wall_density_vs_r, plot_density_vs_r_z,
-    animate_density_vs_r_z, plot_parallel_flow_vs_r0_z, plot_wall_parallel_flow_vs_r,
+    animate_density_vs_r_z, plot_parallel_flow_vs_r0_z, plot_perpendicular_flow_vs_r0_z, 
+    plot_wall_parallel_flow_vs_r,
     plot_parallel_flow_vs_r_z, animate_parallel_flow_vs_r_z,
     plot_parallel_pressure_vs_r0_z, plot_wall_parallel_pressure_vs_r,
     plot_parallel_pressure_vs_r_z, animate_parallel_pressure_vs_r_z,

From 276d6827960fd19019ec195260e86a2aad9af63b Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:05 +0100
Subject: [PATCH 367/428] Update moment_kinetics/src/input_structs.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 moment_kinetics/src/input_structs.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 1554dbd4c..b740615c9 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -318,7 +318,6 @@ mutable struct species_composition
     # wall potential used if electron_physics=boltzmann_electron_response_with_simple_sheath
     phi_wall::mk_float
     # constant for testing nonzero Er
-    #Er_constant::mk_float
     # ratio of the neutral particle mass to the ion mass
     mn_over_mi::mk_float
     # ratio of the electron particle mass to the ion mass

From 9193a563425d9b766514663ac5fe1c6cf216e621 Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:18 +0100
Subject: [PATCH 368/428] Update
 makie_post_processing/makie_post_processing/src/shared_utils.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 makie_post_processing/makie_post_processing/src/shared_utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/makie_post_processing/makie_post_processing/src/shared_utils.jl b/makie_post_processing/makie_post_processing/src/shared_utils.jl
index f168bd457..9c7f6303d 100644
--- a/makie_post_processing/makie_post_processing/src/shared_utils.jl
+++ b/makie_post_processing/makie_post_processing/src/shared_utils.jl
@@ -105,7 +105,7 @@ function get_composition(scan_input)
     # ratio of the electron particle mass to the ion particle mass
     me_over_mi = 1.0/1836.0
     composition = species_composition(n_species, n_ion_species, n_neutral_species,
-        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, #Er_constant,
+        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall,
         mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
     return composition
 

From 185598a7b0cfb47e336386f096dc76871698ee95 Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:25 +0100
Subject: [PATCH 369/428] Update
 makie_post_processing/makie_post_processing/src/shared_utils.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 makie_post_processing/makie_post_processing/src/shared_utils.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/makie_post_processing/makie_post_processing/src/shared_utils.jl b/makie_post_processing/makie_post_processing/src/shared_utils.jl
index 9c7f6303d..32acaf1ec 100644
--- a/makie_post_processing/makie_post_processing/src/shared_utils.jl
+++ b/makie_post_processing/makie_post_processing/src/shared_utils.jl
@@ -89,7 +89,6 @@ function get_composition(scan_input)
     use_test_neutral_wall_pdf = get(scan_input, "use_test_neutral_wall_pdf", false)
     gyrokinetic_ions = get(scan_input, "gyrokinetic_ions", false)
     # constant to be used to test nonzero Er in wall boundary condition
-    #Er_constant = get(scan_input, "Er_constant", 0.0)
     recycling_fraction = get(scan_input, "recycling_fraction", 1.0)
     # constant to be used to control Ez divergences
     epsilon_offset = get(scan_input, "epsilon_offset", 0.001)

From 126cea21c1e0e55923589a9934f052d43edb6e34 Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:32 +0100
Subject: [PATCH 370/428] Update moment_kinetics/src/moment_kinetics_input.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 moment_kinetics/src/moment_kinetics_input.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 9eb9d3b25..68b2366d3 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1027,7 +1027,7 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     recycling_fraction = 1.0
     gyrokinetic_ions = false
     composition = species_composition(n_species, n_ion_species, n_neutral_species,
-        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall, #Er_constant,
+        electron_physics, use_test_neutral_wall_pdf, T_e, T_wall, phi_wall,
         mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
     
     species_ion = Array{species_parameters_mutable,1}(undef,n_ion_species)

From 8e041b7331c16503ccaadd4a4194bba07616d275 Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:40 +0100
Subject: [PATCH 371/428] Update moment_kinetics/src/moment_kinetics_input.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 moment_kinetics/src/moment_kinetics_input.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 68b2366d3..7d9f79a12 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1016,7 +1016,6 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     T_wall = 1.0
     # wall potential at z = 0
     phi_wall = 0.0
-    # constant to test nonzero Er
     #Er_constant = 0.0
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0

From 1a89877da4fe825ac2cdc0589c819d18162c7ad4 Mon Sep 17 00:00:00 2001
From: mrhardman <29800382+mrhardman@users.noreply.github.com>
Date: Wed, 31 Jul 2024 11:44:47 +0100
Subject: [PATCH 372/428] Update moment_kinetics/src/moment_kinetics_input.jl

Co-authored-by: John Omotani <john.omotani@ukaea.uk>
---
 moment_kinetics/src/moment_kinetics_input.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 7d9f79a12..a80c03751 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1016,7 +1016,6 @@ function load_defaults(n_ion_species, n_neutral_species, electron_physics)
     T_wall = 1.0
     # wall potential at z = 0
     phi_wall = 0.0
-    #Er_constant = 0.0
     # ratio of the neutral particle mass to the ion particle mass
     mn_over_mi = 1.0
     # ratio of the electron particle mass to the ion particle mass

From a6e13a981c895db011bb6ce2fe6ba54c5f8a6f5e Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Wed, 31 Jul 2024 11:05:23 +0000
Subject: [PATCH 373/428] Suggested change to dirichlet_bc option for Gauss
 Legendre, do not impose bc on lower endpoint of a cylindrical coordinate.

---
 moment_kinetics/src/gauss_legendre.jl | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/gauss_legendre.jl b/moment_kinetics/src/gauss_legendre.jl
index 539a5fd86..02eaf8b80 100644
--- a/moment_kinetics/src/gauss_legendre.jl
+++ b/moment_kinetics/src/gauss_legendre.jl
@@ -887,9 +887,12 @@ function setup_global_weak_form_matrix!(QQ_global::Array{mk_float,2},
     if dirichlet_bc
         # Make matrix diagonal for first/last grid points so it does not change the values
         # there
-        if coord.irank == 0
-            QQ_global[1,:] .= 0.0
-            QQ_global[1,1] = 1.0
+        if !(coord.name == "vperp") 
+            # modify lower endpoint if not a radial/cylindrical coordinate
+            if coord.irank == 0
+                QQ_global[1,:] .= 0.0
+                QQ_global[1,1] = 1.0
+            end
         end
         if coord.irank == coord.nrank - 1
             QQ_global[end,:] .= 0.0

From 85c2ee746669a4953d791b6dd961afbaccf1b4b2 Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Wed, 31 Jul 2024 11:34:57 +0000
Subject: [PATCH 374/428] Change names of vperp boundary conditions for
 sensible behaviour. "zero" now implies that zeros are imposed at the ends of
 the domain, "zero-impose-regularity" imposes dFdvperp = 0 at vperp = 0.

---
 ...er-planck-relaxation-slowing-down-alphas.toml |  2 +-
 ...k-slowing-down-alphas-no-self-collisions.toml |  2 +-
 ...r-planck-source-sink-slowing-down-alphas.toml |  2 +-
 moment_kinetics/src/boundary_conditions.jl       |  6 +++---
 moment_kinetics/src/fokker_planck_calculus.jl    |  4 ++--
 moment_kinetics/src/moment_kinetics_input.jl     | 10 +++++-----
 .../test/fokker_planck_time_evolution_tests.jl   | 16 ++++++++--------
 7 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
index dfab272e3..e6b75e29c 100644
--- a/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
+++ b/examples/fokker-planck/fokker-planck-relaxation-slowing-down-alphas.toml
@@ -46,7 +46,7 @@ vperp_ngrid = 5
 vperp_nelement = 16
 vperp_L = 1.5
 vperp_discretization = "gausslegendre_pseudospectral"
-vperp_bc = "zero-no-regularity"
+vperp_bc = "zero"
 # Fokker-Planck operator requires the "gausslegendre_pseudospectral
 # options for the vpa and vperp grids
 
diff --git a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
index e3070a6f7..1a95fc210 100644
--- a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
+++ b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas-no-self-collisions.toml
@@ -46,7 +46,7 @@ vperp_ngrid = 5
 vperp_nelement = 16
 vperp_L = 1.5
 vperp_discretization = "gausslegendre_pseudospectral"
-vperp_bc = "zero-no-regularity"
+vperp_bc = "zero"
 # Fokker-Planck operator requires the "gausslegendre_pseudospectral
 # options for the vpa and vperp grids
 
diff --git a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml
index 841af3a1e..cd49c82e7 100644
--- a/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml
+++ b/examples/fokker-planck/fokker-planck-source-sink-slowing-down-alphas.toml
@@ -46,7 +46,7 @@ vperp_ngrid = 5
 vperp_nelement = 16
 vperp_L = 1.5
 vperp_discretization = "gausslegendre_pseudospectral"
-vperp_bc = "zero-no-regularity"
+vperp_bc = "zero"
 # Fokker-Planck operator requires the "gausslegendre_pseudospectral
 # options for the vpa and vperp grids
 
diff --git a/moment_kinetics/src/boundary_conditions.jl b/moment_kinetics/src/boundary_conditions.jl
index 27282d65e..b014abe99 100644
--- a/moment_kinetics/src/boundary_conditions.jl
+++ b/moment_kinetics/src/boundary_conditions.jl
@@ -974,7 +974,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,5}, bc, vpe
 end
 
 function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vperp, vperp_spectral, vperp_advect, diffusion)
-    if bc == "zero" || bc == "zero-no-regularity"
+    if bc == "zero" || bc == "zero-impose-regularity"
         nvperp = vperp.n
         ngrid = vperp.ngrid
         # set zero boundary condition
@@ -984,7 +984,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vpe
             end
         end
         # set regularity condition d F / d vperp = 0 at vperp = 0
-        if bc == "zero" && (vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral")
+        if bc == "zero-impose-regularity" && (vperp.discretization == "gausslegendre_pseudospectral" || vperp.discretization == "chebyshev_pseudospectral")
             D0 = vperp_spectral.radau.D0
             buffer = @view vperp.scratch[1:ngrid-1]
             @loop_r_z_vpa ir iz ivpa begin
@@ -994,7 +994,7 @@ function enforce_vperp_boundary_condition!(f::AbstractArray{mk_float,4}, bc, vpe
                     f[ivpa,1,iz,ir] = -sum(buffer)/D0[1]
                 end
             end
-        elseif bc == "zero-no-regularity"
+        elseif bc == "zero"
             # do nothing
         else
             println("vperp.bc=\"$bc\" not supported by discretization "
diff --git a/moment_kinetics/src/fokker_planck_calculus.jl b/moment_kinetics/src/fokker_planck_calculus.jl
index 4007e6356..f11dcf43c 100644
--- a/moment_kinetics/src/fokker_planck_calculus.jl
+++ b/moment_kinetics/src/fokker_planck_calculus.jl
@@ -2299,12 +2299,12 @@ function enforce_vpavperp_BCs!(pdf,vpa,vperp,vpa_spectral,vperp_spectral)
     # set regularity condition d F / d vperp = 0 at vperp = 0
     # adjust F(vperp = 0) so that d F / d vperp = 0 at vperp = 0
     begin_anyv_vpa_region()
-    if vperp.bc in ("zero", "zero-no-regularity")
+    if vperp.bc in ("zero", "zero-impose-regularity")
         @loop_vpa ivpa begin
             pdf[ivpa,nvperp] = 0.0
         end
     end
-    if vperp.bc == "zero"
+    if vperp.bc == "zero-impose-regularity"
         buffer = @view vperp.scratch[1:ngrid_vperp-1]
         @loop_vpa ivpa begin
             @views @. buffer = D0[2:ngrid_vperp] * pdf[ivpa,2:ngrid_vperp]
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index a80c03751..3bdba79f8 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -1196,8 +1196,8 @@ function check_coordinate_input(coord, coord_name, io)
         println(io,">$coord_name.bc = 'constant'.  enforcing constant incoming BC in $coord_name.")
     elseif coord.bc == "zero"
         println(io,">$coord_name.bc = 'zero'.  enforcing zero incoming BC in $coord_name. Enforcing zero at both boundaries if diffusion operator is present.")
-    elseif coord.bc == "zero-no-regularity"
-        println(io,">$coord_name.bc = 'zero'.  enforcing zero incoming BC in $coord_name. Enforcing zero at both boundaries if diffusion operator is present. Do not enforce dF/dcoord = 0 at origin if coord = vperp.")
+    elseif coord.bc == "zero-impose-regularity"
+        println(io,">$coord_name.bc = 'zero'.  enforcing zero incoming BC in $coord_name. Enforcing zero at both boundaries if diffusion operator is present. Enforce dF/dcoord = 0 at origin if coord = vperp.")
     elseif coord.bc == "zero_gradient"
         println(io,">$coord_name.bc = 'zero_gradient'.  enforcing zero gradients at both limits of $coord_name domain.")
     elseif coord.bc == "both_zero"
@@ -1216,9 +1216,9 @@ function check_coordinate_input(coord, coord_name, io)
                 coord.nelement_global, " elements across the $coord_name domain [",
                 0.0, ",", coord.L, "].")
 
-        if coord.bc != "zero" && coord.n_global > 1 && global_rank[] == 0
-            println("WARNING: regularity condition (df/dvperp=0 at vperp=0) not being "
-                    * "imposed. Collisions or vperp-diffusion may be unstable.")
+        if coord.bc == "zero-impose-regularity" && coord.n_global > 1 && global_rank[] == 0
+            println("WARNING: regularity condition (df/dvperp=0 at vperp=0) being "
+                    * "imposed explicitly.")
         end
     else
         println(io,">using ", coord.ngrid, " grid points per $coord_name element on ",
diff --git a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
index b5d5086d6..a02fddf70 100644
--- a/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
+++ b/moment_kinetics/test/fokker_planck_time_evolution_tests.jl
@@ -34,7 +34,7 @@ struct expected_data
     f_ion::Array{mk_float, 3} # vpa, vperp, time
 end
 
-const expected_base =
+const expected_zero_impose_regularity =
   expected_data(
    [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
    [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
@@ -81,7 +81,7 @@ const expected_base =
     0.005877384425022965 0.00466291282312835 0.0028530945920350282 0.0008130106752860425 2.4399432485774198e-5 -9.743480904212476e-5 0.0;
     0.00017108987342944414 7.097261590252536e-5 -7.822316408657793e-5 -0.0001534897546375058 -9.087984332447822e-5 -3.3079379573126077e-5 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 0.0])
-const expected_no_regularity = 
+const expected_zero = 
 expected_data(
    [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],
    [0.155051025721682, 0.644948974278318, 1.000000000000000, 1.500000000000000, 2.000000000000000, 2.500000000000000, 3.000000000000000],
@@ -435,17 +435,17 @@ function runtests()
         # Benchmark data is taken from this run (GaussLegendre)
         @testset "Gauss Legendre base" begin
             run_name = "gausslegendre_pseudospectral"
-            vperp_bc = "zero"
+            vperp_bc = "zero-impose-regularity"
             run_test(test_input_gauss_legendre,
-             expected_base, 1.0e-14, 1.0e-14;
-             vperp_bc="zero")
+             expected_zero_impose_regularity, 1.0e-14, 1.0e-14;
+             vperp_bc=vperp_bc)
         end
         @testset "Gauss Legendre no enforced regularity condition at vperp = 0" begin
             run_name = "gausslegendre_pseudospectral_no_regularity"
-            vperp_bc = "zero-no-regularity"
+            vperp_bc = "zero"
             run_test(test_input_gauss_legendre,
-            expected_no_regularity,
-             1.0e-14, 1.0e-14; vperp_bc="zero-no-regularity")
+            expected_zero,
+             1.0e-14, 1.0e-14; vperp_bc=vperp_bc)
         end
         @testset "Gauss Legendre no (explicitly) enforced boundary conditions" begin
             run_name = "gausslegendre_pseudospectral_none_bc"

From 6c35c7999976f3611ef3cb15ebc11219ef0505a6 Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Wed, 31 Jul 2024 11:47:32 +0000
Subject: [PATCH 375/428] Change input option in util/precompile_run.jl

---
 util/precompile_run.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/util/precompile_run.jl b/util/precompile_run.jl
index e927edd36..43f293ee5 100644
--- a/util/precompile_run.jl
+++ b/util/precompile_run.jl
@@ -77,7 +77,7 @@ collisions_input2 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
                                                   "fokker_planck_collisions" => Dict{String,Any}("use_fokker_planck" => true, "self_collisions" => true, "slowing_down_test" => true),
                                                   "vperp_discretization" => "gausslegendre_pseudospectral",
                                                   "vpa_discretization" => "gausslegendre_pseudospectral",
-                                                  "vperp_bc" => "zero-no-regularity",
+                                                  "vperp_bc" => "zero-impose-regularity",
                                                  ))
 # add an additional input for every geometry option available in addition to the default
 geo_input1 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,

From 339213ec56d4494c06cec46886bf8ea01aee89bd Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 22:31:35 +0100
Subject: [PATCH 376/428] Ensure correct error estimate can be calculated when
 writing after a failed timestep

---
 .../src/electron_kinetic_equation.jl          |  13 +-
 moment_kinetics/src/runge_kutta.jl            |  12 +-
 moment_kinetics/src/time_advance.jl           | 249 ++++++++++++------
 3 files changed, 180 insertions(+), 94 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index e85b7f154..969bf42a6 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1357,9 +1357,16 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
         push!(total_points, z.n_global * r.n_global)
     end
 
-    adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
-                                       total_points, current_dt, error_norm_method, "",
-                                       0.0, electron=true)
+    adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
+                                       current_dt, error_norm_method, "", 0.0,
+                                       electron=true)
+    if t_params.previous_dt[] == 0.0
+        # Timestep failed, so reset  scratch[t_params.n_rk_stages+1] equal to
+        # scratch[1] to start the timestep over.
+        scratch_temp = scratch[t_params.n_rk_stages+1]
+        scratch[t_params.n_rk_stages+1] = scratch[1]
+        scratch[1] = scratch_temp
+    end
 
     return nothing
 end
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index f6912d379..da2d55ae3 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -997,7 +997,7 @@ end
 
 Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_params`.
 """
-function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
+function adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms,
                                             total_points, current_dt, error_norm_method,
                                             success, nl_max_its_fraction;
                                             electron=false)
@@ -1059,11 +1059,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
     if success != ""
         # Iteration failed in implicit part of timestep try decreasing timestep
 
-        # Set scratch[end] equal to scratch[1] to start the timestep over
-        scratch_temp = scratch[t_params.n_rk_stages+1]
-        scratch[t_params.n_rk_stages+1] = scratch[1]
-        scratch[1] = scratch_temp
-
         @serial_region begin
             t_params.failure_counter[] += 1
 
@@ -1117,11 +1112,6 @@ function adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, er
         #
         # Timestep failed, reduce timestep and re-try
 
-        # Set scratch[end] equal to scratch[1] to start the timestep over
-        scratch_temp = scratch[t_params.n_rk_stages+1]
-        scratch[t_params.n_rk_stages+1] = scratch[1]
-        scratch[1] = scratch_temp
-
         @serial_region begin
             t_params.failure_counter[] += 1
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index b6770f1bd..19038df6b 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2011,6 +2011,24 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             end
         end
 
+        if t_params.previous_dt[] == 0.0
+            # Timestep failed, so reset  scratch[t_params.n_rk_stages+1] equal to
+            # scratch[1] to start the timestep over.
+            scratch_temp = scratch[t_params.n_rk_stages+1]
+            scratch[t_params.n_rk_stages+1] = scratch[1]
+            scratch[1] = scratch_temp
+
+            # Re-update remaining velocity moments that are calculable from the evolved
+            # pdf These need to be re-calculated because `scratch[istage+1]` is now the
+            # state at the beginning of the timestep, because the timestep failed
+            apply_all_bcs_constraints_update_moments!(
+                scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
+                vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
+                collisions, geometry, gyroavs, external_source_settings, num_diss_params,
+                t_params, advance, scratch_dummy, false; pdf_bc_constraints=false,
+                update_electrons=false)
+        end
+
         if finish_now
             break
         end
@@ -2750,9 +2768,9 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         end
     end
 
-    adaptive_timestep_update_t_params!(t_params, scratch, t, CFL_limits, error_norms,
-                                       total_points, current_dt, error_norm_method,
-                                       success, nl_max_its_fraction)
+    adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
+                                       current_dt, error_norm_method, success,
+                                       nl_max_its_fraction)
 
     if composition.electron_physics == kinetic_electrons
         if t_params.previous_dt[] == 0.0
@@ -2773,18 +2791,83 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
                     pdf.electron.norm[ivpa,ivperp,iz,ir]
             end
         end
-    end
 
-    if t_params.previous_dt[] == 0.0
-        # Re-update remaining velocity moments that are calculable from the evolved
-        # pdf These need to be re-calculated because `scratch[istage+1]` is now the
-        # state at the beginning of the timestep, because the timestep failed
-        apply_all_bcs_constraints_update_moments!(
-            scratch[t_params.n_rk_stages+1], pdf, moments, fields, nothing, nothing, vz,
-            vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
-            collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-            t_params, advance, scratch_dummy, false; pdf_bc_constraints=false,
-            update_electrons=false)
+        istage = t_params.n_rk_stages+1
+
+        # update the pdf.norm and moments arrays as needed
+        begin_s_r_z_region()
+        final_scratch = scratch[istage]
+        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+            pdf.ion.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
+        end
+        @loop_s_r_z is ir iz begin
+            moments.ion.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
+            moments.ion.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
+            moments.ion.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
+            moments.ion.pperp[iz,ir,is] = final_scratch.pperp[iz,ir,is]
+        end
+        # No need to synchronize here as we only change electron quantities and previous
+        # region only changed ion quantities.
+        begin_r_z_region(no_synchronize=true)
+        @loop_r_z ir iz begin
+            moments.electron.dens[iz,ir] = final_scratch.electron_density[iz,ir]
+            moments.electron.upar[iz,ir] = final_scratch.electron_upar[iz,ir]
+            moments.electron.ppar[iz,ir] = final_scratch.electron_ppar[iz,ir]
+            moments.electron.temp[iz,ir] = final_scratch.electron_temp[iz,ir]
+        end
+        if composition.n_neutral_species > 0
+            # No need to synchronize here as we only change neutral quantities and previous
+            # region only changed plasma quantities.
+            begin_sn_r_z_region(no_synchronize=true)
+            @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
+                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = final_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn]
+            end
+            @loop_sn_r_z isn ir iz begin
+                moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
+                moments.neutral.uz[iz,ir,isn] = final_scratch.uz_neutral[iz,ir,isn]
+                moments.neutral.pz[iz,ir,isn] = final_scratch.pz_neutral[iz,ir,isn]
+            end
+            # for now update moments.neutral object directly for diagnostic moments
+            # that are not used in Runga-Kutta steps
+            update_neutral_pr!(moments.neutral.pr, moments.neutral.pr_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+            update_neutral_pzeta!(moments.neutral.pzeta, moments.neutral.pzeta_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+            # Update ptot (isotropic pressure)
+            if r.n > 1 #if 2D geometry
+                @loop_sn_r_z isn ir iz begin
+                    moments.neutral.ptot[iz,ir,isn] = (moments.neutral.pz[iz,ir,isn] + moments.neutral.pr[iz,ir,isn] + moments.neutral.pzeta[iz,ir,isn])/3.0
+                end
+            else # 1D model
+                @loop_sn_r_z isn ir iz begin
+                    moments.neutral.ptot[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
+                end
+            end
+            # get particle fluxes (n.b. bad naming convention uz -> means -> n uz here)
+            update_neutral_ur!(moments.neutral.ur, moments.neutral.ur_updated,
+                               moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z, r,
+                               composition)
+            update_neutral_uzeta!(moments.neutral.uzeta, moments.neutral.uzeta_updated,
+                                  moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z,
+                                  r, composition)
+            try #below loop can cause DomainError if ptot < 0 or density < 0, so exit cleanly if possible
+                @loop_sn_r_z isn ir iz begin
+                    # update density using last density from Runga-Kutta stages
+                    moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
+                    # get vth for neutrals
+                    moments.neutral.vth[iz,ir,isn] = sqrt(2.0*moments.neutral.ptot[iz,ir,isn]/moments.neutral.dens[iz,ir,isn])
+                end
+            catch e
+                if global_size[] > 1
+                    println("ERROR: error at line 724 of time_advance.jl")
+                    println(e)
+                    display(stacktrace(catch_backtrace()))
+                    flush(stdout)
+                    flush(stderr)
+                    MPI.Abort(comm_world, 1)
+                end
+                rethrow(e)
+            end
+        end
+
     end
 
     return nothing
@@ -2970,84 +3053,90 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
     elseif success != ""
         error("Implicit part of timestep failed")
     end
+#if global_rank[] == 0
+#    println("loworder ", scratch[2].pdf[92:95,1,1,1,1])
+#    println()
+#end
 
     reset_nonlinear_per_stage_counters(nl_solver_params.ion_advance)
     reset_nonlinear_per_stage_counters(nl_solver_params.vpa_advection)
     reset_nonlinear_per_stage_counters(nl_solver_params.maxwell_diffusion)
 
-    istage = n_rk_stages+1
+    if t_params.previous_dt[] > 0.0
+        istage = n_rk_stages+1
 
-    # update the pdf.norm and moments arrays as needed
-    begin_s_r_z_region()
-    final_scratch = scratch[istage]
-    @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
-        pdf.ion.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
-    end
-    @loop_s_r_z is ir iz begin
-        moments.ion.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
-        moments.ion.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
-        moments.ion.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
-        moments.ion.pperp[iz,ir,is] = final_scratch.pperp[iz,ir,is]
-    end
-    # No need to synchronize here as we only change electron quantities and previous
-    # region only changed ion quantities.
-    begin_r_z_region(no_synchronize=true)
-    @loop_r_z ir iz begin
-        moments.electron.dens[iz,ir] = final_scratch.electron_density[iz,ir]
-        moments.electron.upar[iz,ir] = final_scratch.electron_upar[iz,ir]
-        moments.electron.ppar[iz,ir] = final_scratch.electron_ppar[iz,ir]
-        moments.electron.temp[iz,ir] = final_scratch.electron_temp[iz,ir]
-    end
-    if composition.n_neutral_species > 0
-        # No need to synchronize here as we only change neutral quantities and previous
-        # region only changed plasma quantities.
-        begin_sn_r_z_region(no_synchronize=true)
-        @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
-            pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = final_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn]
+        # update the pdf.norm and moments arrays as needed
+        begin_s_r_z_region()
+        final_scratch = scratch[istage]
+        @loop_s_r_z_vperp_vpa is ir iz ivperp ivpa begin
+            pdf.ion.norm[ivpa,ivperp,iz,ir,is] = final_scratch.pdf[ivpa,ivperp,iz,ir,is]
         end
-        @loop_sn_r_z isn ir iz begin
-            moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
-            moments.neutral.uz[iz,ir,isn] = final_scratch.uz_neutral[iz,ir,isn]
-            moments.neutral.pz[iz,ir,isn] = final_scratch.pz_neutral[iz,ir,isn]
-        end
-        # for now update moments.neutral object directly for diagnostic moments
-        # that are not used in Runga-Kutta steps
-        update_neutral_pr!(moments.neutral.pr, moments.neutral.pr_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
-        update_neutral_pzeta!(moments.neutral.pzeta, moments.neutral.pzeta_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
-        # Update ptot (isotropic pressure)
-        if r.n > 1 #if 2D geometry
-            @loop_sn_r_z isn ir iz begin
-                moments.neutral.ptot[iz,ir,isn] = (moments.neutral.pz[iz,ir,isn] + moments.neutral.pr[iz,ir,isn] + moments.neutral.pzeta[iz,ir,isn])/3.0
-            end
-        else # 1D model
-            @loop_sn_r_z isn ir iz begin
-                moments.neutral.ptot[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
-            end
+        @loop_s_r_z is ir iz begin
+            moments.ion.dens[iz,ir,is] = final_scratch.density[iz,ir,is]
+            moments.ion.upar[iz,ir,is] = final_scratch.upar[iz,ir,is]
+            moments.ion.ppar[iz,ir,is] = final_scratch.ppar[iz,ir,is]
+            moments.ion.pperp[iz,ir,is] = final_scratch.pperp[iz,ir,is]
+        end
+        # No need to synchronize here as we only change electron quantities and previous
+        # region only changed ion quantities.
+        begin_r_z_region(no_synchronize=true)
+        @loop_r_z ir iz begin
+            moments.electron.dens[iz,ir] = final_scratch.electron_density[iz,ir]
+            moments.electron.upar[iz,ir] = final_scratch.electron_upar[iz,ir]
+            moments.electron.ppar[iz,ir] = final_scratch.electron_ppar[iz,ir]
+            moments.electron.temp[iz,ir] = final_scratch.electron_temp[iz,ir]
         end
-        # get particle fluxes (n.b. bad naming convention uz -> means -> n uz here)
-        update_neutral_ur!(moments.neutral.ur, moments.neutral.ur_updated,
-                           moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z, r,
-                           composition)
-        update_neutral_uzeta!(moments.neutral.uzeta, moments.neutral.uzeta_updated,
-                              moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z,
-                              r, composition)
-        try #below loop can cause DomainError if ptot < 0 or density < 0, so exit cleanly if possible
+        if composition.n_neutral_species > 0
+            # No need to synchronize here as we only change neutral quantities and previous
+            # region only changed plasma quantities.
+            begin_sn_r_z_region(no_synchronize=true)
+            @loop_sn_r_z_vzeta_vr_vz isn ir iz ivzeta ivr ivz begin
+                pdf.neutral.norm[ivz,ivr,ivzeta,iz,ir,isn] = final_scratch.pdf_neutral[ivz,ivr,ivzeta,iz,ir,isn]
+            end
             @loop_sn_r_z isn ir iz begin
-                # update density using last density from Runga-Kutta stages
                 moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
-                # get vth for neutrals
-                moments.neutral.vth[iz,ir,isn] = sqrt(2.0*moments.neutral.ptot[iz,ir,isn]/moments.neutral.dens[iz,ir,isn])
+                moments.neutral.uz[iz,ir,isn] = final_scratch.uz_neutral[iz,ir,isn]
+                moments.neutral.pz[iz,ir,isn] = final_scratch.pz_neutral[iz,ir,isn]
             end
-        catch e
-            if global_size[] > 1
-                println("ERROR: error at line 724 of time_advance.jl")
-                println(e)
-                display(stacktrace(catch_backtrace()))
-                flush(stdout)
-                flush(stderr)
-                MPI.Abort(comm_world, 1)
+            # for now update moments.neutral object directly for diagnostic moments
+            # that are not used in Runga-Kutta steps
+            update_neutral_pr!(moments.neutral.pr, moments.neutral.pr_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+            update_neutral_pzeta!(moments.neutral.pzeta, moments.neutral.pzeta_updated, pdf.neutral.norm, vz, vr, vzeta, z, r, composition)
+            # Update ptot (isotropic pressure)
+            if r.n > 1 #if 2D geometry
+                @loop_sn_r_z isn ir iz begin
+                    moments.neutral.ptot[iz,ir,isn] = (moments.neutral.pz[iz,ir,isn] + moments.neutral.pr[iz,ir,isn] + moments.neutral.pzeta[iz,ir,isn])/3.0
+                end
+            else # 1D model
+                @loop_sn_r_z isn ir iz begin
+                    moments.neutral.ptot[iz,ir,isn] = moments.neutral.pz[iz,ir,isn]
+                end
+            end
+            # get particle fluxes (n.b. bad naming convention uz -> means -> n uz here)
+            update_neutral_ur!(moments.neutral.ur, moments.neutral.ur_updated,
+                               moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z, r,
+                               composition)
+            update_neutral_uzeta!(moments.neutral.uzeta, moments.neutral.uzeta_updated,
+                                  moments.neutral.dens, pdf.neutral.norm, vz, vr, vzeta, z,
+                                  r, composition)
+            try #below loop can cause DomainError if ptot < 0 or density < 0, so exit cleanly if possible
+                @loop_sn_r_z isn ir iz begin
+                    # update density using last density from Runga-Kutta stages
+                    moments.neutral.dens[iz,ir,isn] = final_scratch.density_neutral[iz,ir,isn]
+                    # get vth for neutrals
+                    moments.neutral.vth[iz,ir,isn] = sqrt(2.0*moments.neutral.ptot[iz,ir,isn]/moments.neutral.dens[iz,ir,isn])
+                end
+            catch e
+                if global_size[] > 1
+                    println("ERROR: error at line 724 of time_advance.jl")
+                    println(e)
+                    display(stacktrace(catch_backtrace()))
+                    flush(stdout)
+                    flush(stderr)
+                    MPI.Abort(comm_world, 1)
+                end
+                rethrow(e)
             end
-            rethrow(e)
         end
     end
 

From c8486dbe8713d8a75af0188f6010b4d8bf894049 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 22:39:02 +0100
Subject: [PATCH 377/428] When using t_params.write_after_fixed_step_count,
 stop after nstep steps

---
 moment_kinetics/src/time_advance.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 19038df6b..c40d547be 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1812,7 +1812,9 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
         # update the time
         t += t_params.previous_dt[]
 
-        if t ≥ t_params.end_time - epsilon
+        if t ≥ t_params.end_time - epsilon ||
+                (t_params.write_after_fixed_step_count &&
+                 t_params.step_counter[] >= t_params.nstep)
             # Ensure all output is written at the final step
             finish_now = true
         elseif t_params.dt[] < 0.0 || isnan(t_params.dt[]) || isinf(t_params.dt[])

From 81d1e211b07ae5417822dd817161145890824c34 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 22:43:04 +0100
Subject: [PATCH 378/428] Set bcs on loworder pdfs using highorder moments

Ensures that there is not a point that is set to zero by the bc on one
of loworder and highorder pdfs, but not on the other causing a large
(but unhelpful) error estimate.
---
 moment_kinetics/src/time_advance.jl | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index c40d547be..1a3257147 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2613,12 +2613,30 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         end
     end
 
-    # Apply boundary conditions and constraints
+    # Apply boundary conditions and constraints to the loworder approximation.
+    # Need to apply constraints using the high-order moments for consistency, to avoid
+    # potential for spurious error estimates at boundary points.
+    loworder_constraints_scratch =
+        scratch_pdf(scratch[2].pdf, scratch[t_params.n_rk_stages+1].density,
+                    scratch[t_params.n_rk_stages+1].upar,
+                    scratch[t_params.n_rk_stages+1].ppar,
+                    scratch[t_params.n_rk_stages+1].pperp,
+                    scratch[t_params.n_rk_stages+1].temp_z_s,
+                    scratch[t_params.n_rk_stages+1].electron_density,
+                    scratch[t_params.n_rk_stages+1].electron_upar,
+                    scratch[t_params.n_rk_stages+1].electron_ppar,
+                    scratch[t_params.n_rk_stages+1].electron_pperp,
+                    scratch[t_params.n_rk_stages+1].electron_temp,
+                    scratch[2].pdf_neutral,
+                    scratch[t_params.n_rk_stages+1].density_neutral,
+                    scratch[t_params.n_rk_stages+1].uz_neutral,
+                    scratch[t_params.n_rk_stages+1].pz_neutral)
     apply_all_bcs_constraints_update_moments!(
-        scratch[2], pdf, moments, fields, boundary_distributions, scratch_electron, vz,
-        vr, vzeta, vpa, vperp, z, r, spectral_objects, advect_objects, composition,
-        collisions, geometry, gyroavs, external_source_settings, num_diss_params,
-        t_params, advance, scratch_dummy, false; update_electrons=false)
+        loworder_constraints_scratch, pdf, moments, fields, boundary_distributions,
+        scratch_electron, vz, vr, vzeta, vpa, vperp, z, r, spectral_objects,
+        advect_objects, composition, collisions, geometry, gyroavs,
+        external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
+        false; update_electrons=false)
 
     # Re-calculate moment derivatives in the `moments` struct, in case they were changed
     # by the previous call

From d8c5b7981a746c80f585543e2f9df51416c670b4 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 7 Jul 2024 22:45:39 +0100
Subject: [PATCH 379/428] Ignore error from first/last non-zero point on ion bc

This point may have a very large error, but is an isolated point and
it seems not to cause numerical instability to ignore the error there.
Avoids unnecessarily (?) small timesteps being caused by the error
estimate at that point.
---
 moment_kinetics/src/time_advance.jl | 34 ++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 1a3257147..e32c5cc69 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -27,7 +27,7 @@ using ..velocity_moments: update_neutral_pzeta!, update_neutral_pz!, update_neut
 using ..velocity_moments: calculate_ion_moment_derivatives!, calculate_neutral_moment_derivatives!
 using ..velocity_moments: calculate_electron_moment_derivatives!
 using ..velocity_grid_transforms: vzvrvzeta_to_vpavperp!, vpavperp_to_vzvrvzeta!
-using ..boundary_conditions: enforce_boundary_conditions!
+using ..boundary_conditions: enforce_boundary_conditions!, get_ion_z_boundary_cutoff_indices
 using ..boundary_conditions: enforce_neutral_boundary_conditions!
 using ..boundary_conditions: vpagrid_to_dzdt, enforce_v_boundary_condition_local!
 using ..input_structs
@@ -403,6 +403,12 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         end
     end
 
+    if t_input["implicit_vpa_advection"]
+        error("implicit_vpa_advection does not work at the moment. Need to figure out "
+              * "what to do with constraints, as explicit and implicit parts would not "
+              * "preserve constaints separately.")
+    end
+
     if t_input["high_precision_error_sum"]
         error_sum_zero = Float128(0.0)
     else
@@ -2648,6 +2654,32 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         false; pdf_bc_constraints=false, update_electrons=false)
 
     # Calculate the timstep error estimates
+    if z.bc == "wall" && (moments.evolve_upar || moments.evolve_ppar)
+        # Set error on last/first non-zero point in ion distribution function to zero, as
+        # this this point may cause unhelpful timestep failures when the cutoff moves from
+        # one point to another.
+        if z.irank == 0 || z.irank == z.nrank - 1
+            begin_s_r_region()
+            @loop_s_r is ir begin
+                density = @view scratch[t_params.n_rk_stages+1].density[:,ir,is]
+                upar = @view scratch[t_params.n_rk_stages+1].upar[:,ir,is]
+                ppar = @view scratch[t_params.n_rk_stages+1].ppar[:,ir,is]
+                last_negative_vpa_ind, first_positive_vpa_ind =
+                    get_ion_z_boundary_cutoff_indices(density, upar, ppar,
+                                                      moments.evolve_upar,
+                                                      moments.evolve_ppar, z, vpa,
+                                                      1.0e-14)
+                if z.irank == 0
+                    scratch[2].pdf[last_negative_vpa_ind,:,1,ir,is] .=
+                        scratch[t_params.n_rk_stages+1].pdf[last_negative_vpa_ind,:,1,ir,is]
+                end
+                if z.irank == z.nrank - 1
+                    scratch[2].pdf[first_positive_vpa_ind,:,end,ir,is] .=
+                        scratch[t_params.n_rk_stages+1].pdf[first_positive_vpa_ind,:,end,ir,is]
+                end
+            end
+        end
+    end
     ion_pdf_error = local_error_norm(scratch[2].pdf, scratch[t_params.n_rk_stages+1].pdf,
                                      t_params.rtol, t_params.atol;
                                      method=error_norm_method, skip_r_inner=skip_r_inner,

From 39814a3b18ffdfb334c199184136047e287e853c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 8 Jul 2024 11:03:00 +0100
Subject: [PATCH 380/428] Plot CFL limits on a log scale

---
 .../makie_post_processing/src/makie_post_processing.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 9d88f94ad..aaf299fe6 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -7457,7 +7457,7 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
             ##########################
 
             CFL_fig, ax = get_1d_ax(; xlabel="time", ylabel="(grid spacing) / speed")
-            maxval = Inf
+            #maxval = Inf
             for ri ∈ run_info
                 if length(run_info) == 1
                     prefix = ""
@@ -7496,7 +7496,7 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                 end
                 for varname ∈ CFL_vars
                     var = get_variable(ri, varname)
-                    maxval = NaNMath.min(maxval, NaNMath.maximum(var))
+                    #maxval = NaNMath.min(maxval, NaNMath.maximum(var))
                     if occursin("neutral", varname)
                         if varname ∈ implicit_CFL_vars
                             linestyle = :dashdot
@@ -7510,10 +7510,12 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                             linestyle = nothing
                         end
                     end
-                    plot_1d(time, var; ax=ax, label=prefix*electron_prefix*varname, linestyle=linestyle)
+                    plot_1d(time, var; ax=ax, label=prefix*electron_prefix*varname,
+                            linestyle=linestyle, yscale=log10,
+                            transform=x->positive_or_nan(x; epsilon=1.e-20))
                 end
             end
-            ylims!(ax, 0.0, 10.0 * maxval)
+            #ylims!(ax, 0.0, 10.0 * maxval)
             put_legend_right(CFL_fig, ax)
 
             limits_fig, ax = get_1d_ax(; xlabel="time", ylabel="number of limits per factor per output",

From df396b6b5f9a03f10ffa4e6fa3743998c99629e9 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 12 Jul 2024 09:00:39 +0100
Subject: [PATCH 381/428] Set number of BLAS threads to 1 to avoid
 oversubcribing processors

Allowing BLAS to use multiple threads when we are already using all
available processes causes oversubscription and significantly slows down
the code. Therefore pin the number of BLAS threads to 1.
---
 moment_kinetics/src/communication.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/moment_kinetics/src/communication.jl b/moment_kinetics/src/communication.jl
index 2f88c8cd9..be3ed60a3 100644
--- a/moment_kinetics/src/communication.jl
+++ b/moment_kinetics/src/communication.jl
@@ -20,6 +20,7 @@ export setup_distributed_memory_MPI
 export setup_distributed_memory_MPI_for_weights_precomputation
 export _block_synchronize, _anyv_subblock_synchronize
 
+using LinearAlgebra
 using MPI
 using SHA
 
@@ -127,6 +128,10 @@ function __init__()
     global_size[] = MPI.Comm_size(comm_world)
     #block_rank[] = MPI.Comm_rank(comm_block)
     #block_size[] = MPI.Comm_size(comm_block)
+
+    # Ensure BLAS only uses 1 thread, to avoid oversubscribing processes as we are
+    # probably already fully parallelised.
+    BLAS.set_num_threads(1)
 end
 
 """

From 4eb97f9ff456da251dea3fbd5dbcd4bf88527457 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 16 Jul 2024 20:14:26 +0100
Subject: [PATCH 382/428] Fix Jacobian-vector product for updated error
 tolerances

Now that `distributed_norm()`, etc. include error tolerances, vectors
normalised to '1.0' are actually very small, so instead of doing
`x + epsilon * v` for a small `epsilon`, we should do
`x + Jv_scale_factor * v` for a large-ish `Jv_scale_factor`. Otherwise
`x + epsilon * v` would be so small that rounding errors are large
relative errors on the estimate of `J.v`, which would prevent the Newton
solver from converging.
---
 moment_kinetics/src/nonlinear_solvers.jl | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl
index ac8f3b3e1..f35f0ade0 100644
--- a/moment_kinetics/src/nonlinear_solvers.jl
+++ b/moment_kinetics/src/nonlinear_solvers.jl
@@ -83,7 +83,7 @@ function setup_nonlinear_solve(input_dict, coords, outer_coords=(); default_rtol
         atol=default_atol,
         nonlinear_max_iterations=20,
         linear_rtol=1.0e-3,
-        linear_atol=1.0e-15,
+        linear_atol=1.0,
         linear_restart=10,
         linear_max_restarts=0,
         preconditioner_update_interval=300,
@@ -1040,16 +1040,25 @@ function linear_solve!(x, residual_func!, residual0, delta_x, v, w; coords, rtol
     # Solve (approximately?):
     #   J δx = residual0
 
-    tol = max(rtol, atol)
-    epsilon = 1.0e-6 / tol
-    inv_epsilon = 1.0 / epsilon
-
+    Jv_scale_factor = 1.0e3
+    inv_Jv_scale_factor = 1.0 / Jv_scale_factor
+
+    # The vectors `v` that are passed to this function will be normalised so that
+    # `distributed_norm(v) == 1.0`. `distributed_norm()` is defined - including the
+    # relative and absolute tolerances from the Newton iteration - so that a vector with a
+    # norm of 1.0 is 'small' in the sense that a vector with a norm of 1.0 is small enough
+    # relative to `x` to consider the iteration converged. This means that `x+v` would be
+    # very close to `x`, so R(x+v)-R(x) would be likely to be badly affected by rounding
+    # errors, because `v` is so small, relative to `x`. We actually want to multiply `v`
+    # by a large number `Jv_scale_factor` (in constrast to the small `epsilon` in the
+    # 'usual' case where the norm does not include either reative or absolute tolerance)
+    # to ensure that we get a reasonable estimate of J.v.
     function approximate_Jacobian_vector_product!(v)
         right_preconditioner(v)
 
-        parallel_map((x,v) -> x + epsilon * v, v, x, v)
+        parallel_map((x,v) -> x + Jv_scale_factor * v, v, x, v)
         residual_func!(rhs_delta, v)
-        parallel_map((rhs_delta, residual0) -> (rhs_delta - residual0) * inv_epsilon,
+        parallel_map((rhs_delta, residual0) -> (rhs_delta - residual0) * inv_Jv_scale_factor,
                      v, rhs_delta, residual0)
         left_preconditioner(v)
         return v

From adee5d73d3a540420a4ae693008f89bebb3a76de Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 16 Jul 2024 20:19:11 +0100
Subject: [PATCH 383/428] Fix "Fix post-processing for failure caused by
 kinetic electrons diagnostic"

---
 moment_kinetics/src/electron_kinetic_equation.jl | 2 +-
 moment_kinetics/src/runge_kutta.jl               | 3 ++-
 moment_kinetics/src/time_advance.jl              | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 969bf42a6..4629d9827 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1359,7 +1359,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
 
     adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
                                        current_dt, error_norm_method, "", 0.0,
-                                       electron=true)
+                                       composition; electron=true)
     if t_params.previous_dt[] == 0.0
         # Timestep failed, so reset  scratch[t_params.n_rk_stages+1] equal to
         # scratch[1] to start the timestep over.
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index da2d55ae3..24f396699 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -9,6 +9,7 @@ export setup_runge_kutta_coefficients!, rk_update_evolved_moments!,
 
 using ..array_allocation: allocate_float
 using ..communication
+using ..input_structs
 using ..looping
 using ..type_definitions: mk_float
 
@@ -999,7 +1000,7 @@ Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_p
 """
 function adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms,
                                             total_points, current_dt, error_norm_method,
-                                            success, nl_max_its_fraction;
+                                            success, nl_max_its_fraction, composition;
                                             electron=false)
     # Get global minimum of CFL limits
     CFL_limit = nothing
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index e32c5cc69..199283520 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2822,7 +2822,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
 
     adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
                                        current_dt, error_norm_method, success,
-                                       nl_max_its_fraction)
+                                       nl_max_its_fraction, composition)
 
     if composition.electron_physics == kinetic_electrons
         if t_params.previous_dt[] == 0.0

From ac7f0a355820c724ded567a9ede5a56140da8d1f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 16 Jul 2024 20:37:13 +0100
Subject: [PATCH 384/428] Normalise grid to initial_temperature while
 initialising f_neutral

...in simulations with `z.bc = "wall"`. This makes the initialition
(integrations, etc.) more accurate when the initial temperature is much
smaller (or larger) than the reference temperature, which should help
especially for moment-kinetic cases.
---
 moment_kinetics/src/initial_conditions.jl | 102 +++++++++++++++++-----
 1 file changed, 78 insertions(+), 24 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index e1025217b..aeffb8ef3 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -1272,6 +1272,8 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
         vz, vr, vzeta, z, vz_spectral, density, uz, pz, vth, v_norm_fac, evolve_density,
         evolve_upar, evolve_ppar, wall_flux_0, wall_flux_L)
 
+    zero = 1.0e-14
+
     # Reduce the ion flux by `recycling_fraction` to account for ions absorbed by the
     # wall.
     wall_flux_0 *= composition.recycling_fraction
@@ -1338,6 +1340,11 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
                 end
             end
         else
+            # Normalise grid to thermal speed of `initial_temperature`, to avoid
+            # inaccuracy in moment-kinetic cases when `initial_temperature` is very
+            # small or large compared to the reference value.
+            vth_init = sqrt(spec.initial_temperature)
+
             # First create distribution functions at the z-boundary points that obey the
             # boundary conditions.
             if z.irank == 0 && z.irank == z.nrank - 1
@@ -1358,7 +1365,7 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
                     @. pdf[:,ivr,ivzeta,iz] = density[iz] *
                                               exp(-((vz.grid - uz[iz])^2 +
                                                     vzeta.grid[ivzeta]^2 + vr.grid[ivr]^2)
-                                                  / vth[iz]^2) / vth[iz]
+                                                  * vth_init^2 / vth[iz]^2) / vth[iz]
 
                     # Also ensure both species go to zero smoothly at v_z=0 at the
                     # wall, where the boundary conditions require that distribution
@@ -1398,16 +1405,49 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
             @views MPI.Bcast!(temp, z.nrank - 1, z.comm)
             wall_flux_L = temp[]
 
-            knudsen_pdf = boundary_distributions.knudsen
+            # Re-calculate Knudsen distribution instead of using
+            # `boundary_distributions.knudsen`, so that we can include vth_init here.
+            knudsen_pdf = allocate_float(vz.n, vr.n, vzeta.n)
+            knudsen_vtfac = sqrt(composition.T_wall * composition.mn_over_mi)
+            if vzeta.n > 1 && vr.n > 1
+                # 3V specification of neutral wall emission distribution for boundary condition
+                # get the true Knudsen cosine distribution for neutral particle wall emission
+                for ivzeta in 1:vzeta.n
+                    for ivr in 1:vr.n
+                        for ivz in 1:vz.n
+                            v_transverse = sqrt(vzeta.grid[ivzeta]^2 + vr.grid[ivr]^2)
+                            v_normal = abs(vz.grid[ivz])
+                            v_tot = sqrt(v_normal^2 + v_transverse^2)
+                            if  v_tot > zero
+                                prefac = v_normal/v_tot
+                            else
+                                prefac = 0.0
+                            end
+                            knudsen_pdf[ivz,ivr,ivzeta] = (3.0*sqrt(pi)/knudsen_vtfac^4) * prefac *
+                                                             exp(-((v_normal/knudsen_vtfac)^2 + (v_transverse/knudsen_vtfac)^2) * vth_init^2)
+                        end
+                    end
+                end
+            elseif vzeta.n == 1 && vr.n == 1
+                # get the marginalised Knudsen cosine distribution after integrating over
+                # vperp appropriate for 1V model
+                @. knudsen_pdf[:,1,1] = (3.0*pi/knudsen_vtfac^3)*abs(vz.grid*vth_init)*erfc(abs(vz.grid) * vth_init / knudsen_vtfac)
+            end
 
-            zero = 1.0e-14
+            if vzeta.n > 1 || vr.n > 1
+                wgts_3V_vth_init = vth_init
+            else
+                wgts_3V_vth_init = 1.0
+            end
 
             # add this species' contribution to the combined ion/neutral particle flux
             # out of the domain at z=-Lz/2
             @views wall_flux_0 += integrate_over_negative_vz(
-                                      abs.(vz.grid) .* lower_z_pdf_buffer, vz.grid, vz.wgts,
-                                      vz.scratch3, vr.grid, vr.wgts, vzeta.grid,
-                                      vzeta.wgts)
+                                      vth_init .* abs.(vz.grid) .* lower_z_pdf_buffer,
+                                      vth_init .* vz.grid, vth_init .* vz.wgts,
+                                      vz.scratch3, vth_init .* vr.grid,
+                                      wgts_3V_vth_init .* vr.wgts, vth_init .* vzeta.grid,
+                                      wgts_3V_vth_init .* vzeta.wgts)
             # for left boundary in zed (z = -Lz/2), want
             # f_n(z=-Lz/2, v_z > 0) = Γ_0 * f_KW(v_z) * pdf_norm_fac(-Lz/2)
             @loop_vz ivz begin
@@ -1419,9 +1459,11 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
             # add this species' contribution to the combined ion/neutral particle flux
             # out of the domain at z=-Lz/2
             @views wall_flux_L += integrate_over_positive_vz(
-                                      abs.(vz.grid) .* upper_z_pdf_buffer, vz.grid, vz.wgts,
-                                      vz.scratch3, vr.grid, vr.wgts, vzeta.grid,
-                                      vzeta.wgts)
+                                      vth_init .* abs.(vz.grid) .* upper_z_pdf_buffer,
+                                      vth_init .* vz.grid, vth_init .* vz.wgts,
+                                      vz.scratch3, vth_init .* vr.grid,
+                                      wgts_3V_vth_init .* vr.wgts, vth_init .* vzeta.grid,
+                                      wgts_3V_vth_init .* vzeta.wgts)
             # for right boundary in zed (z = Lz/2), want
             # f_n(z=Lz/2, v_z < 0) = Γ_Lz * f_KW(v_z) * pdf_norm_fac(Lz/2)
             @loop_vz ivz begin
@@ -1451,8 +1493,9 @@ function init_neutral_pdf_over_density!(pdf, boundary_distributions, spec, compo
             # Get the unnormalised pdf and the moments of the constructed full-f
             # distribution function (which will be modified from the input moments).
             convert_full_f_neutral_to_normalised!(pdf, density, uz, pz, vth, vzeta, vr,
-                                                  vz, vz_spectral, evolve_density,
-                                                  evolve_upar, evolve_ppar)
+                                                  vz, vz_spectral, vth_init,
+                                                  evolve_density, evolve_upar,
+                                                  evolve_ppar)
 
             if !evolve_density
                 # Need to divide out density to return pdf/density
@@ -1827,19 +1870,30 @@ the moments of `f`.
 Inputs/outputs depend on z, vzeta, vr and vz (should be inside loops over species, r)
 """
 function convert_full_f_neutral_to_normalised!(f, density, uz, pz, vth, vzeta, vr, vz,
-        vz_spectral, evolve_density, evolve_upar, evolve_ppar)
+        vz_spectral, vth_init, evolve_density, evolve_upar, evolve_ppar)
 
+    if vzeta.n > 1 || vr.n > 1
+        wgts_3V_vth_init = vth_init
+    else
+        wgts_3V_vth_init = 1.0
+    end
     @loop_z iz begin
         # Calculate moments
-        @views density[iz] = integrate_over_neutral_vspace(f[:,:,:,iz], vz.grid, 0,
-                                                           vz.wgts, vr.grid, 0, vr.wgts,
-                                                           vzeta.grid, 0, vzeta.wgts)
-        @views uz[iz] = integrate_over_neutral_vspace(f[:,:,:,iz], vz.grid, 1, vz.wgts,
-                                                      vr.grid, 0, vr.wgts, vzeta.grid, 0,
-                                                      vzeta.wgts) / density[iz]
-        @views pz[iz] = integrate_over_neutral_vspace(f[:,:,:,iz], vz.grid, 2, vz.wgts,
-                                                      vr.grid, 0, vr.wgts, vzeta.grid, 0,
-                                                      vzeta.wgts) - density[iz]*uz[iz]^2
+        @views density[iz] = integrate_over_neutral_vspace(
+                                 f[:,:,:,iz], vth_init .* vth_init .* vz.grid, 0,
+                                 vth_init .* vz.wgts, vth_init .* vr.grid, 0,
+                                 wgts_3V_vth_init .* vr.wgts, vth_init .* vzeta.grid, 0,
+                                 wgts_3V_vth_init .* vzeta.wgts)
+        @views uz[iz] = integrate_over_neutral_vspace(
+                            f[:,:,:,iz], vth_init .* vz.grid, 1, vth_init .* vz.wgts,
+                            vth_init .* vr.grid, 0, wgts_3V_vth_init .* vr.wgts,
+                            vth_init .* vzeta.grid, 0, wgts_3V_vth_init .* vzeta.wgts) /
+                        density[iz]
+        @views pz[iz] = integrate_over_neutral_vspace(
+                            f[:,:,:,iz], vth_init .* vz.grid, 2, vth_init .* vz.wgts,
+                            vth_init .* vr.grid, 0, wgts_3V_vth_init .* vr.wgts,
+                            vth_init .* vzeta.grid, 0, wgts_3V_vth_init .* vzeta.wgts) -
+                        density[iz]*uz[iz]^2
         vth[iz] = sqrt(2.0*pz[iz]/density[iz])
 
         # Normalise f
@@ -1854,11 +1908,11 @@ function convert_full_f_neutral_to_normalised!(f, density, uz, pz, vth, vzeta, v
             # The values to interpolate *to* are the v_parallel values corresponding to
             # the w_parallel grid
             vz.scratch .= vpagrid_to_dzdt(vz.grid, vth[iz], uz[iz], evolve_ppar,
-                                          evolve_upar)
+                                          evolve_upar) ./ vth_init
             @loop_vzeta_vr ivzeta ivr begin
                 @views vz.scratch2 .= f[:,ivr,ivzeta,iz] # Copy to use as input to interpolation
-                @views interpolate_to_grid_1d!(f[:,ivr,ivzeta,iz], vz.scratch, vz.scratch2,
-                                               vz, vz_spectral)
+                @views interpolate_to_grid_1d!(f[:,ivr,ivzeta,iz], vz.scratch,
+                                               vz.scratch2, vz, vz_spectral)
             end
         end
     end

From c277d2a912d49a5e3b07fb7afe9f5796d78adbba Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 17 Jul 2024 13:37:36 +0100
Subject: [PATCH 385/428] Fix neutral wall plots for case when vz is not
 identical to vpa

---
 .../src/makie_post_processing.jl                     | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index aaf299fe6..e0abcffea 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -3873,15 +3873,17 @@ function plot_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing, neutra
         upar = get_variable(run_info, "uz_neutral"; it=it, is=is, ir=input.ir0, iz=iz)
         vth = get_variable(run_info, "thermal_speed_neutral"; it=it, is=is, ir=input.ir0,
                            iz=iz)
+        vcoord = run_info.vz
     else
         f = get_variable(run_info, "f"; it=it, is=is, ir=input.ir0, iz=iz,
                          ivperp=input.ivperp0)
         density = get_variable(run_info, "density"; it=it, is=is, ir=input.ir0, iz=iz)
         upar = get_variable(run_info, "parallel_flow"; it=it, is=is, ir=input.ir0, iz=iz)
         vth = get_variable(run_info, "thermal_speed"; it=it, is=is, ir=input.ir0, iz=iz)
+        vcoord = run_info.vpa
     end
 
-    f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(f, run_info.vpa.grid, density, upar, vth,
+    f_unnorm, dzdt = get_unnormalised_f_dzdt_1d(f, vcoord.grid, density, upar, vth,
                                                 run_info.evolve_density,
                                                 run_info.evolve_upar,
                                                 run_info.evolve_ppar)
@@ -4205,6 +4207,7 @@ function animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing,
         density = get_variable(run_info, "density_neutral"; is=is, ir=input.ir0, iz=iz)
         upar = get_variable(run_info, "uz_neutral"; is=is, ir=input.ir0, iz=iz)
         vth = get_variable(run_info, "thermal_speed_neutral"; is=is, ir=input.ir0, iz=iz)
+        vcoord = run_info.vz
     else
         f = VariableCache(run_info, "f", chunk_size_2d; it=nothing, is=is, ir=input.ir0, iz=iz,
                           ivperp=input.ivperp0, ivpa=nothing, ivzeta=nothing, ivr=nothing,
@@ -4212,6 +4215,7 @@ function animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing,
         density = get_variable(run_info, "density"; is=is, ir=input.ir0, iz=iz)
         upar = get_variable(run_info, "parallel_flow"; is=is, ir=input.ir0, iz=iz)
         vth = get_variable(run_info, "thermal_speed"; is=is, ir=input.ir0, iz=iz)
+        vcoord = run_info.vpa
     end
 
     function get_this_f_unnorm(it)
@@ -4219,7 +4223,7 @@ function animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing,
                                          run_info.evolve_density, run_info.evolve_ppar)
 
         if f_over_vpa2
-            dzdt = vpagrid_to_dzdt(run_info.vpa.grid, vth[it], upar[it],
+            dzdt = vpagrid_to_dzdt(vcoord.grid, vth[it], upar[it],
                                    run_info.evolve_ppar, run_info.evolve_upar)
             dzdt2 = dzdt.^2
             for i ∈ eachindex(dzdt2)
@@ -4240,7 +4244,7 @@ function animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing,
     fmin = Inf
     fmax = -Inf
     for it ∈ 1:run_info.nt
-        this_dzdt = vpagrid_to_dzdt(run_info.vpa.grid, vth[it], upar[it],
+        this_dzdt = vpagrid_to_dzdt(vcoord.grid, vth[it], upar[it],
                                     run_info.evolve_ppar, run_info.evolve_upar)
         this_dzdtmin, this_dzdtmax = extrema(this_dzdt)
         dzdtmin = min(dzdtmin, this_dzdtmin)
@@ -4264,7 +4268,7 @@ function animate_f_unnorm_vs_vpa(run_info; f_over_vpa2=false, input=nothing,
                 fmin - 0.01*yheight, fmax + 0.01*yheight)
     end
 
-    dzdt = @lift vpagrid_to_dzdt(run_info.vpa.grid, vth[$frame_index], upar[$frame_index],
+    dzdt = @lift vpagrid_to_dzdt(vcoord.grid, vth[$frame_index], upar[$frame_index],
                                  run_info.evolve_ppar, run_info.evolve_upar)
     f_unnorm = @lift transform.(get_this_f_unnorm($frame_index))
 

From 615b25ed2c51f7a82d2bbe603edea39a9bb43fe7 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 18 Jul 2024 15:24:32 +0100
Subject: [PATCH 386/428] Allow different vpa and vz grids in 1V simulations

Add a couple of calls to interpolate between ion/neutral grids so that
1V simulations can use different grids for vpa and vz. This may be
useful when ions and neutrals have significantly different temperatures.
---
 moment_kinetics/src/charge_exchange.jl | 48 ++++++++++++++++++--------
 moment_kinetics/src/ionization.jl      | 27 ++++++++-------
 moment_kinetics/src/time_advance.jl    | 16 +++------
 3 files changed, 52 insertions(+), 39 deletions(-)

diff --git a/moment_kinetics/src/charge_exchange.jl b/moment_kinetics/src/charge_exchange.jl
index e70782c8c..ac460dba0 100644
--- a/moment_kinetics/src/charge_exchange.jl
+++ b/moment_kinetics/src/charge_exchange.jl
@@ -30,7 +30,7 @@ function ion_charge_exchange_collisions_1V!(f_out, fvec_in, moments, composition
                 fvec_in.density_neutral[:,:,is], fvec_in.upar[:,:,is],
                 fvec_in.uz_neutral[:,:,is], moments.ion.vth[:,:,is],
                 moments.neutral.vth[:,:,is], moments, vpa, vz, charge_exchange_frequency,
-                vz_spectral, dt)
+                vz_spectral, dt; neutrals=false)
         end
     else
         begin_s_r_z_region()
@@ -38,11 +38,17 @@ function ion_charge_exchange_collisions_1V!(f_out, fvec_in, moments, composition
             # apply CX collisions to all ion species
             # for each ion species, obtain affect of charge exchange collisions
             # with the corresponding neutral species
-            @loop_r_z_vpa ir iz ivpa begin
-                f_out[ivpa,1,iz,ir,is] +=
-                    dt*charge_exchange_frequency*(
-                        fvec_in.pdf_neutral[ivpa,1,1,iz,ir,is]*fvec_in.density[iz,ir,is]
-                        - fvec_in.pdf[ivpa,1,iz,ir,is]*fvec_in.density_neutral[iz,ir,is])
+            isn = is
+            @loop_r_z ir iz begin
+                @views interpolate_to_grid_vpa!(vpa.scratch, vpa.grid,
+                                                fvec_in.pdf_neutral[:,1,1,iz,ir,isn], vz,
+                                                vz_spectral)
+                @loop_vpa ivpa begin
+                    f_out[ivpa,1,iz,ir,is] +=
+                        dt*charge_exchange_frequency*(
+                            vpa.scratch[ivpa]*fvec_in.density[iz,ir,is]
+                            - fvec_in.pdf[ivpa,1,iz,ir,is]*fvec_in.density_neutral[iz,ir,is])
+                end
             end
         end
     end
@@ -70,7 +76,7 @@ function neutral_charge_exchange_collisions_1V!(f_neutral_out, fvec_in, moments,
                 fvec_in.pdf[:,1,:,:,isn], fvec_in.density[:,:,isn],
                 fvec_in.uz_neutral[:,:,isn], fvec_in.upar[:,:,isn],
                 moments.neutral.vth[:,:,isn], moments.ion.vth[:,:,isn], moments,
-                vz, vpa, charge_exchange_frequency, vpa_spectral, dt)
+                vz, vpa, charge_exchange_frequency, vpa_spectral, dt; neutrals=true)
         end
     else
         begin_sn_r_z_region()
@@ -78,11 +84,16 @@ function neutral_charge_exchange_collisions_1V!(f_neutral_out, fvec_in, moments,
             # apply CX collisions to all neutral species
             # for each neutral species, obtain affect of charge exchange collisions
             # with the corresponding ion species
-            @loop_r_z_vz ir iz ivz begin
-                f_neutral_out[ivz,1,1,iz,ir,isn] +=
-                    dt*charge_exchange_frequency*(
-                        fvec_in.pdf[ivz,1,iz,ir,isn]*fvec_in.density_neutral[iz,ir,isn]
-                        - fvec_in.pdf_neutral[ivz,1,1,iz,ir,isn]*fvec_in.density[iz,ir,isn])
+            @loop_r_z ir iz begin
+                @views interpolate_to_grid_vpa!(vz.scratch, vz.grid,
+                                                fvec_in.pdf[:,1,iz,ir,isn], vpa,
+                                                vpa_spectral)
+                @loop_vz ivz begin
+                    f_neutral_out[ivz,1,1,iz,ir,isn] +=
+                        dt*charge_exchange_frequency*(
+                            vz.scratch[ivz]*fvec_in.density_neutral[iz,ir,isn]
+                            - fvec_in.pdf_neutral[ivz,1,1,iz,ir,isn]*fvec_in.density[iz,ir,isn])
+                end
             end
         end
     end
@@ -94,7 +105,7 @@ with a single species of the opposite type; e.g., ions with neutrals or neutrals
 """
 function charge_exchange_collisions_single_species!(f_out, pdf_in, pdf_other,
         density_other, upar, upar_other, vth, vth_other, moments, vpa, vpa_other,
-        charge_exchange_frequency, spectral_other, dt)
+        charge_exchange_frequency, spectral_other, dt; neutrals)
     @loop_r_z ir iz begin
         if moments.evolve_ppar
             # will need the ratio of thermal speeds both to interpolate between vpa grids
@@ -142,9 +153,16 @@ function charge_exchange_collisions_single_species!(f_out, pdf_in, pdf_other,
             # no need to interpolate if neither upar or ppar evolved separately from pdf
             vpa.scratch2 .= pdf_other[:,iz,ir]
         end
-        @loop_vpa ivpa begin
-            f_out[ivpa,iz,ir] += dt * charge_exchange_frequency * density_other[iz,ir] *
+        if neutrals
+            @loop_vz ivz begin
+                f_out[ivz,iz,ir] += dt * charge_exchange_frequency * density_other[iz,ir] *
+                (vpa.scratch2[ivz] * vth_ratio - pdf_in[ivz,iz,ir])
+            end
+        else
+            @loop_vpa ivpa begin
+                f_out[ivpa,iz,ir] += dt * charge_exchange_frequency * density_other[iz,ir] *
                 (vpa.scratch2[ivpa] * vth_ratio - pdf_in[ivpa,iz,ir])
+            end
         end
     end
 end
diff --git a/moment_kinetics/src/ionization.jl b/moment_kinetics/src/ionization.jl
index babdd5d9b..9ea348da0 100644
--- a/moment_kinetics/src/ionization.jl
+++ b/moment_kinetics/src/ionization.jl
@@ -78,8 +78,7 @@ function ion_ionization_collisions_1V!(f_out, fvec_in, vz, vpa, vperp, z, r, vz_
     @boundscheck r.n == size(f_out,4) || throw(BoundsError(f_out))
     @boundscheck composition.n_ion_species == size(f_out,5) || throw(BoundsError(f_out))
     
-    
-    begin_r_z_vpa_region()
+    begin_r_z_region()
 
     if moments.evolve_density
         # For now, assume species index `is` corresponds to the neutral
@@ -146,9 +145,14 @@ function ion_ionization_collisions_1V!(f_out, fvec_in, vz, vpa, vperp, z, r, vz_
             #NB: used quasineutrality to replace electron density n_e with ion density
             #NEEDS GENERALISATION TO n_ion_species > 1 (missing species charge: Sum_i Z_i n_i = n_e)
             isn = is
-            @loop_r_z_vpa ir iz ivpa begin
-                # apply ionization collisions to all ion species
-                f_out[ivpa,1,iz,ir,is] += dt*collisions.ionization*fvec_in.pdf_neutral[ivpa,1,1,iz,ir,isn]*fvec_in.density[iz,ir,is]
+            @loop_r_z ir iz begin
+                @views interpolate_to_grid_vpa!(vpa.scratch, vpa.grid,
+                                                fvec_in.pdf_neutral[:,1,1,iz,ir,isn], vz,
+                                                vz_spectral)
+                @loop_vpa ivpa begin
+                    # apply ionization collisions to all ion species
+                    f_out[ivpa,1,iz,ir,is] += dt*collisions.ionization*vpa.scratch[ivpa]*fvec_in.density[iz,ir,is]
+                end
             end
         end
     end
@@ -157,10 +161,9 @@ end
 function neutral_ionization_collisions_1V!(f_neutral_out, fvec_in, vz, vpa, vperp, z, r,
                                            vz_spectral, moments, composition, collisions, dt)
     # This routine assumes a 1D model with:
-    # nvz = nvpa and identical vz and vpa grids
     # nvperp = nvr = nveta = 1
     # constant charge_exchange_frequency independent of species
-    @boundscheck vpa.n == size(f_neutral_out,1) || throw(BoundsError(f_neutral_out))
+    @boundscheck vz.n == size(f_neutral_out,1) || throw(BoundsError(f_neutral_out))
     @boundscheck 1 == size(f_neutral_out,2) || throw(BoundsError(f_neutral_out))
     @boundscheck 1 == size(f_neutral_out,3) || throw(BoundsError(f_neutral_out))
     @boundscheck z.n == size(f_neutral_out,4) || throw(BoundsError(f_neutral_out))
@@ -168,18 +171,18 @@ function neutral_ionization_collisions_1V!(f_neutral_out, fvec_in, vz, vpa, vper
     @boundscheck composition.n_neutral_species == size(f_neutral_out,6) || throw(BoundsError(f_neutral_out))
 
     if !moments.evolve_density
-        begin_r_z_vpa_region()
+        begin_sn_r_z_vz_region()
 
-        @loop_s is begin
+        @loop_sn isn begin
             # ion ionisation rate =   < f_n > n_e R_ion
             # neutral "ionisation" (depopulation) rate =   -  f_n  n_e R_ion
             # no gyroaverage here as 1V code
             #NB: used quasineutrality to replace electron density n_e with ion density
             #NEEDS GENERALISATION TO n_ion_species > 1 (missing species charge: Sum_i Z_i n_i = n_e)
-            isn = is
-            @loop_r_z_vpa ir iz ivpa begin
+            is = isn
+            @loop_r_z_vz ir iz ivz begin
                 # apply ionization collisions to all neutral species
-                f_neutral_out[ivpa,1,1,iz,ir,isn] -= dt*collisions.ionization*fvec_in.pdf_neutral[ivpa,1,1,iz,ir,isn]*fvec_in.density[iz,ir,is]
+                f_neutral_out[ivz,1,1,iz,ir,isn] -= dt*collisions.ionization*fvec_in.pdf_neutral[ivz,1,1,iz,ir,isn]*fvec_in.density[iz,ir,is]
             end
         end
     end
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 199283520..1f3f8189f 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1158,7 +1158,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
             # if charge exchange collision frequency non-zero,
             # account for charge exchange collisions
             if abs(collisions.charge_exchange) > 0.0
-                if vz.n == vpa.n && vperp.n == 1 && vr.n == 1 && vzeta.n == 1
+                if vperp.n == 1 && vr.n == 1 && vzeta.n == 1
                     advance_ion_cx_1V = !t_params.implicit_ion_advance
                     advance_neutral_cx_1V = true
                 elseif vperp.n > 1 && vr.n > 1 && vzeta.n > 1
@@ -1166,8 +1166,6 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                     advance_neutral_cx = true
                 else
                     error("If any perpendicular velocity has length>1 they all must. "
-                          * "If all perpendicular velocities have length=1, then vpa and "
-                          * "vz should be the same.\n"
                           * "vperp.n=$(vperp.n), vr.n=$(vr.n), vzeta.n=$(vzeta.n), "
                           * "vpa.n=$(vpa.n), vz.n=$(vz.n)")
                 end
@@ -1175,7 +1173,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
             # if ionization collision frequency non-zero,
             # account for ionization collisions
             if abs(collisions.ionization) > 0.0
-                if vz.n == vpa.n && vperp.n == 1 && vr.n == 1 && vzeta.n == 1
+                if vperp.n == 1 && vr.n == 1 && vzeta.n == 1
                     advance_ion_ionization_1V = !t_params.implicit_ion_advance
                     advance_neutral_ionization_1V = true
                 elseif vperp.n > 1 && vr.n > 1 && vzeta.n > 1
@@ -1183,8 +1181,6 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                     advance_neutral_ionization = true
                 else
                     error("If any perpendicular velocity has length>1 they all must. "
-                          * "If all perpendicular velocities have length=1, then vpa and "
-                          * "vz should be the same.\n"
                           * "vperp.n=$(vperp.n), vr.n=$(vr.n), vzeta.n=$(vzeta.n), "
                           * "vpa.n=$(vpa.n), vz.n=$(vz.n)")
                 end
@@ -1356,27 +1352,23 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
         advance_z_advection = z.n > 1
         advance_r_advection = r.n > 1
         if abs(collisions.charge_exchange) > 0.0
-            if vz.n == vpa.n && vperp.n == 1 && vr.n == 1 && vzeta.n == 1
+            if vperp.n == 1 && vr.n == 1 && vzeta.n == 1
                 advance_ion_cx_1V = true
             elseif vperp.n > 1 && vr.n > 1 && vzeta.n > 1
                 advance_ion_cx = true
             else
                 error("If any perpendicular velocity has length>1 they all must. "
-                      * "If all perpendicular velocities have length=1, then vpa and "
-                      * "vz should be the same.\n"
                       * "vperp.n=$(vperp.n), vr.n=$(vr.n), vzeta.n=$(vzeta.n), "
                       * "vpa.n=$(vpa.n), vz.n=$(vz.n)")
             end
         end
         if abs(collisions.ionization) > 0.0
-            if vz.n == vpa.n && vperp.n == 1 && vr.n == 1 && vzeta.n == 1
+            if vperp.n == 1 && vr.n == 1 && vzeta.n == 1
                 advance_ion_ionization_1V = true
             elseif vperp.n > 1 && vr.n > 1 && vzeta.n > 1
                 advance_ion_ionization = true
             else
                 error("If any perpendicular velocity has length>1 they all must. "
-                      * "If all perpendicular velocities have length=1, then vpa and "
-                      * "vz should be the same.\n"
                       * "vperp.n=$(vperp.n), vr.n=$(vr.n), vzeta.n=$(vzeta.n), "
                       * "vpa.n=$(vpa.n), vz.n=$(vz.n)")
             end

From 139e8b7b4e3b59169c4b48bf5a3b5660eeb23fbb Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 Jul 2024 09:10:22 +0100
Subject: [PATCH 387/428] Handle error when loading data in
 plots_for_variable()

---
 .../makie_post_processing/src/makie_post_processing.jl   | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index e0abcffea..8de72b013 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1005,7 +1005,14 @@ function plots_for_variable(run_info, variable_name; plot_prefix, has_rdim=true,
     println("Making plots for $variable_name")
     flush(stdout)
 
-    variable = get_variable(run_info, variable_name)
+    variable = nothing
+    try
+        variable = get_variable(run_info, variable_name)
+    catch e
+        println("plots_for_variable() failed for $variable_name - could not load data. "
+                * "Error was $e")
+        return nothing
+    end
 
     if variable_name ∈ em_variables
         species_indices = (nothing,)

From 9ddc035d23af78fe3cc9caeb29cffb1e8162c63c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 Jul 2024 13:51:32 +0100
Subject: [PATCH 388/428] Store time or pseudo-time in t_params structs

Logical, and makes it easier to keep track of cumulative pseudo-time for
electrons.
---
 .../src/electron_kinetic_equation.jl          |  67 +++++-----
 moment_kinetics/src/file_io.jl                |  37 ++++--
 moment_kinetics/src/initial_conditions.jl     |   7 +-
 moment_kinetics/src/input_structs.jl          |   1 +
 moment_kinetics/src/moment_kinetics.jl        |  12 +-
 moment_kinetics/src/runge_kutta.jl            |  10 +-
 moment_kinetics/src/time_advance.jl           | 114 +++++++++---------
 7 files changed, 140 insertions(+), 108 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 4629d9827..2ecf8fce3 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -65,7 +65,7 @@ OUTPUT:
 function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_spectral,
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
+        max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
         residual_tolerance=nothing, evolve_ppar=false)
 
     # set the method to use to solve the electron kinetic equation
@@ -142,7 +142,7 @@ OUTPUT:
 function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, collisions,
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_electron=nothing, initial_time=0.0,
+        max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
         residual_tolerance=nothing, evolve_ppar=false)
 
     begin_r_z_region()
@@ -182,11 +182,16 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
                          buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
-    time = initial_time
-    # Make sure that output times are set relative to this initial_time (the values in
-    # t_params are set relative to 0.0).
-    moments_output_times = t_params.moments_output_times .+ initial_time
-    dfns_output_times = t_params.dfns_output_times .+ initial_time
+    if initial_time !== nothing
+        @serial_region begin
+            t_params.t[] = initial_time
+        end
+        _block_synchronize()
+        # Make sure that output times are set relative to this initial_time (the values in
+        # t_params are set relative to 0.0).
+        moments_output_times = t_params.moments_output_times .+ initial_time
+        dfns_output_times = t_params.dfns_output_times .+ initial_time
+    end
     if io_electron === nothing && t_params.debug_io !== nothing
         # Overwrite the debug output file with the output from this call to
         # update_electron_pdf_with_time_advance!().
@@ -229,15 +234,15 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 # need to exit or handle this appropriately
                 @loop_vpa ivpa begin
                     @loop_z iz begin
-                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", scratch[t_params.n_rk_stages+1].pdf_electron[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", scratch[t_params.n_rk_stages+1].pdf_electron[ivpa, 1, iz, 1], " time: ", t_params.t[], " residual: ", residual[ivpa, 1, iz, 1])
                     end
                     println(io_pdf,"")
                 end
                 @loop_z iz begin
-                    println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                    println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
+                    println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                 end
                 println(io_upar,"")
                 println(io_qpar,"")
@@ -252,7 +257,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     t_params.moments_output_counter[] += 1
     @serial_region begin
         if io_electron !== nothing
-            write_electron_state(scratch, moments, t_params, time, io_electron,
+            write_electron_state(scratch, moments, t_params, io_electron,
                                  t_params.moments_output_counter[], r, z, vperp, vpa)
         end
     end
@@ -394,10 +399,11 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             update_derived_moments_and_derivatives()
 
             if t_params.adaptive && istage == t_params.n_rk_stages
-                electron_adaptive_timestep_update!(scratch, time, t_params, moments,
-                                                   phi, z_advect, vpa_advect, composition,
-                                                   r, z, vperp, vpa, vperp_spectral,
-                                                   vpa_spectral, external_source_settings,
+                electron_adaptive_timestep_update!(scratch, t_params.t[], t_params,
+                                                   moments, phi, z_advect, vpa_advect,
+                                                   composition, r, z, vperp, vpa,
+                                                   vperp_spectral, vpa_spectral,
+                                                   external_source_settings,
                                                    num_diss_params;
                                                    evolve_ppar=evolve_ppar)
                 # Re-do this in case electron_adaptive_timestep_update!() re-arranged the
@@ -415,7 +421,10 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         end
 
         # update the time following the pdf update
-        time += t_params.previous_dt[]
+        @serial_region begin
+            t_params.t[] += t_params.previous_dt[]
+        end
+        _block_synchronize()
 
         residual = -1.0
         if t_params.previous_dt[] > 0.0
@@ -466,9 +475,9 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             begin_serial_region()
             @serial_region begin
                 if z.irank == 0 && z.irank == z.nrank - 1
-                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
+                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", t_params.t[], " dt_electron: ", t_params.dt[], " phi_boundary: ", phi[[1,end],1], " residual: ", residual)
                 elseif z.irank == 0
-                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", time, " dt_electron: ", t_params.dt[], " phi_boundary_lower: ", phi[1,1], " residual: ", residual)
+                    println("iteration: ", t_params.step_counter[] - initial_step_counter, " time: ", t_params.t[], " dt_electron: ", t_params.dt[], " phi_boundary_lower: ", phi[1,1], " residual: ", residual)
                 end
             end
         end
@@ -482,17 +491,17 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     if (mod(t_params.moments_output_counter[], 100) == 0)
                         @loop_vpa ivpa begin
                             @loop_z iz begin
-                                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                                println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", new_pdf[ivpa, 1, iz, 1], " time: ", t_params.t[], " residual: ", residual[ivpa, 1, iz, 1])
                             end
                             println(io_pdf,"")
                         end
                         println(io_pdf,"")
                     end
                     @loop_z iz begin
-                        println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                        println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter)
-                        println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", time, " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
+                        println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                     end
                     println(io_upar,"")
                     println(io_qpar,"")
@@ -504,7 +513,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             @serial_region begin
                 if io_electron !== nothing
                     t_params.write_moments_output[] = false
-                    write_electron_state(scratch, moments, t_params, time, io_electron,
+                    write_electron_state(scratch, moments, t_params, io_electron,
                                          t_params.moments_output_counter[], r, z, vperp,
                                          vpa)
                 end
@@ -530,7 +539,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             if !electron_pdf_converged
                 @loop_vpa ivpa begin
                     @loop_z iz begin
-                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", time, " residual: ", residual[ivpa, 1, iz, 1])
+                        println(io_pdf, "z: ", z.grid[iz], " wpa: ", vpa.grid[ivpa], " pdf: ", pdf[ivpa, 1, iz, 1], " time: ", t_params.t[], " residual: ", residual[ivpa, 1, iz, 1])
                     end
                     println(io_pdf,"")
                 end
@@ -545,7 +554,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         if !electron_pdf_converged || do_debug_io
             if io_electron !== nothing && io_electron !== true
                 t_params.moments_output_counter[] += 1
-                write_electron_state(scratch, moments, t_params, time, io_electron,
+                write_electron_state(scratch, moments, t_params, io_electron,
                                      t_params.moments_output_counter[], r, z, vperp, vpa)
                 finish_electron_io(io_electron)
             end
@@ -556,7 +565,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     else
         success = ""
     end
-    return time, success
+    return success
 end
 
 """
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 1c12e16f5..5b6196dcc 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -191,6 +191,8 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
     dt_before_last_fail::Ttime
     # cumulative number of electron pseudo-timesteps taken
     electron_step_counter::Telectronint
+    # cumulative electron pseudo-time
+    electron_cumulative_pseudotime::Telectrontime
     # current electron pseudo-timestep size
     electron_dt::Telectrontime
     # size of last electron pseudo-timestep before the output was written
@@ -250,13 +252,12 @@ end
 structure containing the data/metadata needed for binary file i/o
 for electron initialization
 """
-struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4,
+struct io_initial_electron_info{Tfile, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4,
                                 Tconstr, Telectrontime, Telectronint, Telectronfailcause,
                                 Tinput}
     # file identifier for the binary file to which data is written
     fid::Tfile
-    # handle for the pseudotime variable
-    pseudotime::Ttime
+    time::Telectrontime
     # handle for the electron distribution function variable
     f_electron::Tfe
     # low-order approximation, used to diagnose timestepping error
@@ -296,6 +297,8 @@ struct io_initial_electron_info{Tfile, Ttime, Tfe, Tmom, Texte1, Texte2, Texte3,
     electron_constraints_C_coefficient::Tconstr
     # cumulative number of electron pseudo-timesteps taken
     electron_step_counter::Telectronint
+    # cumulative electron pseudo-time
+    electron_cumulative_pseudotime::Telectrontime
     # current electron pseudo-timestep size
     electron_dt::Telectrontime
     # size of last electron pseudo-timestep before the output was written
@@ -561,6 +564,7 @@ function reopen_initial_electron_io(file_info)
                                         getvar("electron_constraints_B_coefficient"),
                                         getvar("electron_constraints_C_coefficient"),
                                         getvar("electron_step_counter"),
+                                        getvar("electron_cumulative_pseudotime"),
                                         getvar("electron_dt"),
                                         getvar("electron_previous_dt"),
                                         getvar("electron_failure_counter"),
@@ -1581,6 +1585,10 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
             dynamic, "electron_step_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of electron pseudo-timesteps for the run")
 
+        io_electron_dt = create_dynamic_variable!(
+            dynamic, "electron_cumulative_pseudotime", mk_float; parallel_io=parallel_io,
+            description="cumulative electron pseudo-time")
+
         io_electron_dt = create_dynamic_variable!(
             dynamic, "electron_dt", mk_float; parallel_io=parallel_io,
             description="current electron pseudo-timestep size")
@@ -1614,6 +1622,7 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                         * "timestepping algorithm")
     else
         io_electron_step_counter = nothing
+        io_electron_cumulative_pseudotime = nothing
         io_electron_dt = nothing
         io_electron_previous_dt = nothing
         io_electron_failure_counter = nothing
@@ -2126,6 +2135,7 @@ function reopen_moments_io(file_info)
                                getvar("dt"), getvar("previous_dt"), getvar("failure_counter"),
                                getvar("failure_caused_by"), getvar("limit_caused_by"),
                                getvar("dt_before_last_fail"),getvar("electron_step_counter"),
+                               getvar("electron_cumulative_pseudotime"),
                                getvar("electron_dt"), getvar("electron_previous_dt"),
                                getvar("electron_failure_counter"),
                                getvar("electron_failure_caused_by"),
@@ -2288,6 +2298,7 @@ function reopen_dfns_io(file_info)
                                      getvar("limit_caused_by"),
                                      getvar("dt_before_last_fail"),
                                      getvar("electron_step_counter"),
+                                     getvar("electron_cumulative_pseudotime"),
                                      getvar("electron_dt"),
                                      getvar("electron_previous_dt"),
                                      getvar("electron_failure_counter"),
@@ -2338,7 +2349,7 @@ end
 write time-dependent moments data for ions, electrons and neutrals to the binary output
 file
 """
-function write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_species,
+function write_all_moments_data_to_binary(scratch, moments, fields, n_ion_species,
                                           n_neutral_species, io_or_file_info_moments,
                                           t_idx, time_for_run, t_params, nl_solver_params,
                                           r, z)
@@ -2357,7 +2368,7 @@ function write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_spe
         parallel_io = io_moments.io_input.parallel_io
 
         # add the time for this time slice to the hdf5 file
-        append_to_dynamic_var(io_moments.time, t, t_idx, parallel_io)
+        append_to_dynamic_var(io_moments.time, t_params.t[], t_idx, parallel_io)
 
         write_em_fields_data_to_binary(fields, io_moments, t_idx, r, z)
 
@@ -2623,6 +2634,9 @@ function write_electron_moments_data_to_binary(scratch, moments, t_params, elect
             # Save timestepping info
             append_to_dynamic_var(io_moments.electron_step_counter,
                                   electron_t_params.step_counter[], t_idx, parallel_io)
+            append_to_dynamic_var(io_moments.electron_cumulative_pseudotime,
+                                  electron_t_params.t[], t_idx,
+                                  parallel_io)
             append_to_dynamic_var(io_moments.electron_dt,
                                   electron_t_params.dt_before_output[], t_idx,
                                   parallel_io)
@@ -2755,7 +2769,7 @@ end
 write time-dependent distribution function data for ions, electrons and neutrals to the
 binary output file
 """
-function write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, t,
+function write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
                                        n_ion_species, n_neutral_species,
                                        io_or_file_info_dfns, t_idx, time_for_run,
                                        t_params, nl_solver_params, r, z, vperp, vpa,
@@ -2773,7 +2787,7 @@ function write_all_dfns_data_to_binary(scratch, scratch_electron, moments, field
 
         # Write the moments for this time slice to the output file.
         # This also updates the time.
-        write_all_moments_data_to_binary(scratch, moments, fields, t, n_ion_species,
+        write_all_moments_data_to_binary(scratch, moments, fields, n_ion_species,
                                          n_neutral_species, io_dfns.io_moments, t_idx,
                                          time_for_run, t_params, nl_solver_params, r, z)
 
@@ -2883,12 +2897,12 @@ function write_neutral_dfns_data_to_binary(scratch, t_params, n_neutral_species,
 end
 
 """
-    write_electron_state(scratch_electron, moments, t_params, t, io_initial_electron,
+    write_electron_state(scratch_electron, moments, t_params, io_initial_electron,
                          t_idx, r, z, vperp, vpa; pdf_electron_converged=false)
 
 Write the electron state to an output file.
 """
-function write_electron_state(scratch_electron, moments, t_params, t,
+function write_electron_state(scratch_electron, moments, t_params,
                               io_or_file_info_initial_electron, t_idx, r, z, vperp, vpa;
                               pdf_electron_converged=false)
 
@@ -2906,7 +2920,10 @@ function write_electron_state(scratch_electron, moments, t_params, t,
         parallel_io = io_initial_electron.io_input.parallel_io
 
         # add the pseudo-time for this time slice to the hdf5 file
-        append_to_dynamic_var(io_initial_electron.pseudotime, t, t_idx, parallel_io)
+        append_to_dynamic_var(io_initial_electron.time,
+                              t_params.t[], t_idx, parallel_io)
+        append_to_dynamic_var(io_initial_electron.electron_cumulative_pseudotime,
+                              t_params.t[], t_idx, parallel_io)
 
         write_electron_dfns_data_to_binary(scratch_electron, t_params,
                                            io_initial_electron, t_idx, r, z, vperp, vpa)
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index aeffb8ef3..155a35070 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -696,7 +696,7 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                                     previous_runs_info,
                                                     "initial_electron")
 
-            electron_pseudotime, success =
+            success =
                 @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments,
                                             fields.phi, r, z, vperp, vpa, z_spectral,
                                             vperp_spectral, vpa_spectral, z_advect,
@@ -756,7 +756,7 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                                vpa_advect, gyroavs, scratch_dummy, 0.0,
                                                initialisation_nl_solver_params)
             else
-                electron_pseudotime, success =
+                success =
                     update_electron_pdf!(scratch_electron, pdf.electron.norm, moments,
                                          fields.phi, r, z, vperp, vpa, z_spectral,
                                          vperp_spectral, vpa_spectral, z_advect,
@@ -764,8 +764,7 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                          collisions, composition,
                                          external_source_settings, num_diss_params,
                                          max_electron_pdf_iterations;
-                                         io_electron=io_initial_electron,
-                                         initial_time=electron_pseudotime)
+                                         io_electron=io_initial_electron)
             end
             if success != ""
                 error("!!!max number of iterations for electron pdf update exceeded!!!\n"
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index b78bae8ec..92ef94ca3 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -43,6 +43,7 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     n_variables::mk_int
     nstep::mk_int
     end_time::mk_float
+    t::MPISharedArray{mk_float,1}
     dt::MPISharedArray{mk_float,1}
     previous_dt::MPISharedArray{mk_float,1}
     next_output_time::MPISharedArray{mk_float,1}
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 2f4b9f9f1..e9bc50108 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -360,20 +360,20 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         restart_time_index, previous_runs_info, time_for_setup, t_params,
         nl_solver_params)
     # write initial data to ascii files
-    write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, code_time,
+    write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, t_params.t[],
         composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
-    write_all_moments_data_to_binary(scratch, moments, fields, code_time,
+    write_all_moments_data_to_binary(scratch, moments, fields, t_params.t[],
         composition.n_ion_species, composition.n_neutral_species, io_moments, 1, 0.0,
         t_params, nl_solver_params, r, z)
-    write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, code_time,
-        composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, 0.0,
-        t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
+    write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
+        t_params.t[], composition.n_ion_species, composition.n_neutral_species, io_dfns,
+        1, 0.0, t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 
-    return pdf, scratch, scratch_implicit, scratch_electron, code_time, t_params, vz, vr,
+    return pdf, scratch, scratch_implicit, scratch_electron, t_params, vz, vr,
            vzeta, vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
            advection_structs, composition, collisions, geometry, gyroavs,
            boundary_distributions, external_source_settings, num_diss_params,
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 24f396699..61808f67d 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -998,7 +998,7 @@ end
 
 Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_params`.
 """
-function adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms,
+function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
                                             total_points, current_dt, error_norm_method,
                                             success, nl_max_its_fraction, composition;
                                             electron=false)
@@ -1237,8 +1237,8 @@ function adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms
                 if (t_params.step_counter[] % 1000 == 0) && global_rank[] == 0
                     prefix = electron ? "electron" : "ion"
                     println("$prefix step ", t_params.step_counter[], ": t=",
-                            round(t, sigdigits=6), ", nfail=", t_params.failure_counter[],
-                            ", dt=", t_params.dt[])
+                            round(t_params.t[], sigdigits=6), ", nfail=",
+                            t_params.failure_counter[], ", dt=", t_params.dt[])
                 end
             end
         end
@@ -1248,12 +1248,12 @@ function adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms
         minimum_dt = 1.e-14
         if t_params.dt[] < minimum_dt
             println("Time advance failed: trying to set dt=$(t_params.dt[]) less than "
-                    * "$minimum_dt at t=$t. Ending run.")
+                    * "$minimum_dt at t=$(t_params.t[]). Ending run.")
             # Set dt negative to signal an error
             t_params.dt[] = -1.0
         end
 
-        current_time = t + t_params.previous_dt[]
+        current_time = t_params.t[] + t_params.previous_dt[]
         # Store here to ensure dt_before_output is set correctly when both moments and
         # dfns are written at the same time.
         current_dt = t_params.dt[]
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 1f3f8189f..37cfa0aea 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -338,6 +338,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                 * "`write_after_fixed_step_count=true`.")
     end
 
+    t_shared = allocate_shared_float(1)
     dt_shared = allocate_shared_float(1)
     previous_dt_shared = allocate_shared_float(1)
     next_output_time = allocate_shared_float(1)
@@ -348,6 +349,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
     write_moments_output = allocate_shared_bool(1)
     write_dfns_output = allocate_shared_bool(1)
     if block_rank[] == 0
+        t_shared[] = t_reload === nothing ? 0.0 : code_time
         dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         previous_dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         next_output_time[] = 0.0
@@ -435,11 +437,11 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         debug_io = nothing
         electron_t_params = electron
     end
-    return time_info(n_variables, t_input["nstep"], end_time, dt_shared, previous_dt_shared,
-                     next_output_time, dt_before_output, dt_before_last_fail,
-                     CFL_prefactor, step_to_moments_output, step_to_dfns_output,
-                     write_moments_output, write_dfns_output, Ref(0), Ref(2), Ref(2),
-                     Ref(0), mk_int[], mk_int[], t_input["nwrite"],
+    return time_info(n_variables, t_input["nstep"], end_time, t_shared, dt_shared,
+                     previous_dt_shared, next_output_time, dt_before_output,
+                     dt_before_last_fail, CFL_prefactor, step_to_moments_output,
+                     step_to_dfns_output, write_moments_output, write_dfns_output, Ref(0),
+                     Ref(2), Ref(2), Ref(0), mk_int[], mk_int[], t_input["nwrite"],
                      t_input["nwrite_dfns"], moments_output_times, dfns_output_times,
                      t_input["type"], rk_coefs, rk_coefs_implicit,
                      implicit_coefficient_is_zero, n_rk_stages, rk_order, adaptive,
@@ -1726,7 +1728,7 @@ df/dt + δv⋅∂f/∂z = 0, with δv(z,t)=v(z,t)-v₀(z)
 for prudent choice of v₀, expect δv≪v so that explicit
 time integrator can be used without severe CFL condition
 """
-function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz,
+function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
                        vr, vzeta, vpa, vperp, gyrophase, z, r, moments, fields,
                        spectral_objects, advect_objects, composition, collisions,
                        geometry, gyroavs, boundary_distributions,
@@ -1769,7 +1771,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
     iwrite_dfns = 2
     finish_now = false
     t_params.step_counter[] = 1
-    if t ≥ t_params.end_time - epsilon
+    if t_params.t[] ≥ t_params.end_time - epsilon
         # User must have requested zero output steps, i.e. to just write out the initial
         # profiles
         return nothing
@@ -1790,7 +1792,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
         if t_params.split_operators
             # MRH NOT SUPPORTED
             time_advance_split_operators!(pdf, scratch, scratch_implicit,
-                                          scratch_electron, t, t_params, vpa, z,
+                                          scratch_electron, t_params, vpa, z,
                                           vpa_spectral, z_spectral, moments, fields,
                                           vpa_advect, z_advect, composition, collisions,
                                           external_source_settings, num_diss_params,
@@ -1798,7 +1800,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                                           t_params.step_counter[])
         else
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                                       t, t_params, vz, vr, vzeta, vpa, vperp, gyrophase,
+                                       t_params, vz, vr, vzeta, vpa, vperp, gyrophase,
                                        z, r, moments, fields, spectral_objects,
                                        advect_objects, composition, collisions, geometry,
                                        gyroavs, boundary_distributions,
@@ -1808,9 +1810,12 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                                        diagnostic_checks, t_params.step_counter[])
         end
         # update the time
-        t += t_params.previous_dt[]
+        @serial_region begin
+            t_params.t[] += t_params.previous_dt[]
+        end
+        _block_synchronize()
 
-        if t ≥ t_params.end_time - epsilon ||
+        if t_params.t[] ≥ t_params.end_time - epsilon ||
                 (t_params.write_after_fixed_step_count &&
                  t_params.step_counter[] >= t_params.nstep)
             # Ensure all output is written at the final step
@@ -1820,7 +1825,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             # write output.
             # t_params.dt[] should never be NaN or Inf, so if it is something has gone
             # wrong.
-            println("dt=", t_params.dt[], " at t=$t, terminating run.")
+            println("dt=", t_params.dt[], " at t=", t_params.t[], ", terminating run.")
             finish_now = true
         end
 
@@ -1901,7 +1906,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                 if global_rank[] == 0
                     print("writing moments output ",
                           rpad(string(t_params.moments_output_counter[] - 1), 4), "  ",
-                          "t = ", rpad(string(round(t, sigdigits=6)), 7), "  ",
+                          "t = ", rpad(string(round(t_params.t[], sigdigits=6)), 7), "  ",
                           "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ")
                     if t_params.adaptive
                         print("nfail = ", rpad(string(t_params.failure_counter[]), 7), "  ",
@@ -1910,10 +1915,10 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                     print(Dates.format(now(), dateformat"H:MM:SS"))
                 end
             end
-            write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, t,
+            write_data_to_ascii(pdf, moments, fields, vpa, vperp, z, r, t_params.t[],
                                 composition.n_ion_species, composition.n_neutral_species,
                                 ascii_io)
-            write_all_moments_data_to_binary(scratch, moments, fields, t,
+            write_all_moments_data_to_binary(scratch, moments, fields,
                                              composition.n_ion_species,
                                              composition.n_neutral_species, io_moments,
                                              iwrite_moments, time_for_run, t_params,
@@ -1991,13 +1996,13 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
                 if global_rank[] == 0
                     println("writing distribution functions output ",
                             rpad(string(t_params.dfns_output_counter[] - 1), 4), "  ",
-                            "t = ", rpad(string(round(t, sigdigits=6)), 7), "  ",
+                            "t = ", rpad(string(round(t_params.t[], sigdigits=6)), 7), "  ",
                             "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ",
                             Dates.format(now(), dateformat"H:MM:SS"))
                     flush(stdout)
                 end
             end
-            write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields, t,
+            write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
                                           composition.n_ion_species,
                                           composition.n_neutral_species, io_dfns,
                                           iwrite_dfns, time_for_run, t_params,
@@ -2033,7 +2038,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t, t_pa
             break
         end
         if t_params.adaptive
-            if t >= t_params.end_time - epsilon
+            if t_params.t[] >= t_params.end_time - epsilon
                 break
             end
         else
@@ -2050,7 +2055,7 @@ end
 """
 """
 function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_electron,
-                                       t, t_params, vpa, z, vpa_spectral, z_spectral,
+                                       t_params, vpa, z, vpa_spectral, z_spectral,
                                        moments, fields, vpa_advect, z_advect, composition,
                                        collisions, external_source_settings,
                                        num_diss_params, nl_solver_params, advance,
@@ -2067,7 +2072,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         # advance the operator-split 1D advection equation in vpa
         # vpa-advection only applies for ion species
         advance.vpa_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
             t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
             z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
@@ -2075,7 +2080,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         # z_advection! advances the operator-split 1D advection equation in z
         # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
             t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
             z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
@@ -2085,14 +2090,14 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
             if collisions.charge_exchange > 0.0
                 advance.ion_cx_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    scratch_electron, t_params, vpa, z, vpa_spectral, z_spectral,
                     moments, fields, vpa_advect, z_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
                 advance.ion_cx_collisions = false
                 advance.neutral_cx_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    scratch_electron, t_params, vpa, z, vpa_spectral, z_spectral,
                     moments, fields, vpa_advect, z_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
@@ -2101,14 +2106,14 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
             if collisions.ionization > 0.0
                 advance.ion_ionization_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    scratch_electron, t_params, z, vpa, z_spectral, vpa_spectral,
                     moments, fields, z_advect, vpa_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
                 advance.ion_ionization_collisions = false
                 advance.neutral_ionization_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    scratch_electron, t_params, z, vpa, z_spectral, vpa_spectral,
                     moments, fields, z_advect, vpa_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
@@ -2118,7 +2123,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if collisions.krook.nuii0  > 0.0
             advance.krook_collisions_ii = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, z, vpa, z_spectral, vpa_spectral, moments, fields, z_advect,
+                t_params, z, vpa, z_spectral, vpa_spectral, moments, fields, z_advect,
                 vpa_advect, z_SL, vpa_SL, composition, collisions, sources,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.krook_collisions_ii = false
@@ -2128,7 +2133,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
             advance.source_terms = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.source_terms = false
@@ -2137,7 +2142,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_density
             advance.continuity = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.continuity = false
@@ -2146,7 +2151,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_upar
             advance.force_balance = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.force_balance = false
@@ -2155,7 +2160,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_ppar
             advance.energy = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.energy = false
@@ -2165,7 +2170,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_ppar
             advance.energy = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.energy = false
@@ -2174,7 +2179,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_upar
             advance.force_balance = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.force_balance = false
@@ -2183,7 +2188,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_density
             advance.continuity = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.continuity = false
@@ -2193,7 +2198,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar
             advance.source_terms = true
             time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
-                t, t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
+                t_params, vpa, z, vpa_spectral, z_spectral, moments, fields,
                 vpa_advect, z_advect, composition, collisions, external_source_settings,
                 num_diss_params, nl_solver_params, advance, advance_implicit, istep)
             advance.source_terms = false
@@ -2203,14 +2208,14 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
             if collisions.ionization > 0.0
                 advance.neutral_ionization = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    scratch_electron, t_params, z, vpa, z_spectral, vpa_spectral,
                     moments, fields, z_advect, vpa_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
                 advance.neutral_ionization = false
                 advance.ion_ionization = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, z, vpa, z_spectral, vpa_spectral,
+                    scratch_electron, t_params, z, vpa, z_spectral, vpa_spectral,
                     moments, fields, z_advect, vpa_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
@@ -2219,14 +2224,14 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
             if collisions.charge_exchange > 0.0
                 advance.neutral_cx_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    scratch_electron, t_params, vpa, z, vpa_spectral, z_spectral,
                     moments, fields, vpa_advect, z_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
                 advance.neutral_cx_collisions = false
                 advance.ion_cx_collisions = true
                 time_advance_no_splitting!(pdf, scratch, scratch_implicit,
-                    scratch_electron, t, t_params, vpa, z, vpa_spectral, z_spectral,
+                    scratch_electron, t_params, vpa, z, vpa_spectral, z_spectral,
                     moments, fields, vpa_advect, z_advect, composition, collisions,
                     external_source_settings, num_diss_params, nl_solver_params, advance,
                     advance_implicit, istep)
@@ -2236,7 +2241,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         # z_advection! advances the operator-split 1D advection equation in z
         # apply z-advection operation to all species (ion and neutral)
         advance.z_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
             t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
             z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
@@ -2244,7 +2249,7 @@ function time_advance_split_operators!(pdf, scratch, scratch_implicit, scratch_e
         # advance the operator-split 1D advection equation in vpa
         # vpa-advection only applies for ion species
         advance.vpa_advection = true
-        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+        time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
             t_params, vpa, z, vpa_spectral, z_spectral, moments, fields, vpa_advect,
             z_advect, composition, collisions, external_source_settings, num_diss_params,
             nl_solver_params, advance, advance_implicit, istep)
@@ -2255,7 +2260,7 @@ end
 
 """
 """
-function time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron, t,
+function time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_electron,
                                     t_params, vz, vr, vzeta, vpa, vperp, gyrophase, z, r,
                                     moments, fields, spectral_objects, advect_objects,
                                     composition, collisions, geometry, gyroavs,
@@ -2264,7 +2269,7 @@ function time_advance_no_splitting!(pdf, scratch, scratch_implicit, scratch_elec
                                     advance_implicit, fp_arrays, scratch_dummy,
                                     manufactured_source_list, diagnostic_checks, istep)
 
-    ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz, vr, vzeta,
+    ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz, vr, vzeta,
             vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
             advect_objects, composition, collisions, geometry, gyroavs,
             boundary_distributions, external_source_settings, num_diss_params,
@@ -2488,7 +2493,7 @@ end
 
 """
     adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
-                              t, t_params, moments, fields,
+                              t_params, moments, fields,
                               composition, collisions, geometry,
                               external_source_settings, spectral_objects,
                               advect_objects, gyroavs, num_diss_params, advance,
@@ -2498,7 +2503,7 @@ end
 Check the error estimate for the embedded RK method and adjust the timestep if
 appropriate.
 """
-function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron, t,
+function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
                                    t_params, pdf, moments, fields, boundary_distributions,
                                    composition, collisions, geometry,
                                    external_source_settings, spectral_objects,
@@ -2544,7 +2549,8 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         ion_z_CFL = Inf
         @loop_s is begin
             update_speed_z!(z_advect[is], moments.ion.upar, moments.ion.vth, evolve_upar,
-                            evolve_ppar, fields, vpa, vperp, z, r, t, geometry, is)
+                            evolve_ppar, fields, vpa, vperp, z, r, t_params.t[], geometry,
+                            is)
             this_minimum = get_minimum_CFL_z(z_advect[is].speed, z)
             @serial_region begin
                 ion_z_CFL = min(ion_z_CFL, this_minimum)
@@ -2558,7 +2564,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         begin_r_z_vperp_region()
         ion_vpa_CFL = Inf
         update_speed_vpa!(vpa_advect, fields, scratch[t_params.n_rk_stages+1], moments, vpa, vperp, z, r,
-                          composition, collisions, external_source_settings.ion, t,
+                          composition, collisions, external_source_settings.ion, t_params.t[],
                           geometry)
         @loop_s is begin
             this_minimum = get_minimum_CFL_vpa(vpa_advect[is].speed, vpa)
@@ -2737,7 +2743,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         @loop_sn isn begin
             update_speed_neutral_z!(neutral_z_advect[isn], moments.neutral.uz,
                                     moments.neutral.vth, evolve_upar, evolve_ppar, vz, vr,
-                                    vzeta, z, r, t)
+                                    vzeta, z, r, t_params.t[])
             this_minimum = get_minimum_CFL_neutral_z(neutral_z_advect[isn].speed, z)
             @serial_region begin
                 neutral_z_CFL = min(neutral_z_CFL, this_minimum)
@@ -2812,7 +2818,7 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         end
     end
 
-    adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
+    adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms, total_points,
                                        current_dt, error_norm_method, success,
                                        nl_max_its_fraction, composition)
 
@@ -2919,7 +2925,7 @@ end
 
 """
 """
-function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params, vz, vr,
+function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz, vr,
                  vzeta, vpa, vperp, gyrophase, z, r, moments, fields, spectral_objects,
                  advect_objects, composition, collisions, geometry, gyroavs,
                  boundary_distributions, external_source_settings, num_diss_params,
@@ -2979,7 +2985,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 # stage.
                 euler_time_advance!(scratch_implicit[istage], scratch[istage],
                                     pdf, fields, moments, advect_objects, vz, vr, vzeta,
-                                    vpa, vperp, gyrophase, z, r, t, t_params.dt[],
+                                    vpa, vperp, gyrophase, z, r, t_params.t[], t_params.dt[],
                                     spectral_objects, composition, collisions, geometry,
                                     scratch_dummy, manufactured_source_list,
                                     external_source_settings, num_diss_params,
@@ -3001,8 +3007,8 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
                 nl_success = backward_euler!(scratch_implicit[istage], scratch[istage],
                                              this_scratch_electron,
                                              pdf, fields, moments, advect_objects, vz, vr,
-                                             vzeta, vpa, vperp, gyrophase, z, r, t,
-                                             t_params.dt[] *
+                                             vzeta, vpa, vperp, gyrophase, z, r,
+                                             t_params.t[], t_params.dt[] *
                                              t_params.rk_coefs_implicit[istage,istage],
                                              spectral_objects, composition, collisions,
                                              geometry, scratch_dummy,
@@ -3047,7 +3053,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
         # calculate f^{(1)} = fⁿ + Δt*G[fⁿ] = scratch[2].pdf
         euler_time_advance!(scratch[istage+1], old_scratch, pdf, fields, moments,
                             advect_objects, vz, vr, vzeta, vpa, vperp, gyrophase, z,
-                            r, t, t_params.dt[], spectral_objects, composition,
+                            r, t_params.t[], t_params.dt[], spectral_objects, composition,
                             collisions, geometry, scratch_dummy,
                             manufactured_source_list, external_source_settings,
                             num_diss_params, advance, fp_arrays, istage)
@@ -3088,7 +3094,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t, t_params,
             end
         end
         adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
-                                  t, t_params, pdf, moments, fields,
+                                  t_params, pdf, moments, fields,
                                   boundary_distributions, composition, collisions,
                                   geometry, external_source_settings, spectral_objects,
                                   advect_objects, gyroavs, num_diss_params,

From 4130235d1efb069dda59722e745ab510c28fbec3 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 19 Jul 2024 14:00:35 +0100
Subject: [PATCH 389/428] Use electron_cumulative_pseudotime to postprocess
 electron_average_successful_dt

Fixes the calculation of electron_average_successful_dt in
post-processing, which was previously using the simulation time, which
resulted in incorrect (much too small) plots of the 'electron_dt'.
---
 .../src/electron_kinetic_equation.jl          |  2 +-
 moment_kinetics/src/file_io.jl                | 21 ++++++++++---------
 moment_kinetics/src/initial_conditions.jl     |  2 +-
 moment_kinetics/src/load_data.jl              |  3 ++-
 moment_kinetics/src/moment_kinetics.jl        | 10 ++++-----
 moment_kinetics/src/time_advance.jl           |  4 ++--
 6 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 2ecf8fce3..aa963b227 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -1366,7 +1366,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
         push!(total_points, z.n_global * r.n_global)
     end
 
-    adaptive_timestep_update_t_params!(t_params, t, CFL_limits, error_norms, total_points,
+    adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms, total_points,
                                        current_dt, error_norm_method, "", 0.0,
                                        composition; electron=true)
     if t_params.previous_dt[] == 0.0
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 5b6196dcc..2c339611c 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1029,10 +1029,10 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
         external_source_electron_momentum_amplitude,
         external_source_electron_pressure_amplitude,
         electron_constraints_A_coefficient, electron_constraints_B_coefficient,
-        electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
-        io_electron_previous_dt, io_electron_failure_counter,
-        io_electron_failure_caused_by, io_electron_limit_caused_by,
-        io_electron_dt_before_last_fail =
+        electron_constraints_C_coefficient, io_electron_step_counter,
+        io_electron_cumulative_pseudotime, io_electron_dt, io_electron_previous_dt,
+        io_electron_failure_counter, io_electron_failure_caused_by,
+        io_electron_limit_caused_by, io_electron_dt_before_last_fail =
             define_dynamic_electron_moment_variables!(fid, r, z, parallel_io,
                                                       external_source_settings,
                                                       evolve_density, evolve_upar,
@@ -1158,7 +1158,8 @@ function define_dynamic_moment_variables!(fid, n_ion_species, n_neutral_species,
                                io_time_for_run, io_step_counter, io_dt, io_previous_dt,
                                io_failure_counter, io_failure_caused_by,
                                io_limit_caused_by, io_dt_before_last_fail,
-                               io_electron_step_counter, io_electron_dt,
+                               io_electron_step_counter,
+                               io_electron_cumulative_pseudotime, io_electron_dt,
                                io_electron_previous_dt, io_electron_failure_counter,
                                io_electron_failure_caused_by, io_electron_limit_caused_by,
                                io_electron_dt_before_last_fail, io_nl_solver_diagnostics,
@@ -1585,7 +1586,7 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
             dynamic, "electron_step_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of electron pseudo-timesteps for the run")
 
-        io_electron_dt = create_dynamic_variable!(
+        io_electron_cumulative_pseudotime = create_dynamic_variable!(
             dynamic, "electron_cumulative_pseudotime", mk_float; parallel_io=parallel_io,
             description="cumulative electron pseudo-time")
 
@@ -1640,10 +1641,10 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
            external_source_electron_momentum_amplitude,
            external_source_electron_pressure_amplitude,
            electron_constraints_A_coefficient, electron_constraints_B_coefficient,
-           electron_constraints_C_coefficient, io_electron_step_counter, io_electron_dt,
-           io_electron_previous_dt, io_electron_failure_counter,
-           io_electron_failure_caused_by, io_electron_limit_caused_by,
-           io_electron_dt_before_last_fail
+           electron_constraints_C_coefficient, io_electron_step_counter,
+           io_electron_cumulative_pseudotime, io_electron_dt, io_electron_previous_dt,
+           io_electron_failure_counter, io_electron_failure_caused_by,
+           io_electron_limit_caused_by, io_electron_dt_before_last_fail
 end
 
 """
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 155a35070..aa8caaf06 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -775,7 +775,7 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
             # re-used if the simulation is re-run.
             t_params.electron.moments_output_counter[] += 1
             write_electron_state(scratch_electron, moments, t_params.electron,
-                                 electron_pseudotime, io_initial_electron,
+                                 io_initial_electron,
                                  t_params.electron.moments_output_counter[], r, z, vperp,
                                  vpa; pdf_electron_converged=true)
             finish_electron_io(io_initial_electron)
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 2e89d25e0..3d0923a07 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4666,8 +4666,9 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         electron_steps_per_output = get_variable(run_info, "electron_steps_per_output"; kwargs...)
         electron_failures_per_output = get_variable(run_info, "electron_failures_per_output"; kwargs...)
         electron_successful_steps_per_output = electron_steps_per_output - electron_failures_per_output
+        electron_pseudotime = get_variable("electron_cumulative_pseudotime"; kwargs...)
 
-        delta_t = copy(run_info.time)
+        delta_t = copy(electron_pseudotime)
         for i ∈ length(delta_t):-1:2
             delta_t[i] -= delta_t[i-1]
         end
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index e9bc50108..95a2de557 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -364,12 +364,12 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
-    write_all_moments_data_to_binary(scratch, moments, fields, t_params.t[],
-        composition.n_ion_species, composition.n_neutral_species, io_moments, 1, 0.0,
-        t_params, nl_solver_params, r, z)
+    write_all_moments_data_to_binary(scratch, moments, fields, composition.n_ion_species,
+        composition.n_neutral_species, io_moments, 1, 0.0, t_params, nl_solver_params, r,
+        z)
     write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
-        t_params.t[], composition.n_ion_species, composition.n_neutral_species, io_dfns,
-        1, 0.0, t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
+        composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, 0.0,
+        t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 37cfa0aea..c65b060a1 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -349,7 +349,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
     write_moments_output = allocate_shared_bool(1)
     write_dfns_output = allocate_shared_bool(1)
     if block_rank[] == 0
-        t_shared[] = t_reload === nothing ? 0.0 : code_time
+        t_shared[] = code_time
         dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         previous_dt_shared[] = dt_reload === nothing ? t_input["dt"] : dt_reload
         next_output_time[] = 0.0
@@ -2412,7 +2412,7 @@ function apply_all_bcs_constraints_update_moments!(
         # calculated here would be discarded - we might as well skip calculating it in
         # that case.
         if update_electrons && !t_params.implicit_electron_advance && success == ""
-            _, kinetic_electron_success = update_electron_pdf!(
+            kinetic_electron_success = update_electron_pdf!(
                scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
                z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect,
                scratch_dummy, t_params.electron, collisions, composition,

From fec2bef15affc5cf14b090f86155d591242c7ef8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Jul 2024 14:49:08 +0100
Subject: [PATCH 390/428] Use correct advection structs for electron solves

Don't think this actually caused a bug, because the advection structs
are effectively temporary buffers, and the z and vpa coordinates are
identical for electrons and ions (at least at the moment), but was not
what was intended, so has potential to cause bugs in future.
---
 moment_kinetics/src/time_advance.jl | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index c65b060a1..f3ae1635c 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2329,6 +2329,7 @@ function apply_all_bcs_constraints_update_moments!(
     z_spectral, r_spectral, vpa_spectral, vperp_spectral = spectral_objects.z_spectral, spectral_objects.r_spectral, spectral_objects.vpa_spectral, spectral_objects.vperp_spectral
     vzeta_spectral, vr_spectral, vz_spectral = spectral_objects.vzeta_spectral, spectral_objects.vr_spectral, spectral_objects.vz_spectral
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
+    electron_z_advect, electron_vpa_advect = advect_objects.electron_z_advect, advect_objects.electron_vpa_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
     success = ""
@@ -2414,9 +2415,10 @@ function apply_all_bcs_constraints_update_moments!(
         if update_electrons && !t_params.implicit_electron_advance && success == ""
             kinetic_electron_success = update_electron_pdf!(
                scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
-               z_spectral, vperp_spectral, vpa_spectral, z_advect, vpa_advect,
-               scratch_dummy, t_params.electron, collisions, composition,
-               external_source_settings, num_diss_params, max_electron_pdf_iterations)
+               z_spectral, vperp_spectral, vpa_spectral, electron_z_advect,
+               electron_vpa_advect, scratch_dummy, t_params.electron, collisions,
+               composition, external_source_settings, num_diss_params,
+               max_electron_pdf_iterations)
             success = kinetic_electron_success
         end
     end
@@ -3476,6 +3478,7 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
     vpa_spectral, vperp_spectral, r_spectral, z_spectral = spectral_objects.vpa_spectral, spectral_objects.vperp_spectral, spectral_objects.r_spectral, spectral_objects.z_spectral
     vz_spectral, vr_spectral, vzeta_spectral = spectral_objects.vz_spectral, spectral_objects.vr_spectral, spectral_objects.vzeta_spectral
     vpa_advect, vperp_advect, r_advect, z_advect = advect_objects.vpa_advect, advect_objects.vperp_advect, advect_objects.r_advect, advect_objects.z_advect
+    electron_z_advect, electron_vpa_advect = advect_objects.electron_z_advect, advect_objects.electron_vpa_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
     if composition.electron_physics == kinetic_electrons && advance.electron_energy
@@ -3483,8 +3486,9 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                              moments, fields, collisions, composition,
                                              external_source_settings, num_diss_params, r,
                                              z, vperp, vpa, r_spectral, z_spectral,
-                                             vperp_spectral, vpa_spectral, z_advect,
-                                             vpa_advect, gyroavs, scratch_dummy, dt,
+                                             vperp_spectral, vpa_spectral,
+                                             electron_z_advect, electron_vpa_advect,
+                                             gyroavs, scratch_dummy, dt,
                                              nl_solver_params.electron_advance)
     elseif advance.electron_conduction
         success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,

From 1eb2ff4d37980bc42f66b6343c7a03fc9f8b2848 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Jul 2024 15:55:32 +0100
Subject: [PATCH 391/428] Correction to setting of apply_bc_constraints for
 IMEX schemes

Removes hard-wired assumption that the first step is an explicit one.
---
 moment_kinetics/src/time_advance.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index f3ae1635c..b014e122b 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -3068,7 +3068,8 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
         # If `implicit_coefficient_is_zero` is true for the next stage, then this step is
         # explicit, so we need the bcs and constraints.
         apply_bc_constraints = (t_params.rk_coefs_implicit === nothing
-                                || istage == n_rk_stages
+                                || !t_params.implicit_ion_advance
+                                || (istage == n_rk_stages && t_params.implicit_coefficient_is_zero[1])
                                 || t_params.implicit_coefficient_is_zero[istage+1])
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
         success = apply_all_bcs_constraints_update_moments!(

From 77e058ceddb448d97efc568dfbc55126a4e7f286 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Jul 2024 20:29:12 +0100
Subject: [PATCH 392/428] Use moments_output_counter/dfns_output_counter
 consistently

Make the definitions of t_params.moments_output_counter and
t_params.dfns_output_counter more sensible (they are now initialised to
0 and incremented by 1 every time output is written), and actually use
them as the output index everywhere. This is just a tidy-up for the
ion/neutral output, but is a minor bug-fix for the electron output, as
the updated behaviour was assumed, but the output counters were actually
1 greater, so there was an empty entry written.
---
 moment_kinetics/src/initial_conditions.jl |  2 +
 moment_kinetics/src/moment_kinetics.jl    | 11 ++++--
 moment_kinetics/src/runge_kutta.jl        | 12 +++---
 moment_kinetics/src/time_advance.jl       | 46 ++++++++++++-----------
 4 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index aa8caaf06..10336f40b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -696,6 +696,8 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                                     previous_runs_info,
                                                     "initial_electron")
 
+            # Can't let this counter stay set to 0
+            t_params.electron.dfns_output_counter[] = max(t_params.electron.dfns_output_counter[], 1)
             success =
                 @views update_electron_pdf!(scratch_electron, pdf.electron.norm, moments,
                                             fields.phi, r, z, vperp, vpa, z_spectral,
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 95a2de557..0a0b5362b 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -364,12 +364,15 @@ function setup_moment_kinetics(input_dict::AbstractDict;
         composition.n_ion_species, composition.n_neutral_species, ascii_io)
     # write initial data to binary files
 
+    t_params.moments_output_counter[] += 1
     write_all_moments_data_to_binary(scratch, moments, fields, composition.n_ion_species,
-        composition.n_neutral_species, io_moments, 1, 0.0, t_params, nl_solver_params, r,
-        z)
+        composition.n_neutral_species, io_moments, t_params.moments_output_counter[], 0.0,
+        t_params, nl_solver_params, r, z)
+    t_params.dfns_output_counter[] += 1
     write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
-        composition.n_ion_species, composition.n_neutral_species, io_dfns, 1, 0.0,
-        t_params, nl_solver_params, r, z, vperp, vpa, vzeta, vr, vz)
+        composition.n_ion_species, composition.n_neutral_species, io_dfns,
+        t_params.dfns_output_counter[], 0.0, t_params, nl_solver_params, r, z, vperp, vpa,
+        vzeta, vr, vz)
 
     begin_s_r_z_vperp_region()
 
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 61808f67d..06bf692ae 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1260,11 +1260,11 @@ function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
         if (!t_params.write_after_fixed_step_count
             && !t_params.write_moments_output[]
             && length(t_params.moments_output_times) > 0
-            && (t_params.moments_output_counter[] ≤ length(t_params.moments_output_times) + 1)
-            && (current_time + t_params.dt[] >= t_params.moments_output_times[t_params.moments_output_counter[]-1]))
+            && (t_params.moments_output_counter[] ≤ length(t_params.moments_output_times))
+            && (current_time + t_params.dt[] >= t_params.moments_output_times[t_params.moments_output_counter[]]))
 
             t_params.dt_before_output[] = current_dt
-            t_params.dt[] = t_params.moments_output_times[t_params.moments_output_counter[]-1] - current_time
+            t_params.dt[] = t_params.moments_output_times[t_params.moments_output_counter[]] - current_time
             t_params.step_to_moments_output[] = true
 
             if t_params.dt[] < 0.0
@@ -1275,11 +1275,11 @@ function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
         if (!t_params.write_after_fixed_step_count
             && !t_params.write_dfns_output[]
             && length(t_params.dfns_output_times) > 0
-            && (t_params.dfns_output_counter[] ≤ length(t_params.dfns_output_times) + 1)
-            && (current_time + t_params.dt[] >= t_params.dfns_output_times[t_params.dfns_output_counter[]-1]))
+            && (t_params.dfns_output_counter[] ≤ length(t_params.dfns_output_times))
+            && (current_time + t_params.dt[] >= t_params.dfns_output_times[t_params.dfns_output_counter[]]))
 
             t_params.dt_before_output[] = current_dt
-            t_params.dt[] = t_params.dfns_output_times[t_params.dfns_output_counter[]-1] - current_time
+            t_params.dt[] = t_params.dfns_output_times[t_params.dfns_output_counter[]] - current_time
             t_params.step_to_dfns_output[] = true
 
             if t_params.dt[] < 0.0
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index b014e122b..31ff54dab 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -441,7 +441,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      previous_dt_shared, next_output_time, dt_before_output,
                      dt_before_last_fail, CFL_prefactor, step_to_moments_output,
                      step_to_dfns_output, write_moments_output, write_dfns_output, Ref(0),
-                     Ref(2), Ref(2), Ref(0), mk_int[], mk_int[], t_input["nwrite"],
+                     Ref(0), Ref(0), Ref(0), mk_int[], mk_int[], t_input["nwrite"],
                      t_input["nwrite_dfns"], moments_output_times, dfns_output_times,
                      t_input["type"], rk_coefs, rk_coefs_implicit,
                      implicit_coefficient_is_zero, n_rk_stages, rk_order, adaptive,
@@ -778,16 +778,24 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
                               scratch_electron, nl_solver_params, t_params, t_input,
                               num_diss_params, advection_structs, io_input, input_dict;
                               restart_electron_physics=restart_electron_physics)
-    elseif restarting && composition.electron_physics == kinetic_electrons &&
-           t_params.electron.debug_io !== nothing
-        # Create *.electron_debug.h5 file so that it can be re-opened in
-        # update_electron_pdf!().
-        io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z, r,
-                                        composition, collisions, moments.evolve_density,
-                                        moments.evolve_upar, moments.evolve_ppar,
-                                        external_source_settings, t_params.electron,
-                                        t_params.electron.debug_io[2], -1, nothing,
-                                        "electron_debug")
+    elseif restarting && composition.electron_physics == kinetic_electrons
+        if t_params.electron.debug_io !== nothing
+            # Create *.electron_debug.h5 file so that it can be re-opened in
+            # update_electron_pdf!().
+            io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z, r,
+                                            composition, collisions, moments.evolve_density,
+                                            moments.evolve_upar, moments.evolve_ppar,
+                                            external_source_settings, t_params.electron,
+                                            t_params.electron.debug_io[2], -1, nothing,
+                                            "electron_debug")
+        end
+
+        # No need to do electron I/O (apart from possibly debug I/O) any more, so if
+        # adaptive timestep is used, it does not need to adjust to output times.
+        resize!(t_params.electron.moments_output_times, 0)
+        resize!(t_params.electron.dfns_output_times, 0)
+        t_params.electron.moments_output_counter[] = 1
+        t_params.electron.dfns_output_counter[] = 1
     end
 
     # update the derivatives of the electron moments as these may be needed when
@@ -1767,8 +1775,6 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
     epsilon = 1.e-11
 
     # main time advance loop
-    iwrite_moments = 2
-    iwrite_dfns = 2
     finish_now = false
     t_params.step_counter[] = 1
     if t_params.t[] ≥ t_params.end_time - epsilon
@@ -1905,7 +1911,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
             @serial_region begin
                 if global_rank[] == 0
                     print("writing moments output ",
-                          rpad(string(t_params.moments_output_counter[] - 1), 4), "  ",
+                          rpad(string(t_params.moments_output_counter[]), 4), "  ",
                           "t = ", rpad(string(round(t_params.t[], sigdigits=6)), 7), "  ",
                           "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ")
                     if t_params.adaptive
@@ -1921,7 +1927,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
             write_all_moments_data_to_binary(scratch, moments, fields,
                                              composition.n_ion_species,
                                              composition.n_neutral_species, io_moments,
-                                             iwrite_moments, time_for_run, t_params,
+                                             t_params.moments_output_counter[], time_for_run, t_params,
                                              nl_solver_params, r, z)
 
             if t_params.steady_state_residual
@@ -1978,7 +1984,6 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
                 end
             end
 
-            iwrite_moments += 1
             begin_s_r_z_vperp_region()
             @debug_detect_redundant_block_synchronize begin
                 # Reactivate check for redundant _block_synchronize()
@@ -1995,7 +2000,7 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
             @serial_region begin
                 if global_rank[] == 0
                     println("writing distribution functions output ",
-                            rpad(string(t_params.dfns_output_counter[] - 1), 4), "  ",
+                            rpad(string(t_params.dfns_output_counter[]), 4), "  ",
                             "t = ", rpad(string(round(t_params.t[], sigdigits=6)), 7), "  ",
                             "nstep = ", rpad(string(t_params.step_counter[]), 7), "  ",
                             Dates.format(now(), dateformat"H:MM:SS"))
@@ -2005,10 +2010,9 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
             write_all_dfns_data_to_binary(scratch, scratch_electron, moments, fields,
                                           composition.n_ion_species,
                                           composition.n_neutral_species, io_dfns,
-                                          iwrite_dfns, time_for_run, t_params,
-                                          nl_solver_params, r, z, vperp, vpa, vzeta, vr,
-                                          vz)
-            iwrite_dfns += 1
+                                          t_params.dfns_output_counter[], time_for_run,
+                                          t_params, nl_solver_params, r, z, vperp, vpa,
+                                          vzeta, vr, vz)
             begin_s_r_z_vperp_region()
             @debug_detect_redundant_block_synchronize begin
                 # Reactivate check for redundant _block_synchronize()

From a146f25979ab3fa6ff93a4f04f63ad1bea9ff3ee Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 22 Jul 2024 14:57:08 +0100
Subject: [PATCH 393/428] Implicit solve for electron_ppar as part of electron
 pseudo-timestepping

Activated if `implicit_electron_ppar` is passed in the `[timestepping]`
section.
---
 .../src/electron_kinetic_equation.jl          | 101 ++++++++++++------
 moment_kinetics/src/input_structs.jl          |   1 +
 moment_kinetics/src/moment_kinetics_input.jl  |   1 +
 moment_kinetics/src/time_advance.jl           |  71 +++++++-----
 4 files changed, 116 insertions(+), 58 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index aa963b227..5c9be2f5a 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -59,6 +59,9 @@ The electron kinetic equation is:
     scratch_dummy = dummy arrays to be used for temporary storage
     dt = time step size
     max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
+    ion_dt = if this is passed, the electron pressure is evolved in a form that results in
+             a backward-Euler update on the ion timestep (ion_dt) once the electron
+             pseudo-timestepping reaches steady state.
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
@@ -66,7 +69,7 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
-        residual_tolerance=nothing, evolve_ppar=false)
+        residual_tolerance=nothing, evolve_ppar=false, ion_dt=nothing)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -79,7 +82,7 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
             vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
             external_source_settings, num_diss_params, max_electron_pdf_iterations;
             io_electron=io_electron, initial_time=initial_time,
-            residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar)
+            residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar, ion_dt=ion_dt)
     elseif solution_method == "shooting_method"
         dens = moments.electron.dens
         vthe = moments.electron.vth
@@ -136,6 +139,9 @@ The electron kinetic equation is:
     max_electron_pdf_iterations = maximum number of iterations to use in the solution of the electron kinetic equation
     io_electron = info struct for binary file I/O
     initial_time = initial value for the (pseudo-)time
+    ion_dt = if this is passed, the electron pressure is evolved in a form that results in
+             a backward-Euler update on the ion timestep (ion_dt) once the electron
+             pseudo-timestepping reaches steady state.
 OUTPUT:
     pdf = updated (modified) electron pdf
 """
@@ -143,10 +149,14 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
         max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
-        residual_tolerance=nothing, evolve_ppar=false)
+        residual_tolerance=nothing, evolve_ppar=false, ion_dt=nothing)
 
     begin_r_z_region()
 
+    if ion_dt !== nothing
+        evolve_ppar = true
+    end
+
     # create several (r) dimension dummy arrays for use in taking derivatives
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
@@ -182,6 +192,18 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
                          buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 
+    if !evolve_ppar
+        # ppar is not updated in the pseudo-timestepping loop below. So that we can read
+        # ppar from the scratch structs, copy moments.electron.ppar into all of them.
+        moments_ppar = moments.electron.ppar
+        for istage ∈ 1:t_params.n_rk_stages+1
+            scratch_ppar = scratch[istage].electron_ppar
+            @loop_r_z ir iz begin
+                scratch_ppar[iz,ir] = moments_ppar[iz,ir]
+            end
+        end
+    end
+
     if initial_time !== nothing
         @serial_region begin
             t_params.t[] = initial_time
@@ -241,7 +263,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 @loop_z iz begin
                     println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
                     println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
-                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                    println(io_ppar, "z: ", z.grid[iz], " ppar: ", scratch[t_params.n_rk_stages+1].electron_ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
                     println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                 end
                 println(io_upar,"")
@@ -308,7 +330,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                                     scratch_dummy, collisions,
                                                     composition, external_source_settings,
                                                     num_diss_params, t_params.dt[];
-                                                    evolve_ppar=evolve_ppar)
+                                                    evolve_ppar=evolve_ppar,
+                                                    ion_dt=ion_dt)
             speedup_hack!(scratch[istage+1], scratch[istage], z_speedup_fac, z, vpa;
                           evolve_ppar=evolve_ppar)
 
@@ -317,15 +340,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             if evolve_ppar
                 rk_update_variable!(scratch, nothing, :electron_ppar, t_params, istage)
 
-                begin_r_z_region()
-                moments_struct_ppar = moments.electron.ppar
-                scratch_ppar = scratch[istage+1].electron_ppar
-                @loop_r_z ir iz begin
-                    moments_struct_ppar[iz,ir] = scratch_ppar[iz,ir]
-                end
-                _block_synchronize()
-
-                update_electron_vth_temperature!(moments, moments_struct_ppar,
+                update_electron_vth_temperature!(moments, scratch[istage+1].electron_ppar,
                                                  moments.electron.dens, composition)
             end
 
@@ -340,8 +355,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                 # update the electron heat flux
                 moments.electron.qpar_updated[] = false
                 calculate_electron_qpar_from_pdf!(moments.electron.qpar,
-                                                  moments.electron.ppar, moments.electron.vth,
-                                                  latest_pdf, vpa)
+                                                  scratch[istage+1].electron_ppar,
+                                                  moments.electron.vth, latest_pdf, vpa)
 
                 # compute the z-derivative of the parallel electron heat flux
                 @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar,
@@ -500,7 +515,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                     @loop_z iz begin
                         println(io_upar, "z: ", z.grid[iz], " upar: ", moments.electron.upar[iz,1], " dupar_dz: ", moments.electron.dupar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
                         println(io_qpar, "z: ", z.grid[iz], " qpar: ", moments.electron.qpar[iz,1], " dqpar_dz: ", moments.electron.dqpar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
-                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", moments.electron.ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
+                        println(io_ppar, "z: ", z.grid[iz], " ppar: ", scratch[t_params.n_rk_stages+1].electron_ppar[iz,1], " dppar_dz: ", moments.electron.dppar_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter)
                         println(io_vth, "z: ", z.grid[iz], " vthe: ", moments.electron.vth[iz,1], " dvth_dz: ", moments.electron.dvth_dz[iz,1], " time: ", t_params.t[], " iteration: ", t_params.step_counter[] - initial_step_counter, " dens: ", dens[iz,1])
                     end
                     println(io_upar,"")
@@ -533,6 +548,15 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
         pdf[ivpa,ivperp,iz,ir] = final_scratch_pdf[ivpa,ivperp,iz,ir]
     end
+    if evolve_ppar
+        # Update `moments.electron.ppar` with the final electron pressure
+        begin_r_z_region()
+        scratch_ppar = scratch[t_params.n_rk_stages+1].electron_ppar
+        moments_ppar = moments.electron.ppar
+        @loop_r_z ir iz begin
+            moments_ppar[iz,ir] = scratch_ppar[iz,ir]
+        end
+    end
     begin_serial_region()
     @serial_region begin
         if text_output
@@ -1306,8 +1330,9 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
     # vpa-advection
     begin_r_z_vperp_region()
     update_electron_speed_vpa!(vpa_advect[1], moments.electron.dens,
-                               moments.electron.upar, moments.electron.ppar,
-                               moments, vpa.grid, external_source_settings.electron)
+                               moments.electron.upar,
+                               scratch[t_params.n_rk_stages+1].electron_ppar, moments,
+                               vpa.grid, external_source_settings.electron)
     vpa_CFL = get_minimum_CFL_vpa(vpa_advect[1].speed, vpa)
     if block_rank[] == 0
         push!(CFL_limits, t_params.CFL_prefactor * vpa_CFL)
@@ -1592,12 +1617,8 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
                                                  vpa, z_spectral, vpa_spectral, z_advect,
                                                  vpa_advect, scratch_dummy, collisions,
                                                  composition, external_source_settings,
-                                                 num_diss_params, dt; evolve_ppar=false)
-    if evolve_ppar
-        ppar = fvec_in.electron_ppar
-    else
-        ppar = moments.electron.ppar
-    end
+                                                 num_diss_params, dt; evolve_ppar=false,
+                                                 ion_dt=nothing)
     # add the contribution from the z advection term
     electron_z_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
                           moments.electron.upar, moments.electron.vth, z_advect, z,
@@ -1605,14 +1626,15 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
 
     # add the contribution from the wpa advection term
     electron_vpa_advection!(fvec_out.pdf_electron, fvec_in.pdf_electron,
-                            moments.electron.dens, moments.electron.upar, ppar,
-                            moments, vpa_advect, vpa, vpa_spectral, scratch_dummy, dt,
-                            external_source_settings.electron)
+                            moments.electron.dens, moments.electron.upar,
+                            fvec_in.electron_ppar, moments, vpa_advect, vpa, vpa_spectral,
+                            scratch_dummy, dt, external_source_settings.electron)
 
     # add in the contribution to the residual from the term proportional to the pdf
-    add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, ppar,
-                                    moments.electron.dens, moments.electron.upar, moments,
-                                    vpa.grid, z, dt, external_source_settings.electron)
+    add_contribution_from_pdf_term!(fvec_out.pdf_electron, fvec_in.pdf_electron,
+                                    fvec_in.electron_ppar, moments.electron.dens,
+                                    moments.electron.upar, moments, vpa.grid, z, dt,
+                                    external_source_settings.electron)
 
     # add in numerical dissipation terms
     add_dissipation_term!(fvec_out.pdf_electron, fvec_in.pdf_electron, scratch_dummy,
@@ -1643,6 +1665,23 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
                                   moments.neutral.pz, moments.electron, collisions, dt,
                                   composition, external_source_settings.electron,
                                   num_diss_params, z)
+
+        if ion_dt !== nothing
+            # Add source term to turn steady state solution into a backward-Euler update of
+            # electron_ppar with the ion timestep `ion_dt`.
+            ppar_out = fvec_out.electron_ppar
+            ppar_previous_ion_step = moments.electron.ppar
+            begin_r_z_region()
+            @loop_r_z ir iz begin
+                # At this point, ppar_out = ppar_in + dt*RHS(ppar_in). Here we add a
+                # source/damping term so that in the steady state of the electron
+                # pseudo-timestepping iteration,
+                #   RHS(ppar) - (ppar - ppar_previous_ion_step) / ion_dt = 0,
+                # resulting in a backward-Euler step (as long as the pseudo-timestepping
+                # loop converges).
+                ppar_out[iz,ir] += -dt * (ppar_out[iz,ir] - ppar_previous_ion_step[iz,ir]) / ion_dt
+            end
+        end
     end
 
     return nothing
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 92ef94ca3..0f27e36bb 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -86,6 +86,7 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     implicit_ion_advance::Bool
     implicit_vpa_advection::Bool
     implicit_ion_maxwell_diffusion::Bool
+    implicit_electron_ppar::Bool
     write_after_fixed_step_count::Bool
     error_sum_zero::Terrorsum
     split_operators::Bool
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 937dbcecb..2dd0f6b0d 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -234,6 +234,7 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         implicit_ion_advance=false,
         implicit_vpa_advection=false,
         implicit_ion_maxwell_diffusion=false,
+        implicit_electron_ppar=false,
         write_after_fixed_step_count=false,
         write_error_diagnostics=false,
         write_steady_state_diagnostics=false,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 31ff54dab..6ea09e764 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -396,12 +396,14 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         t_input["implicit_ion_advance"] = false
         t_input["implicit_vpa_advection"] = false
         t_input["implicit_ion_maxwell_diffusion"] = false
+        t_input["implicit_electron_ppar"] = false
     else
         if composition.electron_physics != braginskii_fluid
             t_input["implicit_braginskii_conduction"] = false
         end
         if composition.electron_physics != kinetic_electrons
             t_input["implicit_electron_advance"] = false
+            t_input["implicit_electron_ppar"] = false
         end
     end
 
@@ -455,6 +457,7 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      electron !== nothing && t_input["implicit_ion_advance"],
                      electron !== nothing && t_input["implicit_vpa_advection"],
                      electron !== nothing && t_input["implicit_ion_maxwell_diffusion"],
+                     electron !== nothing && t_input["implicit_electron_ppar"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
                      t_input["split_operators"], t_input["steady_state_residual"],
                      t_input["converged_residual_value"],
@@ -724,6 +727,10 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         error("Cannot use implicit_vpa_advection and implicit_ion_maxwell_diffusion at the same "
               * "time")
     end
+    if nl_solver_electron_advance_params !== nothing && t_params.implicit_electron_ppar
+        error("Cannot use implicit_electron_advance and implicit_electron_ppar at the "
+              * "same time.")
+    end
     nl_solver_params = (electron_conduction=electron_conduction_nl_solve_parameters,
                         electron_advance=nl_solver_electron_advance_params,
                         ion_advance=nl_solver_ion_advance_params,
@@ -1251,7 +1258,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         # if treating the electrons as a fluid with Braginskii closure, or
         # moment-kinetically then advance the electron energy equation
         if composition.electron_physics == kinetic_electrons
-            if !t_params.implicit_electron_advance
+            if !(t_params.implicit_electron_advance || t_params.implicit_electron_ppar)
                 advance_electron_energy = true
                 advance_electron_conduction = true
             end
@@ -1415,7 +1422,7 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
         advance_electron_conduction = true
     end
 
-    if t_params.implicit_electron_advance
+    if (t_params.implicit_electron_advance || t_params.implicit_electron_ppar)
         advance_electron_energy = true
         advance_electron_conduction = true
     end
@@ -2416,7 +2423,9 @@ function apply_all_bcs_constraints_update_moments!(
         # to the beginning of the ion/neutral timestep, so the electron solution
         # calculated here would be discarded - we might as well skip calculating it in
         # that case.
-        if update_electrons && !t_params.implicit_electron_advance && success == ""
+        if update_electrons &&
+                !(t_params.implicit_electron_advance || t_params.implicit_electron_ppar) &&
+                success == ""
             kinetic_electron_success = update_electron_pdf!(
                scratch_electron, pdf.electron.norm, moments, fields.phi, r, z, vperp, vpa,
                z_spectral, vperp_spectral, vpa_spectral, electron_z_advect,
@@ -3007,17 +3016,19 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
                 # rk_coefs_implicit[istage,istage]=a[istage,istage].
                 if scratch_electron === nothing
                     this_scratch_electron = nothing
-                else
+                elseif t_params.implicit_electron_advance
                     this_scratch_electron = scratch_electron[t_params.electron.n_rk_stages+1]
+                else
+                    this_scratch_electron = scratch_electron
                 end
                 nl_success = backward_euler!(scratch_implicit[istage], scratch[istage],
                                              this_scratch_electron,
                                              pdf, fields, moments, advect_objects, vz, vr,
                                              vzeta, vpa, vperp, gyrophase, z, r,
-                                             t_params.t[], t_params.dt[] *
+                                             t_params.dt[] *
                                              t_params.rk_coefs_implicit[istage,istage],
-                                             spectral_objects, composition, collisions,
-                                             geometry, scratch_dummy,
+                                             t_params, spectral_objects, composition,
+                                             collisions, geometry, scratch_dummy,
                                              manufactured_source_list,
                                              external_source_settings, num_diss_params,
                                              gyroavs, nl_solver_params, advance_implicit,
@@ -3075,6 +3086,10 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
                                 || !t_params.implicit_ion_advance
                                 || (istage == n_rk_stages && t_params.implicit_coefficient_is_zero[1])
                                 || t_params.implicit_coefficient_is_zero[istage+1])
+        update_electrons = (t_params.rk_coefs_implicit === nothing
+                            || !(t_params.implicit_electron_advance || t_params.implicit_electron_ppar)
+                            || t_params.implicit_coefficient_is_zero[istage+1]
+                            || (istage == n_rk_stages && t_params.implicit_coefficient_is_zero[1]))
         diagnostic_moments = diagnostic_checks && istage == n_rk_stages
         success = apply_all_bcs_constraints_update_moments!(
             scratch[istage+1], pdf, moments, fields, boundary_distributions,
@@ -3082,7 +3097,7 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
             advect_objects, composition, collisions, geometry, gyroavs,
             external_source_settings, num_diss_params, t_params, advance, scratch_dummy,
             diagnostic_moments; pdf_bc_constraints=apply_bc_constraints,
-            update_electrons=apply_bc_constraints)
+            update_electrons=update_electrons)
         if success != ""
             # Break out of the istage loop, as passing `success != ""` to the
             # adaptive timestep update function will signal a failed timestep, so
@@ -3474,8 +3489,8 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
 end
 
 function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, moments,
-                         advect_objects, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, t,
-                         dt, spectral_objects, composition, collisions, geometry,
+                         advect_objects, vz, vr, vzeta, vpa, vperp, gyrophase, z, r, dt,
+                         t_params, spectral_objects, composition, collisions, geometry,
                          scratch_dummy, manufactured_source_list,
                          external_source_settings, num_diss_params, gyroavs,
                          nl_solver_params, advance, fp_arrays, istage)
@@ -3486,7 +3501,7 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
     electron_z_advect, electron_vpa_advect = advect_objects.electron_z_advect, advect_objects.electron_vpa_advect
     neutral_z_advect, neutral_r_advect, neutral_vz_advect = advect_objects.neutral_z_advect, advect_objects.neutral_r_advect, advect_objects.neutral_vz_advect
 
-    if composition.electron_physics == kinetic_electrons && advance.electron_energy
+    if nl_solver_params.electron_advance !== nothing
         success = implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron,
                                              moments, fields, collisions, composition,
                                              external_source_settings, num_diss_params, r,
@@ -3495,6 +3510,17 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                              electron_z_advect, electron_vpa_advect,
                                              gyroavs, scratch_dummy, dt,
                                              nl_solver_params.electron_advance)
+    elseif t_params.implicit_electron_ppar
+        max_electron_pdf_iterations = 1000
+        electron_success = update_electron_pdf!(scratch_electron, pdf.electron.norm,
+                                                moments, fields.phi, r, z, vperp, vpa,
+                                                z_spectral, vperp_spectral, vpa_spectral,
+                                                electron_z_advect, electron_vpa_advect,
+                                                scratch_dummy, t_params.electron,
+                                                collisions, composition,
+                                                external_source_settings, num_diss_params,
+                                                max_electron_pdf_iterations; ion_dt=dt)
+        success = (electron_success == "")
     elseif advance.electron_conduction
         success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,
                                                   z_spectral, composition, collisions,
@@ -3505,27 +3531,21 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
     if nl_solver_params.ion_advance !== nothing
         success = implicit_ion_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                         advect_objects, vz, vr, vzeta, vpa, vperp,
-                                        gyrophase, z, r, t, dt, spectral_objects,
-                                        composition, collisions, geometry, scratch_dummy,
-                                        manufactured_source_list,
+                                        gyrophase, z, r, t_params.t[], dt,
+                                        spectral_objects, composition, collisions,
+                                        geometry, scratch_dummy, manufactured_source_list,
                                         external_source_settings, num_diss_params,
                                         gyroavs, nl_solver_params.ion_advance, advance,
                                         fp_arrays, istage)
-        if !success
-            return success
-        end
     elseif advance.vpa_advection
         success = implicit_vpa_advection!(fvec_out.pdf, fvec_in, fields, moments,
-                                          z_advect, vpa_advect, vpa, vperp, z, r, dt, t,
-                                          r_spectral, z_spectral, vpa_spectral,
-                                          composition, collisions,
+                                          z_advect, vpa_advect, vpa, vperp, z, r, dt,
+                                          t_params.t[], r_spectral, z_spectral,
+                                          vpa_spectral, composition, collisions,
                                           external_source_settings.ion, geometry,
                                           nl_solver_params.vpa_advection,
                                           advance.vpa_diffusion, num_diss_params, gyroavs,
                                           scratch_dummy)
-        if !success
-            return success
-        end
     elseif advance.mxwl_diff_collisions_ii
         success = implicit_ion_maxwell_diffusion!(fvec_out.pdf, fvec_in, moments,
                                                   z_advect, vpa, vperp, z, r, dt,
@@ -3533,12 +3553,9 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                                   collisions, geometry,
                                                   nl_solver_params.maxwell_diffusion,
                                                   gyroavs, scratch_dummy)
-        if !success
-            return success
-        end
     end
 
-    return true
+    return success
 end
 
 """

From 52232978353a47a67ec7eed96662b90b3a03999d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Jul 2024 10:07:18 +0100
Subject: [PATCH 394/428] Fix titles in comparison plot animations

...although the titles don't seem to be shown (maybe covered up by the
legend??).
---
 .../makie_post_processing/src/makie_post_processing.jl    | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 8de72b013..a284c7294 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -1898,8 +1898,8 @@ for dim ∈ one_dimension_combinations_no_t
                          time = select_slice(run_info[1].time, :t; input=input, it=it)
                          title = lift(i->string("t = ", time[i]), frame_index)
                      else
-                         time = select_slice(ri.time, :t; input=input, it=it)
-                         title = lift(i->join((string("t", irun, " = ", time[i])
+                         title = lift(i->join((string("t", irun, " = ",
+                                                      select_slice(ri.time, :t; input=input, it=it)[i])
                                                for (irun,ri) ∈ enumerate(run_info)), "; "),
                                       frame_index)
                      end
@@ -2136,8 +2136,8 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
 
                      if length(run_info) > 1
                          title = get_variable_symbol(var_name)
-                         time = select_slice(ri.time, :t; input=input, it=it)
-                         subtitles = (lift(i->string(ri.run_name, "\nt = ", time[i]),
+                         subtitles = (lift(i->string(ri.run_name, "\nt = ",
+                                                     select_slice(ri.time, :t; input=input, it=it)[i]),
                                            frame_index)
                                       for ri ∈ run_info)
                      else

From 80ad00a2432cfc724bb232d6b5a209b2b20a64ff Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Jul 2024 20:05:02 +0100
Subject: [PATCH 395/428] Fix post-processing electron moment gradients when
 `ir` is an integer

Also replace remaining uses of `postproc_load_variable()` with
`get_variable()` for consistency.
---
 moment_kinetics/src/load_data.jl | 65 +++++++++++++++++++++++---------
 1 file changed, 47 insertions(+), 18 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 3d0923a07..0c7e35d5a 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4123,43 +4123,72 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
     end
 
     if variable_name == "temperature"
-        vth = postproc_load_variable(run_info, "thermal_speed"; kwargs...)
+        vth = get_variable(run_info, "thermal_speed"; kwargs...)
         variable = vth.^2
     elseif variable_name == "collision_frequency_ii"
-        n = postproc_load_variable(run_info, "density"; kwargs...)
-        vth = postproc_load_variable(run_info, "thermal_speed"; kwargs...)
+        n = get_variable(run_info, "density"; kwargs...)
+        vth = get_variable(run_info, "thermal_speed"; kwargs...)
         variable = get_collision_frequency_ii(run_info.collisions, n, vth)
     elseif variable_name == "collision_frequency_ee"
-        n = postproc_load_variable(run_info, "electron_density"; kwargs...)
-        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        n = get_variable(run_info, "electron_density"; kwargs...)
+        vth = get_variable(run_info, "electron_thermal_speed"; kwargs...)
         variable = get_collision_frequency_ee(run_info.collisions, n, vth)
     elseif variable_name == "collision_frequency_ei"
-        n = postproc_load_variable(run_info, "electron_density"; kwargs...)
-        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        n = get_variable(run_info, "electron_density"; kwargs...)
+        vth = get_variable(run_info, "electron_thermal_speed"; kwargs...)
         variable = get_collision_frequency_ei(run_info.collisions, n, vth)
     elseif variable_name == "electron_temperature"
-        vth = postproc_load_variable(run_info, "electron_thermal_speed"; kwargs...)
+        vth = get_variable(run_info, "electron_thermal_speed"; kwargs...)
         variable = run_info.composition.me_over_mi .* vth.^2
     elseif variable_name == "electron_dudz"
-        upar = postproc_load_variable(run_info, "electron_parallel_flow"; kwargs...)
+        upar = get_variable(run_info, "electron_parallel_flow"; kwargs...)
         variable = similar(upar)
-        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
-            @views derivative!(variable[:,ir,it], upar[:,ir,it], run_info.z, run_info.z_spectral)
+        if :iz ∈ keys(kwargs) && kwargs[:iz] !== nothing
+            error("Cannot take z-derivative when iz!==nothing")
+        end
+        if :ir ∈ keys(kwargs) && isa(kwargs[:ir], mk_int)
+            for it ∈ 1:size(variable, 2)
+                @views derivative!(variable[:,it], upar[:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 3), ir ∈ 1:run_info.r.n
+                @views derivative!(variable[:,ir,it], upar[:,ir,it], run_info.z, run_info.z_spectral)
+            end
         end
     elseif variable_name == "electron_dpdz"
-        ppar = postproc_load_variable(run_info, "electron_parallel_pressure"; kwargs...)
+        ppar = get_variable(run_info, "electron_parallel_pressure"; kwargs...)
         variable = similar(ppar)
-        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
-            @views derivative!(variable[:,ir,it], ppar[:,ir,it], run_info.z, run_info.z_spectral)
+        println("electron_dpdz kwargs ", kwargs)
+        if :iz ∈ keys(kwargs) && kwargs[:iz] !== nothing
+            error("Cannot take z-derivative when iz!==nothing")
+        end
+        if :ir ∈ keys(kwargs) && isa(kwargs[:ir], mk_int)
+            println("check range ", (kwargs[:it] === nothing ? (1:run_info.nt) : (1:length(kwargs[:it]))))
+            for it ∈ 1:size(variable, 2)
+                @views derivative!(variable[:,it], ppar[:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 3), ir ∈ 1:run_info.r.n
+                @views derivative!(variable[:,ir,it], ppar[:,ir,it], run_info.z, run_info.z_spectral)
+            end
         end
     elseif variable_name == "electron_dqdz"
-        qpar = postproc_load_variable(run_info, "electron_parallel_heat_flux"; kwargs...)
+        qpar = get_variable(run_info, "electron_parallel_heat_flux"; kwargs...)
         variable = similar(qpar)
-        for it ∈ 1:run_info.nt, ir ∈ 1:run_info.r.n
-            @views derivative!(variable[:,ir,it], qpar[:,ir,it], run_info.z, run_info.z_spectral)
+        if :iz ∈ keys(kwargs) && kwargs[:iz] !== nothing
+            error("Cannot take z-derivative when iz!==nothing")
+        end
+        if :ir ∈ keys(kwargs) && isa(kwargs[:ir], mk_int)
+            for it ∈ 1:size(variable, 2)
+                @views derivative!(variable[:,it], qpar[:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 3), ir ∈ 1:run_info.r.n
+                @views derivative!(variable[:,ir,it], qpar[:,ir,it], run_info.z, run_info.z_spectral)
+            end
         end
     elseif variable_name == "temperature_neutral"
-        vth = postproc_load_variable(run_info, "thermal_speed_neutral"; kwargs...)
+        vth = get_variable(run_info, "thermal_speed_neutral"; kwargs...)
         variable = vth.^2
     elseif variable_name == "sound_speed"
         T_e = run_info.composition.T_e

From c0d14b02dfc5c68612a048157d147b66021d6007 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 23 Jul 2024 20:07:58 +0100
Subject: [PATCH 396/428] Use `calculate_electron_moment_derivatives!()` in
 `update_electron_pdf_with_time_advance!()`

---
 .../src/electron_kinetic_equation.jl          | 77 ++++++-------------
 moment_kinetics/src/velocity_moments.jl       |  2 +
 2 files changed, 26 insertions(+), 53 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 5c9be2f5a..8b2c01443 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -165,28 +165,22 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     buffer_r_5 = @view scratch_dummy.buffer_rs_5[:,1]
     buffer_r_6 = @view scratch_dummy.buffer_rs_6[:,1]
 
-    # compute the z-derivative of the input electron parallel flow, needed for the electron kinetic equation
-    @views derivative_z!(moments.electron.dupar_dz, moments.electron.upar, buffer_r_1,
-                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-
-    # compute the z-derivative of the input electron parallel pressure, needed for the electron kinetic equation
-    @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, buffer_r_1,
-                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-
     begin_r_z_region()
     @loop_r_z ir iz begin
         # update the electron thermal speed using the updated electron parallel pressure
         moments.electron.vth[iz,ir] = sqrt(abs(2.0 * moments.electron.ppar[iz,ir] /
                                                 (moments.electron.dens[iz,ir] *
                                                  composition.me_over_mi)))
-        # update the z-derivative of the electron thermal speed from the z-derivatives of the electron density
-        # and parallel pressure
-        moments.electron.dvth_dz[iz,ir] =
-            0.5 * moments.electron.vth[iz,ir] *
-            (moments.electron.dppar_dz[iz,ir] / moments.electron.ppar[iz,ir] -
-             moments.electron.ddens_dz[iz,ir] / moments.electron.dens[iz,ir])
         scratch[t_params.n_rk_stages+1].electron_ppar[iz,ir] = moments.electron.ppar[iz,ir]
     end
+    calculate_electron_moment_derivatives!(moments,
+                                           (electron_density=moments.electron.dens,
+                                            electron_upar=moments.electron.upar,
+                                            electron_ppar=moments.electron.ppar),
+                                           scratch_dummy, z, z_spectral,
+                                           num_diss_params.electron.moment_dissipation_coefficient,
+                                           composition.electron_physics)
+
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
@@ -358,26 +352,13 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                                   scratch[istage+1].electron_ppar,
                                                   moments.electron.vth, latest_pdf, vpa)
 
-                # compute the z-derivative of the parallel electron heat flux
-                @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar,
-                                     buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
-                                     z_spectral, z)
-
                 if evolve_ppar
-                    # get an updated iterate of the electron parallel pressure
-                    begin_r_z_region()
-                    # compute the z-derivative of the updated electron parallel pressure
-                    @views derivative_z!(moments.electron.dppar_dz,
-                                         scratch[istage+1].electron_ppar, buffer_r_1, buffer_r_2,
-                                         buffer_r_3, buffer_r_4, z_spectral, z)
-
                     this_ppar = scratch[istage+1].electron_ppar
-                    this_dppar_dz = moments.electron.dppar_dz
-                    this_ddens_dz = moments.electron.ddens_dz
                     this_dens = moments.electron.dens
-                    this_vth = moments.electron.vth
-                    this_dvth_dz = moments.electron.dvth_dz
+                    this_upar = moments.electron.upar
                     if update_vth
+                        begin_r_z_region()
+                        this_vth = moments.electron.vth
                         @loop_r_z ir iz begin
                             # update the electron thermal speed using the updated electron
                             # parallel pressure
@@ -386,29 +367,19 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                                         composition.me_over_mi)))
                         end
                     end
-                    @loop_r_z ir iz begin
-                        # update the z-derivative of the electron thermal speed from the
-                        # z-derivatives of the electron density and parallel pressure
-                        this_dvth_dz[iz,ir] = 0.5 * this_vth[iz,ir] *
-                                              (this_dppar_dz[iz,ir] / this_ppar[iz,ir] -
-                                               this_ddens_dz[iz,ir] / this_dens[iz,ir])
-                    end
-
-                    # centred second derivative for dissipation
-                    if num_diss_params.electron.moment_dissipation_coefficient > 0.0
-                        @views derivative_z!(moments.electron.d2ppar_dz2,
-                                             moments.electron.dppar_dz, buffer_r_1, buffer_r_2,
-                                             buffer_r_3, buffer_r_4, z_spectral, z)
-                        begin_serial_region()
-                        @serial_region begin
-                            if z.irank == 0
-                                moments.electron.d2ppar_dz2[1,:] .= 0.0
-                            end
-                            if z.irank == z.nrank - 1
-                                moments.electron.d2ppar_dz2[end,:] .= 0.0
-                            end
-                        end
-                    end
+                    calculate_electron_moment_derivatives!(
+                        moments,
+                        (electron_density=this_dens,
+                         electron_upar=this_upar,
+                         electron_ppar=this_ppar),
+                        scratch_dummy, z, z_spectral,
+                        num_diss_params.electron.moment_dissipation_coefficient,
+                        composition.electron_physics)
+                else
+                    # compute the z-derivative of the parallel electron heat flux
+                    @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar,
+                                         buffer_r_1, buffer_r_2, buffer_r_3, buffer_r_4,
+                                         z_spectral, z)
                 end
             end
             update_derived_moments_and_derivatives()
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 6cd1f6d34..4881d0aad 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -1000,6 +1000,8 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
     end
     @views derivative_z!(moments.electron.dT_dz, dummy_zr, buffer_r_1,
                             buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    @views derivative_z!(moments.electron.dvth_dz, moments.electron.vth, buffer_r_1,
+                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
 end
 
 """

From bf8fe31a0a9d9bb93702d043d2de7dbf7f1bce5f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Jul 2024 12:24:01 +0100
Subject: [PATCH 397/428] Allow keyboard interrupt to actually stop
 `makie_post_process()`

Previously the exception handling meant that the `InterruptException`
was caught and only stopped the low-level function, so the keyboard
interrupt had to be pressed many times to actually get back to the REPL.
---
 .../src/makie_post_processing.jl              | 103 +++++++++++++-----
 1 file changed, 75 insertions(+), 28 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index a284c7294..f9b39fc77 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -819,6 +819,15 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
     return nothing
 end
 
+function makie_post_processing_error_handler(e::Exception, message::String)
+    if isa(e, InterruptException)
+        rethrow(e)
+    else
+        println(message * "\nError was $e.")
+        return nothing
+    end
+end
+
 """
     get_run_info(run_dir...; itime_min=1, itime_max=0,
                  itime_skip=1, dfns=false, initial_electron=false, do_setup=true,
@@ -1009,9 +1018,9 @@ function plots_for_variable(run_info, variable_name; plot_prefix, has_rdim=true,
     try
         variable = get_variable(run_info, variable_name)
     catch e
-        println("plots_for_variable() failed for $variable_name - could not load data. "
-                * "Error was $e")
-        return nothing
+        return makie_post_processing_error_handler(
+                   e,
+                   "plots_for_variable() failed for $variable_name - could not load data.")
     end
 
     if variable_name ∈ em_variables
@@ -1488,8 +1497,9 @@ for dim ∈ one_dimension_combinations
                      end
                      return fig
                  catch e
-                     println("$($function_name_str) failed for $var_name, is=$is. Error was $e")
-                     return nothing
+                     return makie_post_processing_error_handler(
+                                e,
+                                "$($function_name_str) failed for $var_name, is=$is.")
                  end
              end
 
@@ -1673,8 +1683,9 @@ for (dim1, dim2) ∈ two_dimension_combinations
                      end
                      return fig
                  catch e
-                     println("$($function_name_str) failed for $var_name, is=$is. Error was $e")
-                     return nothing
+                     return makie_post_processing_error_handler(
+                                e,
+                                "$($function_name_str) failed for $var_name, is=$is.")
                  end
              end
 
@@ -1937,8 +1948,9 @@ for dim ∈ one_dimension_combinations_no_t
 
                      return fig
                  catch e
-                     println("$($function_name_str)() failed for $var_name, is=$is. Error was $e")
-                     return nothing
+                     return makie_post_processing_error_handler(
+                                e,
+                                "$($function_name_str)() failed for $var_name, is=$is.")
                  end
              end
 
@@ -2167,8 +2179,9 @@ for (dim1, dim2) ∈ two_dimension_combinations_no_t
 
                      return fig
                  catch e
-                     println("$($function_name_str) failed for $var_name, is=$is. Error was $e")
-                     return nothing
+                     return makie_post_processing_error_handler(
+                                e,
+                                "$($function_name_str) failed for $var_name, is=$is.")
                  end
              end
 
@@ -3656,7 +3669,9 @@ function _save_residual_plots(fig_axes, plot_prefix)
             save(plot_prefix * replace(key, " "=>"_") * ".pdf", fa[1])
         end
     catch e
-        println("Error in _save_residual_plots(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in _save_residual_plots().")
     end
 end
 
@@ -3709,7 +3724,9 @@ function calculate_steady_state_residual(run_info::Tuple, variable_name; is=1,
 
         return fig_axes
     catch e
-        println("Error in calculate_steady_state_residual(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in calculate_steady_state_residual().")
     end
 end
 
@@ -3840,7 +3857,9 @@ function plot_f_unnorm_vs_vpa(run_info::Tuple; f_over_vpa2=false, neutral=false,
 
         return fig
     catch e
-        println("Error in plot_f_unnorm_vs_vpa(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in plot_f_unnorm_vs_vpa().")
     end
 end
 
@@ -4005,7 +4024,9 @@ function plot_f_unnorm_vs_vpa_z(run_info::Tuple; neutral=false, outfile=nothing,
 
         return fig
     catch e
-        println("Error in plot_f_unnorm_vs_vpa_z(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in plot_f_unnorm_vs_vpa_z().")
     end
 end
 
@@ -4171,7 +4192,9 @@ function animate_f_unnorm_vs_vpa(run_info::Tuple; f_over_vpa2=false, neutral=fal
 
         return fig
     catch e
-        println("Error in animate_f_unnorm_vs_vpa(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in animate_f_unnorm_vs_vpa().")
     end
 end
 
@@ -4380,7 +4403,9 @@ function animate_f_unnorm_vs_vpa_z(run_info::Tuple; neutral=false, outfile=nothi
 
         return fig
     catch e
-        println("Error in animate_f_unnorm_vs_vpa_z(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in animate_f_unnorm_vs_vpa_z().")
     end
 end
 
@@ -4755,7 +4780,9 @@ function plot_charged_pdf_2D_at_wall(run_info; plot_prefix, electron=false)
             end
         end
     catch e
-        println("Error in plot_charged_pdf_2D_at_wall(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in plot_charged_pdf_2D_at_wall().")
     end
 
     return nothing
@@ -5046,7 +5073,9 @@ function plot_neutral_pdf_2D_at_wall(run_info; plot_prefix)
             end
         end
     catch e
-        println("Error in plot_neutral_pdf_2D_at_wall(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in plot_neutral_pdf_2D_at_wall().")
     end
 
     return nothing
@@ -5380,7 +5409,9 @@ function constraints_plots(run_info; plot_prefix=plot_prefix)
             #end
         end
     catch e
-        println("Error in constraints_plots(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in constraints_plots().")
     end
 end
 
@@ -5546,7 +5577,9 @@ function Chodura_condition_plots(run_info::Tuple; plot_prefix)
             save(outfile, fig)
         end
     catch e
-        println("Error in Chodura_condition_plots(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in Chodura_condition_plots().")
     end
 
     return nothing
@@ -5736,7 +5769,9 @@ function sound_wave_plots(run_info::Tuple; plot_prefix)
             return fig
         end
     catch e
-        println("Error in sound_wave_plots(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in sound_wave_plots().")
     end
 
     return nothing
@@ -6126,7 +6161,9 @@ function instability2D_plots(run_info, variable_name; plot_prefix, zind=nothing,
             plot_Fourier_1D(variable_Fourier_1D, get_variable_symbol(variable_name),
                             variable_name)
         catch e
-            println("Warning: error in 1D Fourier analysis for $variable_name. Error was $e")
+            return makie_post_processing_error_handler(
+                       e,
+                       "Warning: error in 1D Fourier analysis for $variable_name.")
         end
 
         # Do this to allow memory to be garbage-collected.
@@ -6185,7 +6222,9 @@ function instability2D_plots(run_info, variable_name; plot_prefix, zind=nothing,
             plot_Fourier_2D(variable_Fourier, get_variable_symbol(variable_name),
                             variable_name)
         catch e
-            println("Warning: error in 2D Fourier analysis for $variable_name. Error was $e")
+            return makie_post_processing_error_handler(
+                       e,
+                       "Warning: error in 2D Fourier analysis for $variable_name.")
         end
 
         # Do this to allow memory to be garbage-collected.
@@ -6213,7 +6252,9 @@ function instability2D_plots(run_info, variable_name; plot_prefix, zind=nothing,
                        colorbar_place=colorbar_place, frame_index=frame_index,
                        outfile=outfile)
         catch e
-            println("Warning: error in perturbation animation for $variable_name. Error was $e")
+            return makie_post_processing_error_handler(
+                       e,
+                       "Warning: error in perturbation animation for $variable_name.")
         end
 
         # Do this to allow memory to be garbage-collected (although this is redundant
@@ -7170,7 +7211,9 @@ function manufactured_solutions_analysis(run_info::Tuple; plot_prefix, nvperp)
         return manufactured_solutions_analysis(run_info[1]; plot_prefix=plot_prefix,
                                                nvperp=nvperp)
     catch e
-        println("Error in manufactured_solutions_analysis(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in manufactured_solutions_analysis().")
     end
 end
 
@@ -7259,7 +7302,9 @@ function manufactured_solutions_analysis_dfns(run_info::Tuple; plot_prefix)
     try
         return manufactured_solutions_analysis_dfns(run_info[1]; plot_prefix=plot_prefix)
     catch e
-        println("Error in manufactured_solutions_analysis_dfns(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in manufactured_solutions_analysis_dfns().")
     end
 end
 
@@ -7981,7 +8026,9 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
 
         return steps_fig, dt_fig, CFL_fig
     catch e
-        println("Error in timestep_diagnostics(). Error was ", e)
+        return makie_post_processing_error_handler(
+                   e,
+                   "Error in timestep_diagnostics().")
     end
 end
 

From 87345176f86fbaa8701e037604ab9994e1ef1699 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 26 Jul 2024 12:43:33 +0100
Subject: [PATCH 398/428] Fix output counters when simulation is stopped with
 `stopnow`

---
 moment_kinetics/src/time_advance.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 6ea09e764..ca5529a55 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -1849,8 +1849,8 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
         end
 
         if t_params.adaptive && !t_params.write_after_fixed_step_count
-            write_moments = t_params.write_moments_output[]
-            write_dfns = t_params.write_dfns_output[]
+            write_moments = t_params.write_moments_output[] || finish_now
+            write_dfns = t_params.write_dfns_output[] || finish_now
 
             _block_synchronize()
             @serial_region begin
@@ -1859,9 +1859,11 @@ function time_advance!(pdf, scratch, scratch_implicit, scratch_electron, t_param
             end
         else
             write_moments = (t_params.step_counter[] % t_params.nwrite_moments == 0
-                             || t_params.step_counter[] >= t_params.nstep)
+                             || t_params.step_counter[] >= t_params.nstep
+                             || finish_now)
             write_dfns = (t_params.step_counter[] % t_params.nwrite_dfns == 0
-                          || t_params.step_counter[] >= t_params.nstep)
+                          || t_params.step_counter[] >= t_params.nstep
+                          || finish_now)
         end
         if write_moments
             t_params.moments_output_counter[] += 1

From d37226eca32d04a84130c36ddaf6cd167eb38f4e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 27 Jul 2024 13:49:03 +0100
Subject: [PATCH 399/428] Rename vmax->vcut in electron sheath entrance
 boundary condition

Is a bit clearer, and will help when we want to distinguish between
-vcut and +vcut.
---
 .../src/electron_kinetic_equation.jl          | 128 +++++++++---------
 1 file changed, 64 insertions(+), 64 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 8b2c01443..458351960 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -918,15 +918,15 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # Initial guess for cut-off velocity is result from previous RK stage (which
             # might be the previous timestep if this is the first stage). Recalculate this
             # value from phi.
-            vmax = sqrt(phi[1,ir] / me_over_mi)
+            vcut = sqrt(phi[1,ir] / me_over_mi)
 
-            # -vmax is between vmax_ind-1 and vmax_ind
-            vmax_ind = searchsortedfirst(vpa_unnorm, -vmax)
-            if vmax_ind < 2
-                error("In lower-z electron bc, failed to find vpa=-vmax point, vmax_ind=$vmax_ind")
+            # -vcut is between minus_vcut_ind-1 and minus_vcut_ind
+            minus_vcut_ind = searchsortedfirst(vpa_unnorm, -vcut)
+            if minus_vcut_ind < 2
+                error("In lower-z electron bc, failed to find vpa=-vcut point, minus_vcut_ind=$minus_vcut_ind")
             end
-            if vmax_ind > vpa.n
-                error("In lower-z electron bc, failed to find vpa=-vmax point, vmax_ind=$vmax_ind")
+            if minus_vcut_ind > vpa.n
+                error("In lower-z electron bc, failed to find vpa=-vcut point, minus_vcut_ind=$minus_vcut_ind")
             end
 
             # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
@@ -956,19 +956,19 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[1,ir]
             quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[1,ir]
 
-            function get_integrals_and_derivatives(vmax, vmax_ind)
-                # vmax_fraction is the fraction of the distance between vmax_ind-1 and
-                # vmax_ind where -vmax is.
-                vmax_fraction = (-vmax - vpa_unnorm[vmax_ind-1]) / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind-1])
+            function get_integrals_and_derivatives(vcut, minus_vcut_ind)
+                # vcut_fraction is the fraction of the distance between minus_vcut_ind-1 and
+                # minus_vcut_ind where -vcut is.
+                vcut_fraction = (-vcut - vpa_unnorm[minus_vcut_ind-1]) / (vpa_unnorm[minus_vcut_ind] - vpa_unnorm[minus_vcut_ind-1])
 
                 function get_for_one_moment(integral_pieces, skip_part2=false)
-                    # Integral contribution from the cell containing vmax
-                    integral_vmax_cell = (0.5 * integral_pieces[vmax_ind-1] + 0.5 * integral_pieces[vmax_ind])
+                    # Integral contribution from the cell containing vcut
+                    integral_vcut_cell = (0.5 * integral_pieces[minus_vcut_ind-1] + 0.5 * integral_pieces[minus_vcut_ind])
 
-                    part1 = sum(integral_pieces[1:vmax_ind-2])
-                    part1 += 0.5 * integral_pieces[vmax_ind-1] + vmax_fraction * integral_vmax_cell
-                    # part1prime is d(part1)/d(vmax)
-                    part1prime = -integral_vmax_cell / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind-1])
+                    part1 = sum(integral_pieces[1:minus_vcut_ind-2])
+                    part1 += 0.5 * integral_pieces[minus_vcut_ind-1] + vcut_fraction * integral_vcut_cell
+                    # part1prime is d(part1)/d(vcut)
+                    part1prime = -integral_vcut_cell / (vpa_unnorm[minus_vcut_ind] - vpa_unnorm[minus_vcut_ind-1])
 
                     if skip_part2
                         part2 = nothing
@@ -977,9 +977,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                         # Integral contribution from the cell containing sigma
                         integral_sigma_cell = (0.5 * integral_pieces[sigma_ind-1] + 0.5 * integral_pieces[sigma_ind])
 
-                        part2 = sum(integral_pieces[vmax_ind+1:sigma_ind-2])
-                        part2 += (1.0 - vmax_fraction) * integral_vmax_cell + 0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[sigma_ind-1] + sigma_fraction * integral_sigma_cell
-                        # part2prime is d(part2)/d(vmax)
+                        part2 = sum(integral_pieces[minus_vcut_ind+1:sigma_ind-2])
+                        part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[minus_vcut_ind] + 0.5 * integral_pieces[sigma_ind-1] + sigma_fraction * integral_sigma_cell
+                        # part2prime is d(part2)/d(vcut)
                         part2prime = -part1prime
                     end
 
@@ -1001,11 +1001,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             counter = 1
             A = 1.0
             C = 0.0
-            # Always do at least one update of vmax
-            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+            # Always do at least one update of vcut
+            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, minus_vcut_ind)
             while true
                 # Newton iteration update. Note that primes denote derivatives with
-                # respect to vmax
+                # respect to vcut
                 delta_v = - epsilon / epsilonprime
 
                 # Prevent the step size from getting too big, to make Newton iteration
@@ -1013,10 +1013,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                 delta_v = min(delta_v, 0.1 * vthe[1,ir])
                 delta_v = max(delta_v, -0.1 * vthe[1,ir])
 
-                vmax = vmax + delta_v
-                vmax_ind = searchsortedfirst(vpa_unnorm, -vmax)
+                vcut = vcut + delta_v
+                minus_vcut_ind = searchsortedfirst(vpa_unnorm, -vcut)
 
-                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, minus_vcut_ind)
 
                 if abs(epsilon) < newton_tol * abs(u_over_vt)
                     break
@@ -1053,15 +1053,15 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
             reverse!(reversed_pdf)
 
-            ivpa_max = searchsortedlast(vpa_unnorm, vmax)
+            ivpa_max = searchsortedlast(vpa_unnorm, vcut)
             reversed_pdf[ivpa_max+1:end] .= 0.0
-            # vmax_fraction is the fraction of the distance between ivpa_max and
-            # ivpa_max+1 where vmax is.
-            vmax_fraction = (vmax - vpa_unnorm[ivpa_max]) / (vpa_unnorm[ivpa_max+1] - vpa_unnorm[ivpa_max])
-            reversed_pdf[ivpa_max] *= vmax_fraction
+            # vcut_fraction is the fraction of the distance between ivpa_max and
+            # ivpa_max+1 where vcut is.
+            vcut_fraction = (vcut - vpa_unnorm[ivpa_max]) / (vpa_unnorm[ivpa_max+1] - vpa_unnorm[ivpa_max])
+            reversed_pdf[ivpa_max] *= vcut_fraction
 
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-            phi[1,ir] = me_over_mi * vmax^2
+            phi[1,ir] = me_over_mi * vcut^2
             pdf[sigma_ind:end,1,1,ir] .= reversed_pdf[sigma_ind:end]
 
             moments.electron.constraints_A_coefficient[1,ir] = A
@@ -1095,15 +1095,15 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # Initial guess for cut-off velocity is result from previous RK stage (which
             # might be the previous timestep if this is the first stage). Recalculate this
             # value from phi.
-            vmax = sqrt(phi[end,ir] / me_over_mi)
+            vcut = sqrt(phi[end,ir] / me_over_mi)
 
-            # -vmax is between vmax_ind and vmax_ind+1
-            vmax_ind = searchsortedlast(vpa_unnorm, vmax)
-            if vmax_ind < 1
-                error("In upper-z electron bc, failed to find vpa=vmax point, vmax_ind=$vmax_ind")
+            # vcut is between plus_vcut_ind and plus_vcut_ind+1
+            plus_vcut_ind = searchsortedlast(vpa_unnorm, vcut)
+            if plus_vcut_ind < 1
+                error("In upper-z electron bc, failed to find vpa=vcut point, plus_vcut_ind=$plus_vcut_ind")
             end
-            if vmax_ind > vpa.n - 1
-                error("In upper-z electron bc, failed to find vpa=vmax point, vmax_ind=$vmax_ind")
+            if plus_vcut_ind > vpa.n - 1
+                error("In upper-z electron bc, failed to find vpa=vcut point, plus_vcut_ind=$plus_vcut_ind")
             end
 
             # sigma is the location we use for w_∥(v_∥=0) - set to 0 to ignore the 'upar
@@ -1133,19 +1133,19 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[end,ir]
             quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[end,ir]
 
-            function get_integrals_and_derivatives(vmax, vmax_ind)
-                # vmax_fraction is the fraction of the distance between vmax_ind and
-                # vmax_ind+1 where vmax is.
-                vmax_fraction = (vmax - vpa_unnorm[vmax_ind+1]) / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind+1])
+            function get_integrals_and_derivatives(vcut, plus_vcut_ind)
+                # vcut_fraction is the fraction of the distance between plus_vcut_ind and
+                # plus_vcut_ind+1 where vcut is.
+                vcut_fraction = (vcut - vpa_unnorm[plus_vcut_ind+1]) / (vpa_unnorm[plus_vcut_ind] - vpa_unnorm[plus_vcut_ind+1])
 
                 function get_for_one_moment(integral_pieces, skip_part2=false)
-                    # Integral contribution from the cell containing vmax
-                    integral_vmax_cell = (0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[vmax_ind+1])
+                    # Integral contribution from the cell containing vcut
+                    integral_vcut_cell = (0.5 * integral_pieces[plus_vcut_ind] + 0.5 * integral_pieces[plus_vcut_ind+1])
 
-                    part1 = sum(integral_pieces[vmax_ind+2:end])
-                    part1 += 0.5 * integral_pieces[vmax_ind+1] + vmax_fraction * integral_vmax_cell
-                    # part1prime is d(part1)/d(vmax)
-                    part1prime = integral_vmax_cell / (vpa_unnorm[vmax_ind] - vpa_unnorm[vmax_ind+1])
+                    part1 = sum(integral_pieces[plus_vcut_ind+2:end])
+                    part1 += 0.5 * integral_pieces[plus_vcut_ind+1] + vcut_fraction * integral_vcut_cell
+                    # part1prime is d(part1)/d(vcut)
+                    part1prime = integral_vcut_cell / (vpa_unnorm[plus_vcut_ind] - vpa_unnorm[plus_vcut_ind+1])
 
                     if skip_part2
                         part2 = nothing
@@ -1154,9 +1154,9 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                         # Integral contribution from the cell containing sigma
                         integral_sigma_cell = (0.5 * integral_pieces[sigma_ind] + 0.5 * integral_pieces[sigma_ind+1])
 
-                        part2 = sum(integral_pieces[sigma_ind+2:vmax_ind-1])
-                        part2 += (1.0 - vmax_fraction) * integral_vmax_cell + 0.5 * integral_pieces[vmax_ind] + 0.5 * integral_pieces[sigma_ind+1] + sigma_fraction * integral_sigma_cell
-                        # part2prime is d(part2)/d(vmax)
+                        part2 = sum(integral_pieces[sigma_ind+2:plus_vcut_ind-1])
+                        part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[plus_vcut_ind] + 0.5 * integral_pieces[sigma_ind+1] + sigma_fraction * integral_sigma_cell
+                        # part2prime is d(part2)/d(vcut)
                         part2prime = -part1prime
                     end
 
@@ -1176,11 +1176,11 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             end
 
             counter = 1
-            # Always do at least one update of vmax
-            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+            # Always do at least one update of vcut
+            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, plus_vcut_ind)
             while true
                 # Newton iteration update. Note that primes denote derivatives with
-                # respect to vmax
+                # respect to vcut
                 delta_v = - epsilon / epsilonprime
 
                 # Prevent the step size from getting too big, to make Newton iteration
@@ -1188,10 +1188,10 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                 delta_v = min(delta_v, 0.1 * vthe[1,ir])
                 delta_v = max(delta_v, -0.1 * vthe[1,ir])
 
-                vmax = vmax + delta_v
-                vmax_ind = searchsortedlast(vpa_unnorm, vmax)
+                vcut = vcut + delta_v
+                plus_vcut_ind = searchsortedlast(vpa_unnorm, vcut)
 
-                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vmax, vmax_ind)
+                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, plus_vcut_ind)
 
                 if abs(epsilon) < newton_tol * abs(u_over_vt)
                     break
@@ -1228,15 +1228,15 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
             reverse!(reversed_pdf)
 
-            ivpa_min = searchsortedfirst(vpa_unnorm, -vmax)
+            ivpa_min = searchsortedfirst(vpa_unnorm, -vcut)
             reversed_pdf[1:ivpa_min-1] .= 0.0
-            # vmax_fraction is the fraction of the distance between ivpa_min and
-            # ivpa_min-1 where -vmax is.
-            vmax_fraction = (-vmax - vpa_unnorm[ivpa_min]) / (vpa_unnorm[ivpa_min-1] - vpa_unnorm[ivpa_min])
-            reversed_pdf[ivpa_min] *= vmax_fraction
+            # vcut_fraction is the fraction of the distance between ivpa_min and
+            # ivpa_min-1 where -vcut is.
+            vcut_fraction = (-vcut - vpa_unnorm[ivpa_min]) / (vpa_unnorm[ivpa_min-1] - vpa_unnorm[ivpa_min])
+            reversed_pdf[ivpa_min] *= vcut_fraction
 
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
-            phi[end,ir] = me_over_mi * vmax^2
+            phi[end,ir] = me_over_mi * vcut^2
             pdf[1:sigma_ind,1,end,ir] .= reversed_pdf[1:sigma_ind]
 
             moments.electron.constraints_A_coefficient[end,ir] = A

From 0651621483fe15c13c1fcc43991c6c75c0431b15 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 28 Jul 2024 14:06:34 +0100
Subject: [PATCH 400/428] Apply moment corrections to 'reflected' part of
 electron f at sheath

---
 .../src/electron_kinetic_equation.jl          | 240 +++++++++++++++---
 1 file changed, 200 insertions(+), 40 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 458351960..36d87e88d 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -956,12 +956,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[1,ir]
             quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[1,ir]
 
-            function get_integrals_and_derivatives(vcut, minus_vcut_ind)
+            function get_integrals_and_derivatives_lowerz(vcut, minus_vcut_ind)
                 # vcut_fraction is the fraction of the distance between minus_vcut_ind-1 and
                 # minus_vcut_ind where -vcut is.
                 vcut_fraction = (-vcut - vpa_unnorm[minus_vcut_ind-1]) / (vpa_unnorm[minus_vcut_ind] - vpa_unnorm[minus_vcut_ind-1])
 
-                function get_for_one_moment(integral_pieces, skip_part2=false)
+                function get_for_one_moment(integral_pieces)
                     # Integral contribution from the cell containing vcut
                     integral_vcut_cell = (0.5 * integral_pieces[minus_vcut_ind-1] + 0.5 * integral_pieces[minus_vcut_ind])
 
@@ -970,39 +970,35 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                     # part1prime is d(part1)/d(vcut)
                     part1prime = -integral_vcut_cell / (vpa_unnorm[minus_vcut_ind] - vpa_unnorm[minus_vcut_ind-1])
 
-                    if skip_part2
-                        part2 = nothing
-                        part2prime = nothing
-                    else
-                        # Integral contribution from the cell containing sigma
-                        integral_sigma_cell = (0.5 * integral_pieces[sigma_ind-1] + 0.5 * integral_pieces[sigma_ind])
-
-                        part2 = sum(integral_pieces[minus_vcut_ind+1:sigma_ind-2])
-                        part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[minus_vcut_ind] + 0.5 * integral_pieces[sigma_ind-1] + sigma_fraction * integral_sigma_cell
-                        # part2prime is d(part2)/d(vcut)
-                        part2prime = -part1prime
-                    end
+                    # Integral contribution from the cell containing sigma
+                    integral_sigma_cell = (0.5 * integral_pieces[sigma_ind-1] + 0.5 * integral_pieces[sigma_ind])
+
+                    part2 = sum(integral_pieces[minus_vcut_ind+1:sigma_ind-2])
+                    part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[minus_vcut_ind] + 0.5 * integral_pieces[sigma_ind-1] + sigma_fraction * integral_sigma_cell
+                    # part2prime is d(part2)/d(vcut)
+                    part2prime = -part1prime
 
                     return part1, part1prime, part2, part2prime
                 end
                 a1, a1prime, a2, a2prime = get_for_one_moment(density_integral_pieces)
-                b1, b1prime, _, _ = get_for_one_moment(flow_integral_pieces, true)
+                b1, b1prime, b2, _ = get_for_one_moment(flow_integral_pieces)
                 c1, c1prime, c2, c2prime = get_for_one_moment(energy_integral_pieces)
-                d1, d1prime, _, _ = get_for_one_moment(cubic_integral_pieces, true)
+                d1, d1prime, d2, _ = get_for_one_moment(cubic_integral_pieces)
                 e1, e1prime, e2, e2prime = get_for_one_moment(quartic_integral_pieces)
 
                 return get_residual_and_coefficients_for_bc(a1, a1prime, a2, a2prime, b1,
                                                             b1prime, c1, c1prime, c2,
                                                             c2prime, d1, d1prime, e1,
                                                             e1prime, e2, e2prime,
-                                                            u_over_vt)
+                                                            u_over_vt)...,
+                       a2, b2, c2, d2
             end
 
             counter = 1
             A = 1.0
             C = 0.0
             # Always do at least one update of vcut
-            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, minus_vcut_ind)
+            epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_lowerz(vcut, minus_vcut_ind)
             while true
                 # Newton iteration update. Note that primes denote derivatives with
                 # respect to vcut
@@ -1016,7 +1012,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                 vcut = vcut + delta_v
                 minus_vcut_ind = searchsortedfirst(vpa_unnorm, -vcut)
 
-                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, minus_vcut_ind)
+                epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_lowerz(vcut, minus_vcut_ind)
 
                 if abs(epsilon) < newton_tol * abs(u_over_vt)
                     break
@@ -1067,6 +1063,90 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             moments.electron.constraints_A_coefficient[1,ir] = A
             moments.electron.constraints_B_coefficient[1,ir] = 0.0
             moments.electron.constraints_C_coefficient[1,ir] = C
+
+            # Ensure the part of f for 0≤v_∥≤vcut has its first 3 moments symmetric with
+            # vcut≤v_∥≤0 (i.e. even moments are the same, odd moments are equal but opposite
+            # sign). This should be true analytically because of the definition of the
+            # boundary condition, but would not be numerically true because of the
+            # interpolation.
+
+            # Need to recalculate these with the updated distribution function
+            density_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.wgts / sqrt(pi)
+            flow_integral_pieces = @views @. vpa.scratch4 = density_integral_pieces * vpa_unnorm / vthe[1,ir]
+            energy_integral_pieces = @views @. vpa.scratch5 = flow_integral_pieces * vpa_unnorm / vthe[1,ir]
+            cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[1,ir]
+            quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[1,ir]
+
+            # Update the part2 integrals since we've applied the A and C factors
+            _, _, _, _, a2, b2, c2, d2 = get_integrals_and_derivatives_lowerz(vcut, minus_vcut_ind)
+
+            function get_part3_for_one_moment_lower(integral_pieces)
+                # Integral contribution from the cell containing sigma
+                integral_sigma_cell = (0.5 * integral_pieces[sigma_ind-1] + 0.5 * integral_pieces[sigma_ind])
+
+                @views part3 = sum(integral_pieces[sigma_ind+1:plus_vcut_ind+1])
+                part3 += 0.5 * integral_pieces[sigma_ind] + (1.0 - sigma_fraction) * integral_sigma_cell
+
+                return part3
+            end
+            a3 = get_part3_for_one_moment_lower(density_integral_pieces)
+            b3 = get_part3_for_one_moment_lower(flow_integral_pieces)
+            c3 = get_part3_for_one_moment_lower(energy_integral_pieces)
+            d3 = get_part3_for_one_moment_lower(cubic_integral_pieces)
+
+            correction0_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,1,ir] * vpa.wgts / sqrt(pi) * vpa_unnorm^2 / vthe[1,ir]^2 / (1.0 + vpa_unnorm^2 / vthe[1,ir]^2)
+            for ivpa ∈ 1:sigma_ind
+                # We only add the corrections to 'part3', so zero them out for negative v_∥.
+                # I think this is only actually significant for `sigma_ind-1` and
+                # `sigma_ind`. Even though `sigma_ind` is part of the distribution
+                # function that we are correcting, for v_∥>0, it affects the integral in
+                # the 'sigma_cell' between `sigma_ind-1` and `sigma_ind`, which would
+                # affect the numerically calculated integrals for f(v_∥<0), so if we
+                # 'corrected' its value, those integrals would change and the constraints
+                # would not be exactly satisfied. The difference should be small, as the
+                # correction at that point is multiplied by
+                # v_∥^2/vth^2/(1+v_∥^2/vth^2)≈v_∥^2/vth^2≈0.
+                correction0_integral_pieces[ivpa] = 0.0
+            end
+            correction1_integral_pieces = @views @. vpa.scratch4 = correction0_integral_pieces * vpa_unnorm / vthe[1,ir]
+            correction2_integral_pieces = @views @. vpa.scratch5 = correction1_integral_pieces * vpa_unnorm / vthe[1,ir]
+            correction3_integral_pieces = @views @. vpa.scratch6 = correction2_integral_pieces * vpa_unnorm / vthe[1,ir]
+            correction4_integral_pieces = @views @. vpa.scratch7 = correction3_integral_pieces * vpa_unnorm / vthe[1,ir]
+            correction5_integral_pieces = @views @. vpa.scratch8 = correction4_integral_pieces * vpa_unnorm / vthe[1,ir]
+            correction6_integral_pieces = @views @. vpa.scratch9 = correction5_integral_pieces * vpa_unnorm / vthe[1,ir]
+
+            alpha = get_part3_for_one_moment_lower(correction0_integral_pieces)
+            beta = get_part3_for_one_moment_lower(correction1_integral_pieces)
+            gamma = get_part3_for_one_moment_lower(correction2_integral_pieces)
+            delta = get_part3_for_one_moment_lower(correction3_integral_pieces)
+            epsilon = get_part3_for_one_moment_lower(correction4_integral_pieces)
+            zeta = get_part3_for_one_moment_lower(correction5_integral_pieces)
+            eta = get_part3_for_one_moment_lower(correction6_integral_pieces)
+
+            # Update the v_∥>0 part of f to correct the moments as
+            # f(0<v_∥<vcut) = (1 + (A + B*v/vth + C*v^2/vth^2 + D*v^3/vth^3) * v^2/vth^2 / (1 + v^2/vth^2)) * fhat(0<v_∥<vcut)
+            # Constraints:
+            # ∫dv_∥ v_∥^n f(0<v_∥<vcut) = (-1)^n ∫dv_∥ v_∥^n f(-vcut<v_∥<0)    for n=0,1,2,3
+            # a2 = a3 + alpha*A + beta*B + gamma*C + delta*D
+            # -b2 = b3 + beta*A + gamma*B + delta*C + epsilon*D
+            # c2 = c3 + gamma*A + delta*B + epsilon*C + zeta*D
+            # -d2 = d3 + delta*A + epsilon*B + zeta*C + eta*D
+            solution = [alpha beta    gamma   delta   ;
+                        beta  gamma   delta   epsilon ;
+                        gamma delta   epsilon zeta    ;
+                        delta epsilon zeta    eta
+                       ] \ [a2-a3, -b2-b3, c2-c3, -d2-d3]
+            A, B, C, D = solution
+            for ivpa ∈ sigma_ind+1:plus_vcut_ind+1
+                v_over_vth = vpa_unnorm[ivpa]/vthe[1,ir]
+                pdf[ivpa,1,1,ir] = pdf[ivpa,1,1,ir] +
+                                   (A
+                                    + B * v_over_vth
+                                    + C * v_over_vth^2
+                                    + D * v_over_vth^3) *
+                                   v_over_vth^2 / (1.0 + v_over_vth^2) *
+                                   pdf[ivpa,1,1,ir]
+            end
         end
     end
 
@@ -1133,12 +1213,12 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[end,ir]
             quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[end,ir]
 
-            function get_integrals_and_derivatives(vcut, plus_vcut_ind)
+            function get_integrals_and_derivatives_upperz(vcut, plus_vcut_ind)
                 # vcut_fraction is the fraction of the distance between plus_vcut_ind and
                 # plus_vcut_ind+1 where vcut is.
                 vcut_fraction = (vcut - vpa_unnorm[plus_vcut_ind+1]) / (vpa_unnorm[plus_vcut_ind] - vpa_unnorm[plus_vcut_ind+1])
 
-                function get_for_one_moment(integral_pieces, skip_part2=false)
+                function get_for_one_moment(integral_pieces)
                     # Integral contribution from the cell containing vcut
                     integral_vcut_cell = (0.5 * integral_pieces[plus_vcut_ind] + 0.5 * integral_pieces[plus_vcut_ind+1])
 
@@ -1147,37 +1227,33 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
                     # part1prime is d(part1)/d(vcut)
                     part1prime = integral_vcut_cell / (vpa_unnorm[plus_vcut_ind] - vpa_unnorm[plus_vcut_ind+1])
 
-                    if skip_part2
-                        part2 = nothing
-                        part2prime = nothing
-                    else
-                        # Integral contribution from the cell containing sigma
-                        integral_sigma_cell = (0.5 * integral_pieces[sigma_ind] + 0.5 * integral_pieces[sigma_ind+1])
-
-                        part2 = sum(integral_pieces[sigma_ind+2:plus_vcut_ind-1])
-                        part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[plus_vcut_ind] + 0.5 * integral_pieces[sigma_ind+1] + sigma_fraction * integral_sigma_cell
-                        # part2prime is d(part2)/d(vcut)
-                        part2prime = -part1prime
-                    end
+                    # Integral contribution from the cell containing sigma
+                    integral_sigma_cell = (0.5 * integral_pieces[sigma_ind] + 0.5 * integral_pieces[sigma_ind+1])
+
+                    part2 = sum(integral_pieces[sigma_ind+2:plus_vcut_ind-1])
+                    part2 += (1.0 - vcut_fraction) * integral_vcut_cell + 0.5 * integral_pieces[plus_vcut_ind] + 0.5 * integral_pieces[sigma_ind+1] + sigma_fraction * integral_sigma_cell
+                    # part2prime is d(part2)/d(vcut)
+                    part2prime = -part1prime
 
                     return part1, part1prime, part2, part2prime
                 end
                 a1, a1prime, a2, a2prime = get_for_one_moment(density_integral_pieces)
-                b1, b1prime, _, _ = get_for_one_moment(flow_integral_pieces, true)
+                b1, b1prime, b2, _ = get_for_one_moment(flow_integral_pieces)
                 c1, c1prime, c2, c2prime = get_for_one_moment(energy_integral_pieces)
-                d1, d1prime, _, _ = get_for_one_moment(cubic_integral_pieces, true)
+                d1, d1prime, d2, _ = get_for_one_moment(cubic_integral_pieces)
                 e1, e1prime, e2, e2prime = get_for_one_moment(quartic_integral_pieces)
 
                 return get_residual_and_coefficients_for_bc(a1, a1prime, a2, a2prime, b1,
                                                             b1prime, c1, c1prime, c2,
                                                             c2prime, d1, d1prime, e1,
                                                             e1prime, e2, e2prime,
-                                                            u_over_vt)
+                                                            u_over_vt)...,
+                       a2, b2, c2, d2
             end
 
             counter = 1
             # Always do at least one update of vcut
-            epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, plus_vcut_ind)
+            epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_upperz(vcut, plus_vcut_ind)
             while true
                 # Newton iteration update. Note that primes denote derivatives with
                 # respect to vcut
@@ -1185,13 +1261,13 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
                 # Prevent the step size from getting too big, to make Newton iteration
                 # more robust.
-                delta_v = min(delta_v, 0.1 * vthe[1,ir])
-                delta_v = max(delta_v, -0.1 * vthe[1,ir])
+                delta_v = min(delta_v, 0.1 * vthe[end,ir])
+                delta_v = max(delta_v, -0.1 * vthe[end,ir])
 
                 vcut = vcut + delta_v
                 plus_vcut_ind = searchsortedlast(vpa_unnorm, vcut)
 
-                epsilon, epsilonprime, A, C = get_integrals_and_derivatives(vcut, plus_vcut_ind)
+                epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_upperz(vcut, plus_vcut_ind)
 
                 if abs(epsilon) < newton_tol * abs(u_over_vt)
                     break
@@ -1242,6 +1318,90 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             moments.electron.constraints_A_coefficient[end,ir] = A
             moments.electron.constraints_B_coefficient[end,ir] = 0.0
             moments.electron.constraints_C_coefficient[end,ir] = C
+
+            # Ensure the part of f for -vcut≤v_∥≤0 has its first 3 moments symmetric with
+            # 0≤v_∥≤vcut  (i.e. even moments are the same, odd moments are equal but opposite
+            # sign). This should be true analytically because of the definition of the
+            # boundary condition, but would not be numerically true because of the
+            # interpolation.
+
+            # Need to recalculate these with the updated distribution function
+            density_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.wgts / sqrt(pi)
+            flow_integral_pieces = @views @. vpa.scratch4 = density_integral_pieces * vpa_unnorm / vthe[end,ir]
+            energy_integral_pieces = @views @. vpa.scratch5 = flow_integral_pieces * vpa_unnorm / vthe[end,ir]
+            cubic_integral_pieces = @views @. vpa.scratch6 = energy_integral_pieces * vpa_unnorm / vthe[end,ir]
+            quartic_integral_pieces = @views @. vpa.scratch7 = cubic_integral_pieces * vpa_unnorm / vthe[end,ir]
+
+            # Update the part2 integrals since we've applied the A and C factors
+            _, _, _, _, a2, b2, c2, d2 = get_integrals_and_derivatives_upperz(vcut, plus_vcut_ind)
+
+            function get_part3_for_one_moment_upper(integral_pieces)
+                # Integral contribution from the cell containing sigma
+                integral_sigma_cell = (0.5 * integral_pieces[sigma_ind] + 0.5 * integral_pieces[sigma_ind+1])
+
+                @views part3 = sum(integral_pieces[minus_vcut_ind-1:sigma_ind-1])
+                part3 += 0.5 * integral_pieces[sigma_ind] + (1.0 - sigma_fraction) * integral_sigma_cell
+
+                return part3
+            end
+            a3 = get_part3_for_one_moment_upper(density_integral_pieces)
+            b3 = get_part3_for_one_moment_upper(flow_integral_pieces)
+            c3 = get_part3_for_one_moment_upper(energy_integral_pieces)
+            d3 = get_part3_for_one_moment_upper(cubic_integral_pieces)
+
+            correction0_integral_pieces = @views @. vpa.scratch3 = pdf[:,1,end,ir] * vpa.wgts / sqrt(pi) * vpa_unnorm^2 / vthe[end,ir]^2 / (1.0 + vpa_unnorm^2 / vthe[end,ir]^2)
+            for ivpa ∈ sigma_ind:vpa.n
+                # We only add the corrections to 'part3', so zero them out for positive v_∥.
+                # I think this is only actually significant for `sigma_ind` and
+                # `sigma_ind+1`. Even though `sigma_ind` is part of the distribution
+                # function that we are correcting, for v_∥<0, it affects the integral in
+                # the 'sigma_cell' between `sigma_ind` and `sigma_ind+1`, which would
+                # affect the numerically calculated integrals for f(v_∥>0), so if we
+                # 'corrected' its value, those integrals would change and the constraints
+                # would not be exactly satisfied. The difference should be small, as the
+                # correction at that point is multiplied by
+                # v_∥^2/vth^2/(1+v_∥^2/vth^2)≈v_∥^2/vth^2≈0.
+                correction0_integral_pieces[ivpa] = 0.0
+            end
+            correction1_integral_pieces = @views @. vpa.scratch4 = correction0_integral_pieces * vpa_unnorm / vthe[end,ir]
+            correction2_integral_pieces = @views @. vpa.scratch5 = correction1_integral_pieces * vpa_unnorm / vthe[end,ir]
+            correction3_integral_pieces = @views @. vpa.scratch6 = correction2_integral_pieces * vpa_unnorm / vthe[end,ir]
+            correction4_integral_pieces = @views @. vpa.scratch7 = correction3_integral_pieces * vpa_unnorm / vthe[end,ir]
+            correction5_integral_pieces = @views @. vpa.scratch8 = correction4_integral_pieces * vpa_unnorm / vthe[end,ir]
+            correction6_integral_pieces = @views @. vpa.scratch9 = correction5_integral_pieces * vpa_unnorm / vthe[end,ir]
+
+            alpha = get_part3_for_one_moment_upper(correction0_integral_pieces)
+            beta = get_part3_for_one_moment_upper(correction1_integral_pieces)
+            gamma = get_part3_for_one_moment_upper(correction2_integral_pieces)
+            delta = get_part3_for_one_moment_upper(correction3_integral_pieces)
+            epsilon = get_part3_for_one_moment_upper(correction4_integral_pieces)
+            zeta = get_part3_for_one_moment_upper(correction5_integral_pieces)
+            eta = get_part3_for_one_moment_upper(correction6_integral_pieces)
+
+            # Update the v_∥>0 part of f to correct the moments as
+            # f(0<v_∥<vcut) = (1 + (A + B*v/vth + C*v^2/vth^2 + D*v^3/vth^3) * v^2/vth^2 / (1 + v^2/vth^2)) * fhat(0<v_∥<vcut)
+            # Constraints:
+            # ∫dv_∥ v_∥^n f(0<v_∥<vcut) = (-1)^n ∫dv_∥ v_∥^n f(-vcut<v_∥<0)    for n=0,1,2,3
+            # a2 = a3 + alpha*A + beta*B + gamma*C + delta*D
+            # -b2 = b3 + beta*A + gamma*B + delta*C + epsilon*D
+            # c2 = c3 + gamma*A + delta*B + epsilon*C + zeta*D
+            # -d2 = d3 + delta*A + epsilon*B + zeta*C + eta*D
+            solution = [alpha beta    gamma   delta   ;
+                        beta  gamma   delta   epsilon ;
+                        gamma delta   epsilon zeta    ;
+                        delta epsilon zeta    eta
+                       ] \ [a2-a3, -b2-b3, c2-c3, -d2-d3]
+            A, B, C, D = solution
+            for ivpa ∈ minus_vcut_ind-1:sigma_ind-1
+                v_over_vth = vpa_unnorm[ivpa]/vthe[end,ir]
+                pdf[ivpa,1,end,ir] = pdf[ivpa,1,end,ir] +
+                                   (A
+                                    + B * v_over_vth
+                                    + C * v_over_vth^2
+                                    + D * v_over_vth^3) *
+                                   v_over_vth^2 / (1.0 + v_over_vth^2) *
+                                   pdf[ivpa,1,end,ir]
+            end
         end
     end
 

From aa2654e7b1bfb2b2ed8732686cb6df1e01c09a82 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sun, 28 Jul 2024 17:16:34 +0100
Subject: [PATCH 401/428] Impose reversed f before constraint corrections in
 electron sheath bc

This allows integrals used in the constraint corrections to use a
consistent value of the distribution function at the grid point next to
v_parallel=0 that contributes to the numerical integral over the
opposite-sign v_parallel.
---
 .../src/electron_kinetic_equation.jl          | 118 ++++++++++--------
 1 file changed, 64 insertions(+), 54 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 36d87e88d..e5f5ad1fe 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -946,6 +946,27 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # sigma_ind where sigma is.
             sigma_fraction = (sigma - vpa_unnorm[sigma_ind-1]) / (vpa_unnorm[sigma_ind] - vpa_unnorm[sigma_ind-1])
 
+            # Want to construct the w-grid corresponding to -vpa.
+            #   wpa(vpa) = (vpa - upar)/vth
+            #   ⇒ vpa = vth*wpa(vpa) + upar
+            #   wpa(-vpa) = (-vpa - upar)/vth
+            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+            #             = (-vth*wpa - 2*upar)/vth
+            #             = -wpa - 2*upar/vth
+            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+            # Need to reverse vpa.grid because the grid passed as the second argument of
+            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
+            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
+            reverse!(reversed_wpa_of_minus_vpa)
+
+            # interpolate the pdf onto this grid
+            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+            reverse!(reversed_pdf)
+            pdf[sigma_ind:end,1,1,ir] .= reversed_pdf[sigma_ind:end]
+
             # Per-grid-point contributions to moment integrals
             # Note that we need to include the normalisation factor of 1/sqrt(pi) that
             # would be factored in by integrate_over_vspace(). This will need to
@@ -1029,36 +1050,20 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # constraints.
             @. pdf[:,1,1,ir] *= A + C * vpa_unnorm^2 / vthe[1,ir]^2
 
-            # Want to construct the w-grid corresponding to -vpa.
-            #   wpa(vpa) = (vpa - upar)/vth
-            #   ⇒ vpa = vth*wpa(vpa) + upar
-            #   wpa(-vpa) = (-vpa - upar)/vth
-            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
-            #             = (-vth*wpa - 2*upar)/vth
-            #             = -wpa - 2*upar/vth
-            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-            # Need to reverse vpa.grid because the grid passed as the second argument of
-            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
-            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
-            reverse!(reversed_wpa_of_minus_vpa)
-
-            # interpolate the pdf onto this grid
-            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
-            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,1,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-            reverse!(reversed_pdf)
-
-            ivpa_max = searchsortedlast(vpa_unnorm, vcut)
-            reversed_pdf[ivpa_max+1:end] .= 0.0
-            # vcut_fraction is the fraction of the distance between ivpa_max and
-            # ivpa_max+1 where vcut is.
-            vcut_fraction = (vcut - vpa_unnorm[ivpa_max]) / (vpa_unnorm[ivpa_max+1] - vpa_unnorm[ivpa_max])
-            reversed_pdf[ivpa_max] *= vcut_fraction
+            plus_vcut_ind = searchsortedlast(vpa_unnorm, vcut)
+            pdf[plus_vcut_ind+2:end,1,1,ir] .= 0.0
+            # vcut_fraction is the fraction of the distance between plus_vcut_ind and
+            # plus_vcut_ind+1 where vcut is.
+            vcut_fraction = (vcut - vpa_unnorm[plus_vcut_ind]) / (vpa_unnorm[plus_vcut_ind+1] - vpa_unnorm[plus_vcut_ind])
+            if vcut_fraction > 0.5
+                pdf[plus_vcut_ind+1,1,1,ir] *= vcut_fraction - 0.5
+            else
+                pdf[plus_vcut_ind+1,1,1,ir] = 0.0
+                pdf[plus_vcut_ind+1,1,1,ir] *= vcut_fraction + 0.5
+            end
 
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
             phi[1,ir] = me_over_mi * vcut^2
-            pdf[sigma_ind:end,1,1,ir] .= reversed_pdf[sigma_ind:end]
 
             moments.electron.constraints_A_coefficient[1,ir] = A
             moments.electron.constraints_B_coefficient[1,ir] = 0.0
@@ -1203,6 +1208,27 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # sigma_ind where sigma is.
             sigma_fraction = (sigma - vpa_unnorm[sigma_ind+1]) / (vpa_unnorm[sigma_ind] - vpa_unnorm[sigma_ind+1])
 
+            # Want to construct the w-grid corresponding to -vpa.
+            #   wpa(vpa) = (vpa - upar)/vth
+            #   ⇒ vpa = vth*wpa(vpa) + upar
+            #   wpa(-vpa) = (-vpa - upar)/vth
+            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
+            #             = (-vth*wpa - 2*upar)/vth
+            #             = -wpa - 2*upar/vth
+            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
+            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
+            # Need to reverse vpa.grid because the grid passed as the second argument of
+            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
+            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
+            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
+            reverse!(reversed_wpa_of_minus_vpa)
+
+            # interpolate the pdf onto this grid
+            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
+            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
+            reverse!(reversed_pdf)
+            pdf[1:sigma_ind,1,end,ir] .= reversed_pdf[1:sigma_ind]
+
             # Per-grid-point contributions to moment integrals
             # Note that we need to include the normalisation factor of 1/sqrt(pi) that
             # would be factored in by integrate_over_vspace(). This will need to
@@ -1284,36 +1310,20 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
             # constraints.
             @. pdf[:,1,end,ir] *= A + C * vpa_unnorm^2 / vthe[end,ir]^2
 
-            # Want to construct the w-grid corresponding to -vpa.
-            #   wpa(vpa) = (vpa - upar)/vth
-            #   ⇒ vpa = vth*wpa(vpa) + upar
-            #   wpa(-vpa) = (-vpa - upar)/vth
-            #             = (-(vth*wpa(vpa) + upar) - upar)/vth
-            #             = (-vth*wpa - 2*upar)/vth
-            #             = -wpa - 2*upar/vth
-            # [Note that `vpa.grid` is slightly mis-named here - it contains the values of
-            #  wpa(+vpa) as we are using a 'moment kinetic' approach.]
-            # Need to reverse vpa.grid because the grid passed as the second argument of
-            # interpolate_to_grid_1d!() needs to be sorted in increasing order.
-            reversed_wpa_of_minus_vpa = @. vpa.scratch3 = -vpa.grid + 2.0 * sigma
-            #reversed_wpa_of_minus_vpa = vpa.scratch3 .= .-vpa.grid
-            reverse!(reversed_wpa_of_minus_vpa)
-
-            # interpolate the pdf onto this grid
-            #@views interpolate_to_grid_1d!(interpolated_pdf, wpa_values, pdf[:,1,1,ir], vpa, vpa_spectral)
-            @views interpolate_to_grid_1d!(reversed_pdf, reversed_wpa_of_minus_vpa, pdf[:,1,end,ir], vpa, vpa_spectral) # Could make this more efficient by only interpolating to the points needed below, by taking an appropriate view of wpa_of_minus_vpa. Also, in the element containing vpa=0, this interpolation depends on the values that will be replaced by the reflected, interpolated values, which is not ideal (maybe this element should be treated specially first?).
-            reverse!(reversed_pdf)
-
-            ivpa_min = searchsortedfirst(vpa_unnorm, -vcut)
-            reversed_pdf[1:ivpa_min-1] .= 0.0
-            # vcut_fraction is the fraction of the distance between ivpa_min and
-            # ivpa_min-1 where -vcut is.
-            vcut_fraction = (-vcut - vpa_unnorm[ivpa_min]) / (vpa_unnorm[ivpa_min-1] - vpa_unnorm[ivpa_min])
-            reversed_pdf[ivpa_min] *= vcut_fraction
+            minus_vcut_ind = searchsortedfirst(vpa_unnorm, -vcut)
+            pdf[1:minus_vcut_ind-2,1,end,ir] .= 0.0
+            # vcut_fraction is the fraction of the distance between minus_vcut_ind and
+            # minus_vcut_ind-1 where -vcut is.
+            vcut_fraction = (-vcut - vpa_unnorm[minus_vcut_ind]) / (vpa_unnorm[minus_vcut_ind-1] - vpa_unnorm[minus_vcut_ind])
+            if vcut_fraction > 0.5
+                pdf[minus_vcut_ind-1,1,end,ir] *= vcut_fraction - 0.5
+            else
+                pdf[minus_vcut_ind-1,1,end,ir] = 0.0
+                pdf[minus_vcut_ind,1,end,ir] *= vcut_fraction + 0.5
+            end
 
             # update the electrostatic potential at the boundary to be the value corresponding to the updated cutoff velocity
             phi[end,ir] = me_over_mi * vcut^2
-            pdf[1:sigma_ind,1,end,ir] .= reversed_pdf[1:sigma_ind]
 
             moments.electron.constraints_A_coefficient[end,ir] = A
             moments.electron.constraints_B_coefficient[end,ir] = 0.0

From 4fd0701833f73d82aef5630fa0731a3cd050b0ef Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 29 Jul 2024 10:35:42 +0100
Subject: [PATCH 402/428] Use fixed, small tolerance for Newton iteration in
 electron sheath bc

Previously was setting the Newton iteration tolerance proportional to
the electron time-solver tolerance, but it is better to use a fixed,
small value (now 1.0e-11). The electron time-solver tolerance might be
set to a large value, in order to set the electron timestep only by CFL
limits, not by truncation error, and anyway the Newton iteration is used
to impose the moment constraints, which should be satisfied 'exactly'.
Even small errors, as they only affect a single grid point (at the
sheath entrance boundary) can cause significant errors in gradients that
are then used elsewhere.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 4 ++--
 moment_kinetics/src/initial_conditions.jl        | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index e5f5ad1fe..cd88b6bd0 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -341,7 +341,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             apply_electron_bc_and_constraints!(scratch[istage+1], phi, moments, z, vperp,
                                                vpa, vperp_spectral, vpa_spectral,
                                                vpa_advect, num_diss_params, composition,
-                                               t_params.rtol)
+                                               1.0e-11)
 
             latest_pdf = scratch[istage+1].pdf_electron
             
@@ -1502,7 +1502,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
     apply_electron_bc_and_constraints!(scratch[t_params.n_rk_stages+1], phi, moments, z,
                                        vperp, vpa, vperp_spectral, vpa_spectral,
                                        vpa_advect, num_diss_params, composition,
-                                       t_params.rtol)
+                                       1.0e-11)
     if evolve_ppar
         # Reset vth in the `moments` struct to the result consistent with full-accuracy RK
         # solution.
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 10336f40b..d729b9e40 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -336,7 +336,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
             [(speed=speed,)], moments, num_diss_params,
-            composition.me_over_mi, scratch_dummy; newton_tol=0.1*t_params.electron.rtol)
+            composition.me_over_mi, scratch_dummy; newton_tol=1.0e-11)
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux

From dec970c9580af375f825a271653f717389ff71e2 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 29 Jul 2024 10:44:18 +0100
Subject: [PATCH 403/428] Don't measure electron sheath bc Newton tolerance
 relative to u_over_vt

Simpler to just use a fixed absolute tolerance - the integrals are
normalised so this should be fine.

Tolerance now hard-coded to 1.0e-13.
---
 .../src/electron_kinetic_equation.jl          | 22 +++++++++----------
 moment_kinetics/src/initial_conditions.jl     |  6 ++---
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index cd88b6bd0..83ede627b 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -340,8 +340,7 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
             apply_electron_bc_and_constraints!(scratch[istage+1], phi, moments, z, vperp,
                                                vpa, vperp_spectral, vpa_spectral,
-                                               vpa_advect, num_diss_params, composition,
-                                               1.0e-11)
+                                               vpa_advect, num_diss_params, composition)
 
             latest_pdf = scratch[istage+1].pdf_electron
             
@@ -617,8 +616,7 @@ function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, mo
 
         apply_electron_bc_and_constraints!(new_scratch_electron, fields.phi, moments, z,
                                            vperp, vpa, vperp_spectral, vpa_spectral,
-                                           vpa_advect, num_diss_params, composition,
-                                           nl_solver_params.rtol)
+                                           vpa_advect, num_diss_params, composition)
 
         # Only the first entry in the `electron_pdf_substruct` will be used, so does not
         # matter what we put in the second and third except that they have the right type.
@@ -803,7 +801,7 @@ end
 
 function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp, vpa,
                                             vperp_spectral, vpa_spectral, vpa_advect,
-                                            num_diss_params, composition, rtol)
+                                            num_diss_params, composition)
     latest_pdf = this_scratch.pdf_electron
 
     begin_r_z_vperp_vpa_region()
@@ -817,7 +815,7 @@ function apply_electron_bc_and_constraints!(this_scratch, phi, moments, z, vperp
                                                 vperp_spectral, vpa_spectral, vpa_advect,
                                                 moments,
                                                 num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                composition.me_over_mi, 0.1 * rtol)
+                                                composition.me_over_mi)
 
     begin_r_z_region()
     A = moments.electron.constraints_A_coefficient
@@ -840,7 +838,10 @@ end
 function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vperp, vpa,
                                                      vperp_spectral, vpa_spectral,
                                                      vpa_adv, moments, vpa_diffusion,
-                                                     me_over_mi, newton_tol)
+                                                     me_over_mi)
+
+    newton_tol = 1.0e-13
+
     # Enforce velocity-space boundary conditions
     if vpa.n > 1
         begin_r_z_vperp_region()
@@ -1035,7 +1036,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
                 epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_lowerz(vcut, minus_vcut_ind)
 
-                if abs(epsilon) < newton_tol * abs(u_over_vt)
+                if abs(epsilon) < newton_tol
                     break
                 end
 
@@ -1295,7 +1296,7 @@ function enforce_boundary_condition_on_electron_pdf!(pdf, phi, vthe, upar, z, vp
 
                 epsilon, epsilonprime, A, C, a2, b2, c2, d2 = get_integrals_and_derivatives_upperz(vcut, plus_vcut_ind)
 
-                if abs(epsilon) < newton_tol * abs(u_over_vt)
+                if abs(epsilon) < newton_tol
                     break
                 end
 
@@ -1501,8 +1502,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
     end
     apply_electron_bc_and_constraints!(scratch[t_params.n_rk_stages+1], phi, moments, z,
                                        vperp, vpa, vperp_spectral, vpa_spectral,
-                                       vpa_advect, num_diss_params, composition,
-                                       1.0e-11)
+                                       vpa_advect, num_diss_params, composition)
     if evolve_ppar
         # Reset vth in the `moments` struct to the result consistent with full-accuracy RK
         # solution.
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index d729b9e40..af8e1f5f0 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -336,7 +336,7 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
             pdf.electron.norm, fields.phi, moments.electron.dens, moments.electron.upar,
             moments.electron.vth, z, vpa, vperp, vperp_spectral, vpa_spectral,
             [(speed=speed,)], moments, num_diss_params,
-            composition.me_over_mi, scratch_dummy; newton_tol=1.0e-11)
+            composition.me_over_mi, scratch_dummy)
     end
     # calculate the initial electron parallel heat flux;
     # if using kinetic electrons, this relies on the electron pdf, which itself relies on the electron heat flux
@@ -1538,7 +1538,7 @@ this 'initital' value for the electron will just be the first guess in an iterat
 """
 function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upar, vth, z,
         vpa, vperp, vperp_spectral, vpa_spectral, vpa_advect, moments, num_diss_params,
-        me_over_mi, scratch_dummy; restart_from_boltzmann=false, newton_tol)
+        me_over_mi, scratch_dummy; restart_from_boltzmann=false)
 
     if z.bc == "wall"
         begin_r_region()
@@ -1557,7 +1557,7 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
                                                     vperp_spectral, vpa_spectral,
                                                     vpa_advect, moments,
                                                     num_diss_params.electron.vpa_dissipation_coefficient > 0.0,
-                                                    me_over_mi, newton_tol)
+                                                    me_over_mi)
 
         # Distribute the z-boundary pdf values to every process
         begin_serial_region()

From 1596e81af8da7a74494c1a42a600ae48e034d248 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 29 Jul 2024 10:47:13 +0100
Subject: [PATCH 404/428] Diagnostic for number of electron steps per ion step

---
 .../makie_post_processing/src/makie_post_processing.jl      | 6 ++++++
 moment_kinetics/src/load_data.jl                            | 5 +++++
 2 files changed, 11 insertions(+)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index f9b39fc77..4ec288b08 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -7755,6 +7755,12 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                         plot_1d(time, linear_iterations, label=prefix * " " * p * " L per NL", ax=ax)
                     end
                 end
+
+                if ri.composition.electron_physics == kinetic_electrons
+                    has_nl_solver = true
+                    electron_steps_per_ion_step = get_variable(ri, "electron_steps_per_ion_step")
+                    plot_1d(time, electron_steps_per_ion_step, label=prefix * " electron steps per solve", ax=ax)
+                end
             end
 
             if has_nl_solver
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 0c7e35d5a..f7fb0171b 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -47,6 +47,7 @@ const timestep_diagnostic_variables = ("time_for_run", "step_counter", "dt",
                                        "average_successful_dt", "electron_step_counter",
                                        "electron_dt", "electron_failure_counter",
                                        "electron_failure_caused_by",
+                                       "electron_steps_per_ion_step",
                                        "electron_steps_per_output",
                                        "electron_failures_per_output",
                                        "electron_failure_caused_by_per_output",
@@ -4683,6 +4684,10 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             # Don't want a meaningless Inf...
             variable[1] = 0.0
         end
+    elseif variable_name == "electron_steps_per_ion_step"
+        electron_steps_per_output = get_variable(run_info, "electron_steps_per_output"; kwargs...)
+        ion_steps_per_output = get_variable(run_info, "steps_per_output"; kwargs...)
+        variable = electron_steps_per_output ./ ion_steps_per_output
     elseif variable_name == "electron_steps_per_output"
         variable = get_per_step_from_cumulative_variable(run_info, "electron_step_counter"; kwargs...)
     elseif variable_name == "electron_failures_per_output"

From 3b67b6dfd2ea69c21037f7d79e63ef6f29fcc979 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Mon, 29 Jul 2024 20:20:07 +0100
Subject: [PATCH 405/428] Avoid array copy when broadcasting in
 calculate_phi_from_Epar!()

The array allocation is not important, but the copy is not allowed when
`phi` is an `MPIDebugSharedArray`.
---
 moment_kinetics/src/em_fields.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index c6df3a405..ccc46b748 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -179,7 +179,7 @@ function calculate_phi_from_Epar!(phi, Epar, r, z)
         # Need to broadcast the lower-z boundary value, because we only communicate
         # delta_phi below, rather than passing the boundary values directly from block to
         # block.
-        phi[1,:] .= MPI.bcast(@view(phi[1,:]), z.comm; root=0)
+        MPI.Bcast!(@view(phi[1,:]), z.comm; root=0)
 
         if z.irank == z.nrank - 1
             # Don't want to change the upper-z boundary value, so save it here so we can
@@ -195,12 +195,12 @@ function calculate_phi_from_Epar!(phi, Epar, r, z)
 
         # Add contributions to integral along z from processes at smaller z-values than
         # this one.
-        this_delta_phi = phi[end,:] .- phi[1,:]
+        this_delta_phi = r.scratch2 .= phi[end,:] .- phi[1,:]
         for irank ∈ 0:z.nrank-2
-            delta_phi = MPI.bcast(this_delta_phi, z.comm; root=irank)
+            MPI.Bcast!(this_delta_phi, z.comm; root=irank)
             if z.irank > irank
                 @loop_r_z ir iz begin
-                    phi[iz,ir] += delta_phi[ir]
+                    phi[iz,ir] += this_delta_phi[ir]
                 end
             end
         end

From ec82cd9ed071118e66afa5375d7670268436184b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 30 Jul 2024 10:49:08 +0100
Subject: [PATCH 406/428] Limit 'sim time' for each electron solve, not number
 of iterations

When using adaptive timestepping, if the timestep becomes very small it
is very unlikely that the electron solve will converge in a fixed
(small-ish) number of iterations. This makes the electron solve
fragile. It should be more robust to fix the 'sim time' (i.e. normalised
physical time that the simulation is run for) for each electron solve,
which is what is implemented in this commit.
---
 .../src/electron_kinetic_equation.jl          | 27 ++++++++++++++-----
 moment_kinetics/src/initial_conditions.jl     | 15 +++++++----
 moment_kinetics/src/time_advance.jl           | 15 ++++++++---
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 83ede627b..6f492f97e 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -68,8 +68,9 @@ OUTPUT:
 function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_spectral,
         vperp_spectral, vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
         collisions, composition, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
-        residual_tolerance=nothing, evolve_ppar=false, ion_dt=nothing)
+        max_electron_pdf_iterations, max_electron_sim_time; io_electron=nothing,
+        initial_time=nothing, residual_tolerance=nothing, evolve_ppar=false,
+        ion_dt=nothing)
 
     # set the method to use to solve the electron kinetic equation
     solution_method = "artificial_time_derivative"
@@ -80,8 +81,8 @@ function update_electron_pdf!(scratch, pdf, moments, phi, r, z, vperp, vpa, z_sp
         return update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi,
             collisions, composition, r, z, vperp, vpa, z_spectral, vperp_spectral,
             vpa_spectral, z_advect, vpa_advect, scratch_dummy, t_params,
-            external_source_settings, num_diss_params, max_electron_pdf_iterations;
-            io_electron=io_electron, initial_time=initial_time,
+            external_source_settings, num_diss_params, max_electron_pdf_iterations,
+            max_electron_sim_time; io_electron=io_electron, initial_time=initial_time,
             residual_tolerance=residual_tolerance, evolve_ppar=evolve_ppar, ion_dt=ion_dt)
     elseif solution_method == "shooting_method"
         dens = moments.electron.dens
@@ -148,8 +149,17 @@ OUTPUT:
 function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, collisions,
         composition, r, z, vperp, vpa, z_spectral, vperp_spectral, vpa_spectral, z_advect,
         vpa_advect, scratch_dummy, t_params, external_source_settings, num_diss_params,
-        max_electron_pdf_iterations; io_electron=nothing, initial_time=nothing,
-        residual_tolerance=nothing, evolve_ppar=false, ion_dt=nothing)
+        max_electron_pdf_iterations, max_electron_sim_time; io_electron=nothing,
+        initial_time=nothing, residual_tolerance=nothing, evolve_ppar=false,
+        ion_dt=nothing)
+
+    if max_electron_pdf_iterations !== nothing && max_electron_sim_time !== nothing
+        error("Cannot use both max_electron_pdf_iterations=$max_electron_pdf_iterations "
+              * "and max_electron_sim_time=$max_electron_sim_time at the same time")
+    end
+    if max_electron_pdf_iterations === nothing && max_electron_sim_time === nothing
+        error("Must set one of max_electron_pdf_iterations and max_electron_sim_time")
+    end
 
     begin_r_z_region()
 
@@ -207,6 +217,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
         # t_params are set relative to 0.0).
         moments_output_times = t_params.moments_output_times .+ initial_time
         dfns_output_times = t_params.dfns_output_times .+ initial_time
+    else
+        initial_time = t_params.t[]
     end
     if io_electron === nothing && t_params.debug_io !== nothing
         # Overwrite the debug output file with the output from this call to
@@ -279,7 +291,8 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
     end
     # evolve (artificially) in time until the residual is less than the tolerance
     while (!electron_pdf_converged
-           && (t_params.step_counter[] - initial_step_counter < max_electron_pdf_iterations)
+           && (max_electron_pdf_iterations === nothing || t_params.step_counter[] - initial_step_counter < max_electron_pdf_iterations)
+           && (max_electron_sim_time === nothing || t_params.t[] - initial_time < max_electron_sim_time)
            && t_params.dt[] > 0.0 && !isnan(t_params.dt[]))
 
         # Set the initial values for the next step to the final values from the previous
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index af8e1f5f0..c9809732e 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -659,9 +659,12 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
         # First run with evolve_ppar=true to get electron_ppar close to steady state.
         # electron_ppar does not have to be exactly steady state as it will be
         # time-evolved along with the ions.
-        max_electron_pdf_iterations = 2000000
-        #max_electron_pdf_iterations = 500000
-        #max_electron_pdf_iterations = 10000
+        #max_electron_pdf_iterations = 2000000
+        ##max_electron_pdf_iterations = 500000
+        ##max_electron_pdf_iterations = 10000
+        #max_electron_sim_time = nothing
+        max_electron_pdf_iterations = nothing
+        max_electron_sim_time = 2.0
         if t_params.electron.debug_io !== nothing
             io_electron = setup_electron_io(t_params.electron.debug_io[1], vpa, vperp, z,
                                             r, composition, collisions,
@@ -705,7 +708,8 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                             vpa_advect, scratch_dummy, t_params.electron,
                                             collisions, composition,
                                             external_source_settings, num_diss_params,
-                                            max_electron_pdf_iterations;
+                                            max_electron_pdf_iterations,
+                                            max_electron_sim_time;
                                             io_electron=io_initial_electron,
                                             initial_time=code_time,
                                             residual_tolerance=t_input["initialization_residual_value"],
@@ -765,7 +769,8 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
                                          vpa_advect, scratch_dummy, t_params.electron,
                                          collisions, composition,
                                          external_source_settings, num_diss_params,
-                                         max_electron_pdf_iterations;
+                                         max_electron_pdf_iterations,
+                                         max_electron_sim_time;
                                          io_electron=io_initial_electron)
             end
             if success != ""
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index ca5529a55..cf0db6a19 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2396,7 +2396,10 @@ function apply_all_bcs_constraints_update_moments!(
                                            num_diss_params.electron.moment_dissipation_coefficient, 
                                            composition.electron_physics)
     if composition.electron_physics == kinetic_electrons
-        max_electron_pdf_iterations = 1000
+        #max_electron_pdf_iterations = 1000
+        #max_electron_sim_time = nothing
+        max_electron_pdf_iterations = nothing
+        max_electron_sim_time = 1.0e-3
 
         # Copy ion and electron moments from `scratch` into `moments` to be used in
         # electron kinetic equation update
@@ -2433,7 +2436,7 @@ function apply_all_bcs_constraints_update_moments!(
                z_spectral, vperp_spectral, vpa_spectral, electron_z_advect,
                electron_vpa_advect, scratch_dummy, t_params.electron, collisions,
                composition, external_source_settings, num_diss_params,
-               max_electron_pdf_iterations)
+               max_electron_pdf_iterations, max_electron_sim_time)
             success = kinetic_electron_success
         end
     end
@@ -3513,7 +3516,10 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                              gyroavs, scratch_dummy, dt,
                                              nl_solver_params.electron_advance)
     elseif t_params.implicit_electron_ppar
-        max_electron_pdf_iterations = 1000
+        #max_electron_pdf_iterations = 1000
+        #max_electron_sim_time = nothing
+        max_electron_pdf_iterations = nothing
+        max_electron_sim_time = 1.0e-3
         electron_success = update_electron_pdf!(scratch_electron, pdf.electron.norm,
                                                 moments, fields.phi, r, z, vperp, vpa,
                                                 z_spectral, vperp_spectral, vpa_spectral,
@@ -3521,7 +3527,8 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                                 scratch_dummy, t_params.electron,
                                                 collisions, composition,
                                                 external_source_settings, num_diss_params,
-                                                max_electron_pdf_iterations; ion_dt=dt)
+                                                max_electron_pdf_iterations,
+                                                max_electron_sim_time; ion_dt=dt)
         success = (electron_success == "")
     elseif advance.electron_conduction
         success = implicit_braginskii_conduction!(fvec_out, fvec_in, moments, z, r, dt,

From 18b42f199662226e55ee850b5d5ec472300be536 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Tue, 30 Jul 2024 10:49:08 +0100
Subject: [PATCH 407/428] Option to use electron temperature equation in
 kinetic electron sims

Using a temperature equation sometimes seems to be more numerically
stable.  It seems to help keep T_e (and q_e) smoother, which is useful
as v_th=sqrt(2*T_e/m_e) is used to scale the grid for the distribution
function.

Currently a slightly hacky implementation which converts electron_ppar
to electron temperature, does the forward-Euler update, then converts
back.
---
 .../src/makie_post_processing.jl              |  36 ++-
 .../src/electron_fluid_equations.jl           | 234 ++++++++++++------
 .../src/electron_kinetic_equation.jl          |  10 +-
 moment_kinetics/src/em_fields.jl              |   6 +-
 moment_kinetics/src/external_sources.jl       |   3 +-
 moment_kinetics/src/file_io.jl                |   5 +-
 moment_kinetics/src/initial_conditions.jl     |  20 +-
 moment_kinetics/src/input_structs.jl          |   2 +
 moment_kinetics/src/load_data.jl              |  10 +-
 moment_kinetics/src/moment_kinetics.jl        |   3 +-
 moment_kinetics/src/runge_kutta.jl            |   4 +-
 moment_kinetics/src/time_advance.jl           |  78 +++---
 moment_kinetics/src/velocity_moments.jl       |   8 +-
 13 files changed, 278 insertions(+), 141 deletions(-)

diff --git a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
index 4ec288b08..fb5aeefa0 100644
--- a/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
+++ b/makie_post_processing/makie_post_processing/src/makie_post_processing.jl
@@ -210,11 +210,13 @@ function makie_post_process(run_dir::Union{String,Tuple},
     # Only plot electron stuff if some runs have electrons
     if any(ri !== nothing for ri ∈ run_info_moments)
         has_electrons = any(r.composition.electron_physics
-                            ∈ (braginskii_fluid, kinetic_electrons)
+                            ∈ (braginskii_fluid, kinetic_electrons,
+                               kinetic_electrons_with_temperature_equation)
                             for r in run_info_moments)
     else
         has_electrons = any(r.composition.electron_physics
-                            ∈ (braginskii_fluid, kinetic_electrons)
+                            ∈ (braginskii_fluid, kinetic_electrons,
+                               kinetic_electrons_with_temperature_equation)
                             for r in run_info_dfns)
     end
 
@@ -267,8 +269,8 @@ function makie_post_process(run_dir::Union{String,Tuple},
     end
 
     timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix)
-    if any((ri.composition.electron_physics ==
-                moment_kinetics.input_structs.kinetic_electrons
+    if any((ri.composition.electron_physics ∈ (kinetic_electrons,
+                                               kinetic_electrons_with_temperature_equation)
             && !ri.t_input["implicit_electron_advance"]) for ri ∈ run_info)
         timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix, electron=true)
     end
@@ -5198,7 +5200,9 @@ function constraints_plots(run_info; plot_prefix=plot_prefix)
             end
 
             # Electrons
-            #if any(ri.composition.electron_physics == kinetic_electrons for ri ∈ run_info)
+            #if any(ri.composition.electron_physics ∈ (kinetic_electrons,
+            #                                          kinetic_electrons_with_temperature_equation)
+            #       for ri ∈ run_info)
 
             #    fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
             #    for ri ∈ run_info
@@ -5360,7 +5364,9 @@ function constraints_plots(run_info; plot_prefix=plot_prefix)
             end
 
             # Electrons
-            #if any(ri.composition.electron_physics == kinetic_electrons for ri ∈ run_info)
+            #if any(ri.composition.electron_physics ∈ (kinetic_electrons,
+            #                                          kinetic_electrons_with_temperature_equation)
+            #       for ri ∈ run_info)
 
             #    frame_index = Observable(1)
             #    fig, ax = get_1d_ax(; xlabel="z", ylabel="constraint coefficient")
@@ -7436,7 +7442,8 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                             ax=ax_failures)
                 end
                 if electron || ri.composition.electron_physics ∈ (braginskii_fluid,
-                                                                  kinetic_electrons)
+                                                                  kinetic_electrons,
+                                                                  kinetic_electrons_with_temperature_equation)
                     # Electron parallel pressure failure counter
                     counter += 1
                     plot_1d(time, @view failure_caused_by_per_output[counter,:];
@@ -7477,7 +7484,8 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                                 linestyle=:dot,
                                 label=prefix * "nonlinear iteration convergence failure", ax=ax_failures)
                     end
-                    if ri.composition.electron_physics == kinetic_electrons
+                    if ri.composition.electron_physics ∈ (kinetic_electrons,
+                                                          kinetic_electrons_with_temperature_equation)
                         # Kinetic electron iteration failed to converge
                         counter += 1
                         plot_1d(time, @view failure_caused_by_per_output[counter,:];
@@ -7646,7 +7654,9 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                             label=prefix * "ion ppar RK accuracy", ax=ax,
                             linestyle=:dash)
                 end
-                if electron || ri.composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+                if electron || ri.composition.electron_physics ∈ (braginskii_fluid,
+                                                                  kinetic_electrons,
+                                                                  kinetic_electrons_with_temperature_equation)
                     counter += 1
                     plot_1d(time, @view limit_caused_by_per_output[counter,:];
                             label=prefix * "electron ppar RK accuracy", ax=ax,
@@ -7756,7 +7766,8 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                     end
                 end
 
-                if ri.composition.electron_physics == kinetic_electrons
+                if ri.composition.electron_physics ∈ (kinetic_electrons,
+                                                      kinetic_electrons_with_temperature_equation)
                     has_nl_solver = true
                     electron_steps_per_ion_step = get_variable(ri, "electron_steps_per_ion_step")
                     plot_1d(time, electron_steps_per_ion_step, label=prefix * " electron steps per solve", ax=ax)
@@ -7821,8 +7832,9 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
                                                 :leftspinevisible=>false,
                                                 :rightspinevisible=>false))
             end
-            if electron || any(ri.composition.electron_physics == kinetic_electrons for ri
-                               ∈ run_info)
+            if electron || any(ri.composition.electron_physics ∈ (kinetic_electrons,
+                                                                  kinetic_electrons_with_temperature_equation)
+                               for ri ∈ run_info)
                 data = get_variable(run_info, "CFL_electron_z")
                 datamin = minimum(minimum(d) for d ∈ data)
                 animate_vs_vpa_z(run_info, "CFL_electron_z"; data=data, it=it,
diff --git a/moment_kinetics/src/electron_fluid_equations.jl b/moment_kinetics/src/electron_fluid_equations.jl
index 55d14655b..395a94f4c 100644
--- a/moment_kinetics/src/electron_fluid_equations.jl
+++ b/moment_kinetics/src/electron_fluid_equations.jl
@@ -12,8 +12,7 @@ export update_electron_vth_temperature!
 using ..communication
 using ..derivatives: derivative_z!
 using ..looping
-using ..input_structs: boltzmann_electron_response_with_simple_sheath
-using ..input_structs: braginskii_fluid, kinetic_electrons
+using ..input_structs
 using ..moment_kinetics_structs: electron_pdf_substruct, moments_electron_substruct
 using ..nonlinear_solvers
 using ..type_definitions: mk_float
@@ -75,7 +74,9 @@ function calculate_electron_upar_from_charge_conservation!(upar_e, updated, dens
         end
         # if using a simple logical sheath model, then the electron parallel current at the boundaries in zed
         # is equal and opposite to the ion parallel current
-        if electron_model ∈ (boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons)
+        if electron_model ∈ (boltzmann_electron_response_with_simple_sheath,
+                             braginskii_fluid, kinetic_electrons,
+                             kinetic_electrons_with_temperature_equation)
             boundary_flux = r.scratch_shared
             boundary_ion_flux = r.scratch_shared2
             if z.irank == 0
@@ -123,7 +124,8 @@ function calculate_electron_moments!(scratch, pdf, moments, composition, collisi
     calculate_electron_upar_from_charge_conservation!(
         scratch.electron_upar, moments.electron.upar_updated, scratch.electron_density,
         scratch.upar, scratch.density, composition.electron_physics, r, z)
-    if composition.electron_physics ∉ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∉ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         begin_r_z_region()
         @loop_r_z ir iz begin
             scratch.electron_ppar[iz,ir] = 0.5 * composition.me_over_mi *
@@ -147,85 +149,168 @@ function calculate_electron_moments!(scratch, pdf, moments, composition, collisi
 end
 
 """
-use the electron energy equation to evolve the electron temperature via
-an explicit time advance.
+use the electron energy or temperature equation to evolve the electron temperature via an
+explicit time advance.
 NB: so far, this is only set up for 1D problem, where we can assume
 an isotropic distribution in f_e so that p_e = n_e T_e = ppar_e
 """
 function electron_energy_equation!(ppar_out, ppar_in, electron_density, electron_upar,
-                                   ion_upar, ion_ppar, density_neutral, uz_neutral,
-                                   pz_neutral, moments, collisions, dt, composition,
-                                   electron_source_settings, num_diss_params, z;
-                                   conduction=true)
-    begin_r_z_region()
-    # define some abbreviated variables for convenient use in rest of function
-    me_over_mi = composition.me_over_mi
-    nu_ei = collisions.nu_ei
-    # calculate contribution to rhs of energy equation (formulated in terms of pressure)
-    # arising from derivatives of ppar, qpar and upar
-    @loop_r_z ir iz begin
-        ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
-                               + 3*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
-    end
-    if conduction
+                                   ion_density, ion_upar, ion_ppar, density_neutral,
+                                   uz_neutral, pz_neutral, moments, collisions, dt,
+                                   composition, electron_source_settings, num_diss_params,
+                                   z; conduction=true)
+    if composition.electron_physics == kinetic_electrons_with_temperature_equation
+        # Hacky way to implement temperature equation:
+        #  - convert ppar to T by dividing by density
+        #  - update T with a forward-Euler step using the temperature equation
+        #  - multiply by density to get back to ppar (should this be new density rather than
+        #    old density? For initial testing, only looking at the electron initialisation
+        #    where density is not updated, this does not matter).
+
+        begin_r_z_region()
+        # define some abbreviated variables for convenient use in rest of function
+        me_over_mi = composition.me_over_mi
+        nu_ei = collisions.nu_ei
+        T_in = moments.temp
+        # calculate contribution to rhs of energy equation (formulated in terms of pressure)
+        # arising from derivatives of ppar, qpar and upar
         @loop_r_z ir iz begin
-            ppar_out[iz,ir] -= dt*moments.dqpar_dz[iz,ir]
+            # Convert ppar_out to temperature for most of this function
+            ppar_out[iz,ir] *= 2.0 / electron_density[iz,ir]
+            ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dT_dz[iz,ir]
+                                   + 2.0*T_in[iz,ir]*moments.dupar_dz[iz,ir])
         end
-    end
-    # @loop_r_z ir iz begin
-    #     ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
-    #                 + (2/3)*moments.dqpar_dz[iz,ir]
-    #                 + (5/3)*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
-    # end
-    # compute the contribution to the rhs of the energy equation
-    # arising from artificial diffusion
-    diffusion_coefficient = num_diss_params.electron.moment_dissipation_coefficient
-    if diffusion_coefficient > 0.0
-        @loop_r_z ir iz begin
-            ppar_out[iz,ir] += dt*diffusion_coefficient*moments.d2ppar_dz2[iz,ir]
+        if conduction
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] -= 2.0 * dt*moments.dqpar_dz[iz,ir] / electron_density[iz,ir]
+            end
         end
-    end
-    # compute the contribution to the rhs of the energy equation
-    # arising from electron-ion collisions
-    if nu_ei > 0.0
-        @loop_s_r_z is ir iz begin
-            ppar_out[iz,ir] += dt * (2 * me_over_mi * nu_ei * (ion_ppar[iz,ir,is] - ppar_in[iz,ir]))
-            ppar_out[iz,ir] += dt * ((2/3) * moments.parallel_friction[iz,ir]
-                                     * (ion_upar[iz,ir,is]-electron_upar[iz,ir]))
+        # compute the contribution to the rhs of the energy equation
+        # arising from artificial diffusion
+        diffusion_coefficient = num_diss_params.electron.moment_dissipation_coefficient
+        if diffusion_coefficient > 0.0
+            error("diffusion not implemented for electron temperature equation yet")
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] += dt*diffusion_coefficient*moments.d2T_dz2[iz,ir]
+            end
         end
-    end
-    # add in contributions due to charge exchange/ionization collisions
-    if composition.n_neutral_species > 0
-        if abs(collisions.charge_exchange_electron) > 0.0
-            @loop_sn_r_z isn ir iz begin
-                ppar_out[iz,ir] +=
-                    dt * me_over_mi * collisions.charge_exchange_electron * (
-                    2*(electron_density[iz,ir]*pz_neutral[iz,ir,isn] -
-                    density_neutral[iz,ir,isn]*ppar_in[iz,ir]) +
-                    (2/3)*electron_density[iz,ir]*density_neutral[iz,ir,isn] *
-                    (uz_neutral[iz,ir,isn] - electron_upar[iz,ir])^2)
+        # compute the contribution to the rhs of the energy equation
+        # arising from electron-ion collisions
+        if nu_ei > 0.0
+            @loop_s_r_z is ir iz begin
+                ppar_out[iz,ir] += dt * 2.0 * (2 * me_over_mi * nu_ei * (2.0*ion_ppar[iz,ir,is]/ion_density[iz,ir,is] - T_in[iz,ir]))
+                ppar_out[iz,ir] += dt * 2.0 * ((2/3) * moments.parallel_friction[iz,ir]
+                                               * (ion_upar[iz,ir,is]-electron_upar[iz,ir])) / electron_density[iz,ir]
             end
         end
-        if abs(collisions.ionization_electron) > 0.0
-            # @loop_s_r_z is ir iz begin
-            #     ppar_out[iz,ir] +=
-            #         dt * collisions.ionization_electron * density_neutral[iz,ir,is] * (
-            #         ppar_in[iz,ir] -
-            #         (2/3)*electron_density[iz,ir] * collisions.ionization_energy)
-            # end
-            @loop_sn_r_z isn ir iz begin
-                ppar_out[iz,ir] +=
-                    dt * collisions.ionization_electron * density_neutral[iz,ir,isn] * (
-                    ppar_in[iz,ir] -
-                    electron_density[iz,ir] * collisions.ionization_energy)
+        # add in contributions due to charge exchange/ionization collisions
+        if composition.n_neutral_species > 0
+            if abs(collisions.charge_exchange_electron) > 0.0
+                @loop_sn_r_z isn ir iz begin
+                    ppar_out[iz,ir] +=
+                        dt * 2.0 * me_over_mi * collisions.charge_exchange_electron * (
+                            2*(pz_neutral[iz,ir,isn] -
+                               density_neutral[iz,ir,isn]*ppar_in[iz,ir]/electron_density[iz,ir]) +
+                            (2/3)*density_neutral[iz,ir,isn] *
+                            (uz_neutral[iz,ir,isn] - electron_upar[iz,ir])^2)
+                end
+            end
+            if abs(collisions.ionization_electron) > 0.0
+                @loop_sn_r_z isn ir iz begin
+                    ppar_out[iz,ir] +=
+                        dt * 2.0 * collisions.ionization_electron * density_neutral[iz,ir,isn] * (
+                            ppar_in[iz,ir] / electron_density[iz,ir]  -
+                            collisions.ionization_energy)
+                end
             end
         end
-    end
 
-    if electron_source_settings.active
-        source_amplitude = moments.external_source_pressure_amplitude
+        if electron_source_settings.active
+            pressure_source_amplitude = moments.external_source_pressure_amplitude
+            density_source_amplitude = moments.external_source_density_amplitude
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] += dt * (2.0 * pressure_source_amplitude[iz,ir]
+                                         - T_in[iz,ir] * density_source_amplitude[iz,ir]) /
+                                        electron_density[iz,ir]
+            end
+        end
+
+        # Now that forward-Euler step for temperature is finished, convert ppar_out back to
+        # pressure.
+        @loop_r_z ir iz begin
+            ppar_out[iz,ir] *= 0.5 * electron_density[iz,ir]
+        end
+    else
+        begin_r_z_region()
+        # define some abbreviated variables for convenient use in rest of function
+        me_over_mi = composition.me_over_mi
+        nu_ei = collisions.nu_ei
+        # calculate contribution to rhs of energy equation (formulated in terms of pressure)
+        # arising from derivatives of ppar, qpar and upar
         @loop_r_z ir iz begin
-            ppar_out[iz,ir] += dt * source_amplitude[iz,ir]
+            ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
+                                   + 3*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
+        end
+        if conduction
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] -= dt*moments.dqpar_dz[iz,ir]
+            end
+        end
+        # @loop_r_z ir iz begin
+        #     ppar_out[iz,ir] -= dt*(electron_upar[iz,ir]*moments.dppar_dz[iz,ir]
+        #                 + (2/3)*moments.dqpar_dz[iz,ir]
+        #                 + (5/3)*ppar_in[iz,ir]*moments.dupar_dz[iz,ir])
+        # end
+        # compute the contribution to the rhs of the energy equation
+        # arising from artificial diffusion
+        diffusion_coefficient = num_diss_params.electron.moment_dissipation_coefficient
+        if diffusion_coefficient > 0.0
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] += dt*diffusion_coefficient*moments.d2ppar_dz2[iz,ir]
+            end
+        end
+        # compute the contribution to the rhs of the energy equation
+        # arising from electron-ion collisions
+        if nu_ei > 0.0
+            @loop_s_r_z is ir iz begin
+                ppar_out[iz,ir] += dt * (2 * me_over_mi * nu_ei * (ion_ppar[iz,ir,is] - ppar_in[iz,ir]))
+                ppar_out[iz,ir] += dt * ((2/3) * moments.parallel_friction[iz,ir]
+                                         * (ion_upar[iz,ir,is]-electron_upar[iz,ir]))
+            end
+        end
+        # add in contributions due to charge exchange/ionization collisions
+        if composition.n_neutral_species > 0
+            if abs(collisions.charge_exchange_electron) > 0.0
+                @loop_sn_r_z isn ir iz begin
+                    ppar_out[iz,ir] +=
+                        dt * me_over_mi * collisions.charge_exchange_electron * (
+                        2*(electron_density[iz,ir]*pz_neutral[iz,ir,isn] -
+                        density_neutral[iz,ir,isn]*ppar_in[iz,ir]) +
+                        (2/3)*electron_density[iz,ir]*density_neutral[iz,ir,isn] *
+                        (uz_neutral[iz,ir,isn] - electron_upar[iz,ir])^2)
+                end
+            end
+            if abs(collisions.ionization_electron) > 0.0
+                # @loop_s_r_z is ir iz begin
+                #     ppar_out[iz,ir] +=
+                #         dt * collisions.ionization_electron * density_neutral[iz,ir,is] * (
+                #         ppar_in[iz,ir] -
+                #         (2/3)*electron_density[iz,ir] * collisions.ionization_energy)
+                # end
+                @loop_sn_r_z isn ir iz begin
+                    ppar_out[iz,ir] +=
+                        dt * collisions.ionization_electron * density_neutral[iz,ir,isn] * (
+                        ppar_in[iz,ir] -
+                        electron_density[iz,ir] * collisions.ionization_energy)
+                end
+            end
+        end
+
+        if electron_source_settings.active
+            source_amplitude = moments.external_source_pressure_amplitude
+            @loop_r_z ir iz begin
+                ppar_out[iz,ir] += dt * source_amplitude[iz,ir]
+            end
         end
     end
 
@@ -254,10 +339,10 @@ function electron_energy_residual!(residual, electron_ppar_out, fvec_in, moments
         residual[iz,ir] = electron_ppar_in[iz,ir]
     end
     electron_energy_equation!(residual, electron_ppar_out,
-                              fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
-                              fvec_in.ppar, fvec_in.density_neutral,
-                              fvec_in.uz_neutral, fvec_in.pz_neutral,
-                              moments.electron, collisions, dt, composition,
+                              fvec_in.density, fvec_in.electron_upar, fvec_in.density,
+                              fvec_in.upar, fvec_in.ppar, fvec_in.density_neutral,
+                              fvec_in.uz_neutral, fvec_in.pz_neutral, moments.electron,
+                              collisions, dt, composition,
                               external_source_settings.electron, num_diss_params, z)
     # Now
     #   residual = f_in + dt*RHS(f_out)
@@ -473,7 +558,8 @@ function calculate_electron_qpar!(electron_moments, pdf, ppar_e, upar_e, upar_i,
                     qpar_e[iz,ir] -= (1/2) * 3.16 * ppar_e[iz,ir] / (me_over_mi * nu_ei) * dTe_dz[iz,ir] 
                 end
             end
-        elseif electron_model == kinetic_electrons
+        elseif electron_model ∈ (kinetic_electrons,
+                                 kinetic_electrons_with_temperature_equation)
             # use the modified electron pdf to calculate the electron heat flux
             if isa(pdf, electron_pdf_substruct)
                 electron_pdf = pdf.norm
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 6f492f97e..dd828ea3a 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -608,10 +608,10 @@ function implicit_electron_advance!(fvec_out, fvec_in, pdf, scratch_electron, mo
                                            num_diss_params.electron.moment_dissipation_coefficient,
                                            composition.electron_physics)
     electron_energy_equation!(electron_ppar_out, fvec_in.electron_ppar,
-                              fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
-                              fvec_in.ppar, fvec_in.density_neutral,
-                              fvec_in.uz_neutral, fvec_in.pz_neutral,
-                              moments.electron, collisions, dt, composition,
+                              fvec_in.density, fvec_in.electron_upar, fvec_in.density,
+                              fvec_in.upar, fvec_in.ppar, fvec_in.density_neutral,
+                              fvec_in.uz_neutral, fvec_in.pz_neutral, moments.electron,
+                              collisions, dt, composition,
                               external_source_settings.electron, num_diss_params, z)
 
     function residual_func!(residual, new_variables)
@@ -1814,7 +1814,7 @@ function electron_kinetic_equation_euler_update!(fvec_out, fvec_in, moments, z,
     if evolve_ppar
         electron_energy_equation!(fvec_out.electron_ppar, fvec_in.electron_ppar,
                                   moments.electron.dens, moments.electron.upar,
-                                  moments.ion.upar, moments.ion.ppar,
+                                  moments.ion.dens, moments.ion.upar, moments.ion.ppar,
                                   moments.neutral.dens, moments.neutral.uz,
                                   moments.neutral.pz, moments.electron, collisions, dt,
                                   composition, external_source_settings.electron,
diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index ccc46b748..6c20ca9cc 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -112,7 +112,8 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments
         @loop_r_z ir iz begin
             fields.phi[iz,ir] = composition.T_e * log(dens_e[iz,ir] / N_e[ir])
         end
-    elseif composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    elseif composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
         calculate_Epar_from_electron_force_balance!(fields.Ez, dens_e, moments.electron.dppar_dz,
             collisions.nu_ei, moments.electron.parallel_friction,
             composition.n_neutral_species, collisions.charge_exchange_electron, composition.me_over_mi,
@@ -142,7 +143,8 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments
         end
     end
     # if advancing electron fluid equations, solve for Ez directly from force balance
-    if composition.electron_physics ∉ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∉ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         if z.n > 1
             # Ez = - d phi / dz
             @views derivative_z!(fields.Ez,-fields.phi,
diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index be2206f94..98fe4169c 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -227,7 +227,8 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
     end
 
     ion_settings = get_settings(false)
-    if electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                           kinetic_electrons_with_temperature_equation)
         electron_settings = get_electron_settings(ion_settings)
     else
         electron_settings = (active=false,)
diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 2c339611c..2515be182 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1581,7 +1581,7 @@ function define_dynamic_electron_moment_variables!(fid, r::coordinate, z::coordi
                                parallel_io=parallel_io,
                                description="'C' coefficient enforcing pressure constraint for electrons")
 
-    if electron_physics == kinetic_electrons
+    if electron_physics ∈ (kinetic_electrons, kinetic_electrons_with_temperature_equation)
         io_electron_step_counter = create_dynamic_variable!(
             dynamic, "electron_step_counter", mk_int; parallel_io=parallel_io,
             description="cumulative number of electron pseudo-timesteps for the run")
@@ -1885,7 +1885,8 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
             io_f_start_last_timestep = nothing
         end
 
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics ∈ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             # io_f_electron is the handle for the electron pdf
             io_f_electron = create_dynamic_variable!(dynamic, "f_electron", mk_float, vpa,
                                                      vperp, z, r;
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index c9809732e..acfd10e35 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -41,7 +41,7 @@ using ..electron_fluid_equations: calculate_electron_upar_from_charge_conservati
 using ..electron_fluid_equations: calculate_electron_qpar!, electron_fluid_qpar_boundary_condition!
 using ..electron_fluid_equations: calculate_electron_parallel_friction_force!
 using ..electron_kinetic_equation: update_electron_pdf!, enforce_boundary_condition_on_electron_pdf!
-using ..input_structs: boltzmann_electron_response_with_simple_sheath, braginskii_fluid, kinetic_electrons
+using ..input_structs
 using ..derivatives: derivative_z!
 using ..utils: get_default_restart_filename, get_prefix_iblock_and_move_existing_file,
                get_backup_filename
@@ -106,7 +106,8 @@ function create_pdf(composition, r, z, vperp, vpa, vzeta, vr, vz)
     pdf_neutral_norm = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_neutral_species)
     # buffer array is for neutral-ion collisions, not for storing neutral pdf
     pdf_neutral_buffer = allocate_shared_float(vz.n, vr.n, vzeta.n, z.n, r.n, composition.n_ion_species)
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         pdf_electron_norm = allocate_shared_float(vpa.n, vperp.n, z.n, r.n)
         # MB: not sure if pdf_electron_buffer will ever be needed, but create for now
         # to emulate ion and neutral behaviour
@@ -218,7 +219,8 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
             moments.electron.constraints_A_coefficient .= 1.0
             moments.electron.constraints_B_coefficient .= 0.0
             moments.electron.constraints_C_coefficient .= 0.0
-            if composition.electron_physics == kinetic_electrons
+            if composition.electron_physics ∈ (kinetic_electrons,
+                                               kinetic_electrons_with_temperature_equation)
                 pdf.electron.norm .= 0.0
             end
         end
@@ -289,7 +291,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         @loop_r_z ir iz begin
             moments.electron.ppar[iz,ir] = 0.5 * moments.electron.dens[iz,ir] * moments.electron.temp[iz,ir]
         end
-    elseif restart_electron_physics ∉ (braginskii_fluid, kinetic_electrons)
+    elseif restart_electron_physics ∉ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         # Restarting from Boltzmann electron simulation, so start with constant
         # electron temperature
         begin_serial_region()
@@ -325,7 +328,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
     @views derivative_z!(moments.electron.dppar_dz, moments.electron.ppar, 
         scratch_dummy.buffer_rs_1[:,1], scratch_dummy.buffer_rs_2[:,1], scratch_dummy.buffer_rs_3[:,1],
         scratch_dummy.buffer_rs_4[:,1], z_spectral, z)
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         # Initialise the array for the electron pdf
         begin_serial_region()
         speed = @view scratch_dummy.buffer_vpavperpzrs_1[:,:,:,:,1]
@@ -344,7 +348,8 @@ function initialize_electrons!(pdf, moments, fields, geometry, composition, r, z
         electron_fluid_qpar_boundary_condition!(
             moments.electron.ppar, moments.electron.upar, moments.electron.dens,
             moments.electron, z)
-        if restart_electron_physics ∉ (nothing, braginskii_fluid, kinetic_electrons)
+        if restart_electron_physics ∉ (nothing, braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
             # Restarting from Boltzmann. If we use an exactly constant T_e profile,
             # qpar for the electrons will be non-zero only at the boundary points,
             # which will crash the code unless the timestep is insanely small. Give
@@ -570,7 +575,8 @@ function initialize_electron_pdf!(scratch, scratch_electron, pdf, moments, field
 
     # now that the initial electron pdf is given, the electron parallel heat flux should be updated
     # if using kinetic electrons
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         begin_serial_region()
         if t_input["no_restart"]
             restart_filename = nothing
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 0f27e36bb..21b7857d7 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -175,12 +175,14 @@ end
     boltzmann_electron_response_with_simple_sheath 
     braginskii_fluid
     kinetic_electrons
+    kinetic_electrons_with_temperature_equation
 end
 export electron_physics_type
 export boltzmann_electron_response
 export boltzmann_electron_response_with_simple_sheath
 export braginskii_fluid
 export kinetic_electrons
+export kinetic_electrons_with_temperature_equation
 
 """
 """
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index f7fb0171b..29c6d38c4 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -872,7 +872,9 @@ function reload_evolving_fields!(pdf, moments, fields, boundary_distributions,
             else
                 restart_electron_physics = boltzmann_electron_response
             end
-            if pdf.electron !== nothing && restart_electron_physics == kinetic_electrons
+            if pdf.electron !== nothing &&
+                    restart_electron_physics ∈ (kinetic_electrons,
+                                                kinetic_electrons_with_temperature_equation)
                 pdf.electron.norm .=
                     reload_electron_pdf(dynamic, time_index, moments, r, z, vperp, vpa,
                                         r_range, z_range, vperp_range, vpa_range,
@@ -3530,10 +3532,12 @@ function get_run_info_no_setup(run_dir::Union{AbstractString,Tuple{AbstractStrin
         if evolve_ppar
             push!(evolving_variables, "parallel_pressure")
         end
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics ∈ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             push!(evolving_variables, "f_electron")
         end
-        if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+        if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             push!(evolving_variables, "electron_parallel_pressure")
         end
         if composition.n_neutral_species > 0
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 0a0b5362b..3712ba994 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -315,7 +315,8 @@ function setup_moment_kinetics(input_dict::AbstractDict;
             @. moments.electron.temp = composition.me_over_mi * moments.electron.vth^2
             @. moments.electron.ppar = 0.5 * moments.electron.dens * moments.electron.temp
         end
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics ∈ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             begin_r_z_vperp_vpa_region()
             @loop_r_z_vperp_vpa ir iz ivperp ivpa begin
                 pdf.electron.pdf_before_ion_timestep[ivpa,ivperp,iz,ir] =
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 06bf692ae..2ee8d346f 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1087,7 +1087,9 @@ function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
             # error norm here.
             # Could do with a better way to sort the different possible types of
             # convergence failure...
-            if t_params.rk_coefs_implicit !== nothing && composition.electron_physics == kinetic_electrons
+            if t_params.rk_coefs_implicit !== nothing &&
+                    composition.electron_physics ∈ (kinetic_electrons,
+                                                    kinetic_electrons_with_temperature_equation)
                 if success == "nonlinear-solver"
                     t_params.failure_caused_by[end-1] += 1
                 elseif success == "kinetic-electrons"
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index cf0db6a19..923406ed8 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -401,7 +401,8 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         if composition.electron_physics != braginskii_fluid
             t_input["implicit_braginskii_conduction"] = false
         end
-        if composition.electron_physics != kinetic_electrons
+        if composition.electron_physics ∉ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             t_input["implicit_electron_advance"] = false
             t_input["implicit_electron_ppar"] = false
         end
@@ -497,7 +498,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         dt_reload = nothing
     end
 
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         electron_t_params = setup_time_info(t_input["electron_t_input"], 2, 0.0,
                                             electron_dt_reload,
                                             electron_dt_before_last_fail_reload,
@@ -538,7 +540,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         # ion pressure
         n_variables += 1
     end
-    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         # electron pressure
         n_variables += 1
     end
@@ -596,11 +599,13 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         push!(t_params.limit_caused_by, 0) # RK accuracy
         push!(t_params.failure_caused_by, 0)
     end
-    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         # electron pressure
         push!(t_params.limit_caused_by, 0) # RK accuracy
         push!(t_params.failure_caused_by, 0) # RK accuracy for electron_ppar
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics ∈ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             push!(t_params.failure_caused_by, 0) # Convergence failure for kinetic electron solve
         end
     end
@@ -747,7 +752,8 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     else
         scratch_implicit = nothing
     end
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         scratch_electron = setup_electron_scratch_arrays(moments, pdf,
                                                          t_params.electron.n_rk_stages+1)
     else
@@ -769,23 +775,11 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
 
     # Now that `t_params` and `scratch` have been created, initialize electrons if
     # necessary
-    if composition.electron_physics != restart_electron_physics
-        begin_serial_region()
-        @serial_region begin
-            # zero-initialise phi here, because the boundary points of phi are used as an
-            # effective 'cache' for the sheath-boundary cutoff speed for the electrons, so
-            # needs to be initialised to something, but phi cannot be calculated properly
-            # until after the electrons are initialised.
-            fields.phi .= 0.0
-        end
-        initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
-                              vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
-                              vperp_spectral, vpa_spectral, collisions, gyroavs,
-                              external_source_settings, scratch_dummy, scratch,
-                              scratch_electron, nl_solver_params, t_params, t_input,
-                              num_diss_params, advection_structs, io_input, input_dict;
-                              restart_electron_physics=restart_electron_physics)
-    elseif restarting && composition.electron_physics == kinetic_electrons
+    if restarting &&
+            composition.electron_physics ∈ (kinetic_electrons,
+                                            kinetic_electrons_with_temperature_equation) &&
+            restart_electron_physics ∈ (kinetic_electrons,
+                                        kinetic_electrons_with_temperature_equation)
         if t_params.electron.debug_io !== nothing
             # Create *.electron_debug.h5 file so that it can be re-opened in
             # update_electron_pdf!().
@@ -803,6 +797,22 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         resize!(t_params.electron.dfns_output_times, 0)
         t_params.electron.moments_output_counter[] = 1
         t_params.electron.dfns_output_counter[] = 1
+    elseif composition.electron_physics != restart_electron_physics
+        begin_serial_region()
+        @serial_region begin
+            # zero-initialise phi here, because the boundary points of phi are used as an
+            # effective 'cache' for the sheath-boundary cutoff speed for the electrons, so
+            # needs to be initialised to something, but phi cannot be calculated properly
+            # until after the electrons are initialised.
+            fields.phi .= 0.0
+        end
+        initialize_electrons!(pdf, moments, fields, geometry, composition, r, z,
+                              vperp, vpa, vzeta, vr, vz, z_spectral, r_spectral,
+                              vperp_spectral, vpa_spectral, collisions, gyroavs,
+                              external_source_settings, scratch_dummy, scratch,
+                              scratch_electron, nl_solver_params, t_params, t_input,
+                              num_diss_params, advection_structs, io_input, input_dict;
+                              restart_electron_physics=restart_electron_physics)
     end
 
     # update the derivatives of the electron moments as these may be needed when
@@ -1257,7 +1267,8 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         end
         # if treating the electrons as a fluid with Braginskii closure, or
         # moment-kinetically then advance the electron energy equation
-        if composition.electron_physics == kinetic_electrons
+        if composition.electron_physics ∈ (kinetic_electrons,
+                                           kinetic_electrons_with_temperature_equation)
             if !(t_params.implicit_electron_advance || t_params.implicit_electron_ppar)
                 advance_electron_energy = true
                 advance_electron_conduction = true
@@ -2312,7 +2323,8 @@ function rk_update!(scratch, scratch_implicit, moments, t_params, istage, compos
     # use Runge Kutta to update any velocity moments evolved separately from the pdf
     rk_update_evolved_moments!(scratch, scratch_implicit, moments, t_params, istage)
 
-    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         rk_update_variable!(scratch, scratch_implicit, :electron_ppar, t_params, istage)
     end
 
@@ -2395,7 +2407,8 @@ function apply_all_bcs_constraints_update_moments!(
                                            z_spectral,
                                            num_diss_params.electron.moment_dissipation_coefficient, 
                                            composition.electron_physics)
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         #max_electron_pdf_iterations = 1000
         #max_electron_sim_time = nothing
         max_electron_pdf_iterations = nothing
@@ -2616,7 +2629,8 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         begin_s_r_z_region()
         rk_loworder_solution!(scratch, scratch_implicit, :ppar, t_params)
     end
-    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         begin_r_z_region()
         rk_loworder_solution!(scratch, scratch_implicit, :electron_ppar, t_params)
     end
@@ -2740,7 +2754,8 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
         push!(total_points, z.n_global * r.n_global * n_ion_species)
     end
 
-    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons)
+    if composition.electron_physics ∈ (braginskii_fluid, kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         begin_r_z_region()
         electron_p_err = local_error_norm(scratch[2].electron_ppar,
                                           scratch[t_params.n_rk_stages+1].electron_ppar,
@@ -2842,7 +2857,8 @@ function adaptive_timestep_update!(scratch, scratch_implicit, scratch_electron,
                                        current_dt, error_norm_method, success,
                                        nl_max_its_fraction, composition)
 
-    if composition.electron_physics == kinetic_electrons
+    if composition.electron_physics ∈ (kinetic_electrons,
+                                       kinetic_electrons_with_temperature_equation)
         if t_params.previous_dt[] == 0.0
             # Reset electron pdf to its value at the beginning of this step.
             begin_r_z_vperp_vpa_region()
@@ -3470,8 +3486,8 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
     end
     if advance.electron_energy
         electron_energy_equation!(fvec_out.electron_ppar, fvec_in.electron_ppar,
-                                  fvec_in.density, fvec_in.electron_upar, fvec_in.upar,
-                                  fvec_in.ppar, fvec_in.density_neutral,
+                                  fvec_in.density, fvec_in.electron_upar, fvec_in.density,
+                                  fvec_in.upar, fvec_in.ppar, fvec_in.density_neutral,
                                   fvec_in.uz_neutral, fvec_in.pz_neutral,
                                   moments.electron, collisions, dt, composition,
                                   external_source_settings.electron, num_diss_params, z;
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 4881d0aad..29684c26a 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -41,7 +41,7 @@ using ..derivatives: derivative_z!, second_derivative_z!
 using ..derivatives: derivative_r!, second_derivative_r!
 using ..looping
 using ..gyroaverages: gyro_operators, gyroaverage_pdf!
-using ..input_structs: braginskii_fluid, kinetic_electrons
+using ..input_structs
 using ..moment_kinetics_structs: moments_ion_substruct, moments_electron_substruct,
                                  moments_neutral_substruct
 
@@ -236,7 +236,8 @@ function create_moments_electron(nz, nr, electron_model, num_diss_params)
     # need dupar/dz to obtain, e.g., the updated electron temperature
     dupar_dz = allocate_shared_float(nz, nr)
     dppar_dz = allocate_shared_float(nz, nr)
-    if electron_model ∈ (braginskii_fluid, kinetic_electrons)
+    if electron_model ∈ (braginskii_fluid, kinetic_electrons,
+                         kinetic_electrons_with_temperature_equation)
         dppar_dz_upwind = allocate_shared_float(nz, nr)
         dT_dz_upwind = allocate_shared_float(nz, nr)
     else
@@ -970,6 +971,7 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
     upar = scratch.electron_upar
     ppar = scratch.electron_ppar
     qpar = moments.electron.qpar
+    vth = moments.electron.vth
     dummy_zr = @view scratch_dummy.dummy_zrs[:,:,1]
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
@@ -993,6 +995,8 @@ function calculate_electron_moment_derivatives!(moments, scratch, scratch_dummy,
                             buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
     @views derivative_z!(moments.electron.dqpar_dz, qpar, buffer_r_1,
                             buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+    @views derivative_z!(moments.electron.dvth_dz, vth, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
     # calculate the zed derivative of the electron temperature
     @loop_r_z ir iz begin
         # store the temperature in dummy_zr

From eebe4fcd049b16acbbb5a10a9a5a642d166464cf Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 25 Jul 2024 13:38:31 +0100
Subject: [PATCH 408/428] Update example kinetic electron input file

---
 ...ecyclefraction0.5_split3_kinetic-vpadiss0.toml | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
index 7828c61fa..0a66ed75e 100644
--- a/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
+++ b/examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-vpadiss0.toml
@@ -1,7 +1,7 @@
 #runtime_plots = true
 n_ion_species = 1
 n_neutral_species = 1
-electron_physics = "kinetic_electrons"
+electron_physics = "kinetic_electrons_with_temperature_equation"
 evolve_moments_density = true
 evolve_moments_parallel_flow = true
 evolve_moments_parallel_pressure = true
@@ -54,7 +54,7 @@ z_discretization = "chebyshev_pseudospectral"
 z_element_spacing_option = "sqrt"
 vpa_ngrid = 6
 vpa_nelement = 63
-vpa_L = 36.0
+vpa_L = 48.0
 vpa_bc = "zero"
 vpa_discretization = "chebyshev_pseudospectral"
 vz_ngrid = 6
@@ -66,13 +66,13 @@ vz_discretization = "chebyshev_pseudospectral"
 [timestepping]
 type = "Fekete4(3)"
 #nstep = 50000
-nstep = 1000000
+nstep = 10000000
 dt = 1.0e-6
 minimum_dt = 1.0e-6
 max_increase_factor_near_last_fail = 1.001
 last_fail_proximity_factor = 1.1
-nwrite = 10000
-nwrite_dfns = 10000
+nwrite = 100000
+nwrite_dfns = 100000
 steady_state_residual = true
 converged_residual_value = 1.0e-3
 
@@ -85,7 +85,8 @@ nwrite_dfns = 100000
 #type = "SSPRK4"
 type = "Fekete4(3)"
 rtol = 1.0e-3
-minimum_dt = 1.0e-8
+atol = 1.0e-14
+minimum_dt = 1.0e-9
 initialization_residual_value = 2.5
 converged_residual_value = 0.1 #1.0e-3
 
@@ -103,7 +104,7 @@ source_T = 2.0
 force_minimum_pdf_value = 0.0
 
 [electron_numerical_dissipation]
-vpa_dissipation_coefficient = 2.0
+#vpa_dissipation_coefficient = 2.0
 force_minimum_pdf_value = 0.0
 
 [neutral_numerical_dissipation]

From 309cd96d5cfd7196c17d0050445eabe8f088dc9c Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 14:00:03 +0100
Subject: [PATCH 409/428] Revert "Implicit solve option for Maxwell diffusion"

This reverts commit aefc2787142a70311a2fd3aa3facb2a4c2957da6.
---
 moment_kinetics/src/coordinates.jl           |  13 +-
 moment_kinetics/src/input_structs.jl         |   1 -
 moment_kinetics/src/maxwell_diffusion.jl     | 345 +++++--------------
 moment_kinetics/src/moment_kinetics_input.jl |   1 -
 moment_kinetics/src/time_advance.jl          |  54 +--
 5 files changed, 107 insertions(+), 307 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index 6f723ffa9..8d6591ab8 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -86,10 +86,6 @@ struct coordinate{T <: AbstractVector{mk_float}}
     scratch6::Array{mk_float,1}
     # scratch7 is an array used for intermediate calculations requiring n entries
     scratch7::Array{mk_float,1}
-    # scratch8 is an array used for intermediate calculations requiring n entries
-    scratch8::Array{mk_float,1}
-    # scratch9 is an array used for intermediate calculations requiring n entries
-    scratch9::Array{mk_float,1}
     # scratch_shared is a shared-memory array used for intermediate calculations requiring
     # n entries
     scratch_shared::T
@@ -236,11 +232,10 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch),
         copy(scratch), copy(scratch), copy(scratch), copy(scratch), copy(scratch),
-        copy(scratch), copy(scratch), scratch_shared, scratch_shared2, scratch_2d,
-        copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
-        local_io_range, global_io_range, element_scale, element_shift,
-        input.element_spacing_option, element_boundaries, radau_first_element,
-        other_nodes, one_over_denominator)
+        scratch_shared, scratch_shared2, scratch_2d, copy(scratch_2d), advection,
+        send_buffer, receive_buffer, input.comm, local_io_range, global_io_range,
+        element_scale, element_shift, input.element_spacing_option, element_boundaries,
+        radau_first_element, other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 21b7857d7..34197d53b 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -85,7 +85,6 @@ struct time_info{Terrorsum <: Real, T_debug_output, T_electron, Trkimp, Timpzero
     implicit_electron_advance::Bool
     implicit_ion_advance::Bool
     implicit_vpa_advection::Bool
-    implicit_ion_maxwell_diffusion::Bool
     implicit_electron_ppar::Bool
     write_after_fixed_step_count::Bool
     error_sum_zero::Terrorsum
diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 431bec193..69ede0588 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -13,17 +13,12 @@ most valid here.
 
 module maxwell_diffusion
 
-export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!, implicit_ion_maxwell_diffusion!
+export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
 
 using ..looping
 using ..input_structs: mxwl_diff_collisions_input, set_defaults_and_check_section!
-using ..boundary_conditions: enforce_v_boundary_condition_local!, vpagrid_to_dzdt
 using ..calculus: derivative!, second_derivative!
-using ..moment_constraints: moment_constraints_on_residual!
-using ..moment_kinetics_structs: scratch_pdf
-using ..nonlinear_solvers: newton_solve!
 using ..reference_parameters: get_reference_collision_frequency_ii
-using ..velocity_moments: update_derived_moments!, calculate_ion_moment_derivatives!
 
 """
 Function for reading Maxwell diffusion operator input parameters. 
@@ -71,31 +66,13 @@ function setup_mxwl_diff_collisions_input(toml_input::Dict, reference_params)
     return mxwl_diff_collisions_input(; input...)
 end
 
-function ion_vpa_maxwell_diffusion_inner!(f_out, f_in, n, upar, vth, vpa, spectral,
-                                          diffusion_coefficient, dt, ::Val{true},
-                                          ::Val{true}, ::Val{true})
-    second_derivative!(vpa.scratch2, f_in, vpa, spectral)
-    @. vpa.scratch = vpa.grid * f_in
-    derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-    @. f_out += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
-end
-
-function ion_vpa_maxwell_diffusion_inner!(f_out, f_in, n, upar, vth, vpa, spectral,
-                                          diffusion_coefficient, dt, ::Val{false},
-                                          ::Val{false}, ::Val{false})
-    second_derivative!(vpa.scratch2, f_in, vpa, spectral)
-    @. vpa.scratch = (vpa.grid - upar) * f_in
-    derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-    @. f_out += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
-end
-
 """
 Calculate the Maxwellian associated with the current ion pdf moments, and then 
 subtract this from current pdf. Then take second derivative of this function
 to act as the diffusion operator. 
 """
-function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral, dt,
-                                    diffusion_coefficient)
+function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::T_spectral, 
+                                       dt, diffusion_coefficient) where T_spectral
     
     # If negative input (should be -1.0), then none of this diffusion will happen. 
     # This number can be put in as some parameter in the input file called something
@@ -108,10 +85,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral,
         error("Maxwell diffusion not implemented for 2V moment-kinetic cases yet")
     end
 
-    evolve_density = Val(moments.evolve_density)
-    evolve_upar = Val(moments.evolve_upar)
-    evolve_ppar = Val(moments.evolve_ppar)
-
     # Otherwise, build the maxwellian function (which is going to be subtracted from 
     # the current distribution) using the moments of the distribution (so that the 
     # operator itself conserves the moments), and then this result will be the one 
@@ -125,76 +98,97 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral,
     # - upar: working in peculiar velocity space, so no upar subtraction from vpa 
     # - ppar: normalisation by vth, in 1D is 1/vth prefactor, and grid is normalised by vth,
     # hence no 1/vth^2 term in the exponent.
-    @loop_s_r_z_vperp is ir iz ivperp begin
-        #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-        #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-        #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-        #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-        @views ion_vpa_maxwell_diffusion_inner!(f_out[:,ivperp,iz,ir,is],
-                                                f_in.pdf[:,ivperp,iz,ir,is],
-                                                f_in.density[iz,ir,is],
-                                                f_in.upar[iz,ir,is],
-                                                moments.ion.vth[iz,ir,is], vpa, spectral,
-                                                diffusion_coefficient, dt, evolve_density,
-                                                evolve_upar, evolve_ppar)
+    if moments.evolve_density && moments.evolve_upar && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+            #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
+            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
+            @views @. vpa.scratch = (vpa.grid - upar / vth) * f_in.pdf[:,ivperp,iz,ir,is]
+            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
+        end
+    elseif moments.evolve_density && moments.evolve_upar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            1.0 / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    elseif moments.evolve_density && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2))
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    elseif moments.evolve_upar && moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    elseif moments.evolve_density
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            1.0 / vth * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    elseif moments.evolve_upar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            vth = moments.ion.vth[iz,ir,is]
+            n = f_in.density[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    elseif moments.evolve_ppar
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            n * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
+        end
+        error("hack not implemented for this case")
+    else
+        # Drift kinetic version is the only one that currently can support 2V. 
+        @loop_s_r_z_vperp is ir iz ivperp begin
+            n = f_in.density[iz,ir,is]
+            upar = f_in.upar[iz,ir,is]
+            vth = moments.ion.vth[iz,ir,is]
+            #if vperp.n == 1
+            #    vth_prefactor = 1.0 / vth
+            #else
+            #    vth_prefactor = 1.0 / vth^3
+            #end
+            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
+            #                exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
+            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
+            @views @. vpa.scratch = (vpa.grid - upar) * f_in.pdf[:,ivperp,iz,ir,is]
+            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
+        end
     end
-    #elseif moments.evolve_density && moments.evolve_upar
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        vth = moments.ion.vth[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        1.0 / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
-    #elseif moments.evolve_density && moments.evolve_ppar
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        upar = f_in.upar[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2))
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
-    #elseif moments.evolve_upar && moments.evolve_ppar
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        n = f_in.density[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        n * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
-    #elseif moments.evolve_density
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        vth = moments.ion.vth[iz,ir,is]
-    #        upar = f_in.upar[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        1.0 / vth * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
-    #elseif moments.evolve_upar
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        vth = moments.ion.vth[iz,ir,is]
-    #        n = f_in.density[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        n / vth * exp(- ((vpa.grid[:])^2 + (vperp.grid[ivperp])^2)/(vth^2) )
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
-    #elseif moments.evolve_ppar
-    #    @loop_s_r_z_vperp is ir iz ivperp begin
-    #        n = f_in.density[iz,ir,is]
-    #        upar = f_in.upar[iz,ir,is]
-    #        @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-    #                        n * exp(- ((vpa.grid[:] - upar)^2 + (vperp.grid[ivperp])^2) )
-    #        second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-    #        @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-    #    end
-    #    error("hack not implemented for this case")
     return nothing
 end
 
@@ -278,155 +272,4 @@ function neutral_vz_maxwell_diffusion!(f_out, f_in, moments, vzeta, vr, vz, spec
     return nothing
 end
 
-"""
-"""
-function implicit_ion_maxwell_diffusion!(f_out, fvec_in, moments, z_advect, vpa, vperp, z,
-                                         r, dt, r_spectral, vpa_spectral, composition,
-                                         collisions, geometry, nl_solver_params, gyroavs,
-                                         scratch_dummy)
-    if vperp.n > 1 && (moments.evolve_density || moments.evolve_upar || moments.evolve_ppar)
-        error("Moment constraints in implicit_maxwell_diffusion!() do not support 2V runs yet")
-    end
-
-    # Ensure moments are consistent with f_new
-    new_scratch = scratch_pdf(f_out, fvec_in.density, fvec_in.upar, fvec_in.ppar,
-                              fvec_in.pperp, fvec_in.temp_z_s, fvec_in.electron_density,
-                              fvec_in.electron_upar, fvec_in.electron_ppar,
-                              fvec_in.electron_pperp, fvec_in.electron_temp,
-                              fvec_in.pdf_neutral, fvec_in.density_neutral,
-                              fvec_in.uz_neutral, fvec_in.pz_neutral)
-    update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
-                            r_spectral, geometry, gyroavs, scratch_dummy, z_advect, false)
-
-    begin_s_r_z_vperp_region()
-
-    evolve_density = Val(moments.evolve_density)
-    evolve_upar = Val(moments.evolve_upar)
-    evolve_ppar = Val(moments.evolve_ppar)
-    coords = (vpa=vpa,)
-    vpa_bc = vpa.bc
-    D_ii = collisions.mxwl_diff.D_ii
-    zero = 1.0e-14
-    @loop_s is begin
-        @loop_r_z_vperp ir iz ivperp begin
-            f_old_no_bc = @view fvec_in.pdf[:,ivperp,iz,ir,is]
-            this_f_out = @view f_out[:,ivperp,iz,ir,is]
-            n = fvec_in.density[iz,ir,is]
-            upar = fvec_in.upar[iz,ir,is]
-            vth = moments.ion.vth[iz,ir,is]
-
-            if z.irank == 0 && iz == 1
-                @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, moments.ion.vth[iz,ir,is],
-                                                 fvec_in.upar[iz,ir,is],
-                                                 moments.evolve_ppar,
-                                                 moments.evolve_upar)
-                icut_lower_z = vpa.n
-                for ivpa ∈ vpa.n:-1:1
-                    # for left boundary in zed (z = -Lz/2), want
-                    # f(z=-Lz/2, v_parallel > 0) = 0
-                    if vpa.scratch[ivpa] ≤ zero
-                        icut_lower_z = ivpa + 1
-                        break
-                    end
-                end
-            end
-            if z.irank == z.nrank - 1 && iz == z.n
-                @. vpa.scratch = vpagrid_to_dzdt(vpa.grid, moments.ion.vth[iz,ir,is],
-                                                 fvec_in.upar[iz,ir,is],
-                                                 moments.evolve_ppar,
-                                                 moments.evolve_upar)
-                icut_upper_z = 0
-                for ivpa ∈ 1:vpa.n
-                    # for right boundary in zed (z = Lz/2), want
-                    # f(z=Lz/2, v_parallel < 0) = 0
-                    if vpa.scratch[ivpa] ≥ -zero
-                        icut_upper_z = ivpa - 1
-                        break
-                    end
-                end
-            end
-
-            function apply_bc!(x)
-                # Boundary condition
-                enforce_v_boundary_condition_local!(x, vpa_bc, nothing, true,
-                                                    vpa, vpa_spectral)
-
-                if z.bc == "wall"
-                    # Wall boundary conditions. Note that as density, upar, ppar do not
-                    # change in this implicit step, f_new, f_old, and residual should all
-                    # be zero at exactly the same set of grid points, so it is reasonable
-                    # to zero-out `residual` to impose the boundary condition. We impose
-                    # this after subtracting f_old in case rounding errors, etc. mean that
-                    # at some point f_old had a different boundary condition cut-off
-                    # index.
-                    if z.irank == 0 && iz == 1
-                        x[icut_lower_z:end] .= 0.0
-                    end
-                    # absolute velocity at right boundary
-                    if z.irank == z.nrank - 1 && iz == z.n
-                        x[1:icut_upper_z] .= 0.0
-                    end
-                end
-            end
-
-            # Need to apply 'new' boundary conditions to `f_old`, so that by imposing them
-            # on `residual`, they are automatically imposed on `f_new`.
-            f_old = vpa.scratch9 .= f_old_no_bc
-            apply_bc!(f_old)
-
-            left_preconditioner = identity
-            right_preconditioner = identity
-
-            # Define a function whose input is `f_new`, so that when it's output
-            # `residual` is zero, f_new is the result of a backward-Euler timestep:
-            #   (f_new - f_old) / dt = RHS(f_new)
-            # ⇒ f_new - f_old - dt*RHS(f_new) = 0
-            function residual_func!(residual, f_new)
-                apply_bc!(f_new)
-                residual .= f_old
-                ion_vpa_maxwell_diffusion_inner!(residual, f_new, n, upar, vth, vpa,
-                                                 vpa_spectral, D_ii, dt, evolve_density,
-                                                 evolve_upar, evolve_ppar)
-
-                # Now
-                #   residual = f_old + dt*RHS(f_new)
-                # so update to desired residual
-                @. residual = f_new - residual
-
-                apply_bc!(residual)
-                moment_constraints_on_residual!(residual, f_new, moments, vpa)
-            end
-
-            # Buffers
-            # Note vpa.scratch, vpa.scratch2 and vpa.scratch3 are used by advance_f!, so
-            # we cannot use it here.
-            residual = vpa.scratch4
-            delta_x = vpa.scratch5
-            rhs_delta = vpa.scratch6
-            v = vpa.scratch7
-            w = vpa.scratch8
-
-            # Use forward-Euler step for initial guess
-            # By passing this_f_out, which is equal to f_old at this point, the 'residual'
-            # is
-            #   f_new - f_old - dt*RHS(f_old) = -dt*RHS(f_old)
-            # so to get a forward-Euler step we have to subtract this 'residual'
-            residual_func!(residual, this_f_out)
-            this_f_out .-= residual
-
-            success = newton_solve!(this_f_out, residual_func!, residual, delta_x,
-                                    rhs_delta, v, w, nl_solver_params, coords=coords,
-                                    left_preconditioner=left_preconditioner,
-                                    right_preconditioner=right_preconditioner)
-            if !success
-                return success
-            end
-        end
-    end
-
-    nl_solver_params.stage_counter[] += 1
-
-    return true
-end
-
-end # maxwell_diffusion
+end # maxwell_diffusion
\ No newline at end of file
diff --git a/moment_kinetics/src/moment_kinetics_input.jl b/moment_kinetics/src/moment_kinetics_input.jl
index 2dd0f6b0d..2ac8e08ec 100644
--- a/moment_kinetics/src/moment_kinetics_input.jl
+++ b/moment_kinetics/src/moment_kinetics_input.jl
@@ -233,7 +233,6 @@ function mk_input(scan_input=Dict(); save_inputs_to_txt=false, ignore_MPI=true)
         implicit_electron_advance=true,
         implicit_ion_advance=false,
         implicit_vpa_advection=false,
-        implicit_ion_maxwell_diffusion=false,
         implicit_electron_ppar=false,
         write_after_fixed_step_count=false,
         write_error_diagnostics=false,
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 923406ed8..0dee6283d 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -51,8 +51,7 @@ using ..ionization: ion_ionization_collisions_1V!, neutral_ionization_collisions
                     ion_ionization_collisions_3V!, neutral_ionization_collisions_3V!,
                     constant_ionization_source!
 using ..krook_collisions: krook_collisions!
-using ..maxwell_diffusion: ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!,
-                           implicit_ion_maxwell_diffusion!
+using ..maxwell_diffusion: ion_vpa_maxwell_diffusion!, neutral_vz_maxwell_diffusion!
 using ..external_sources
 using ..nonlinear_solvers
 using ..numerical_dissipation: vpa_boundary_buffer_decay!,
@@ -395,7 +394,6 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
         t_input["implicit_electron_advance"] = false
         t_input["implicit_ion_advance"] = false
         t_input["implicit_vpa_advection"] = false
-        t_input["implicit_ion_maxwell_diffusion"] = false
         t_input["implicit_electron_ppar"] = false
     else
         if composition.electron_physics != braginskii_fluid
@@ -457,7 +455,6 @@ function setup_time_info(t_input, n_variables, code_time, dt_reload,
                      electron !== nothing && t_input["implicit_electron_advance"],
                      electron !== nothing && t_input["implicit_ion_advance"],
                      electron !== nothing && t_input["implicit_vpa_advection"],
-                     electron !== nothing && t_input["implicit_ion_maxwell_diffusion"],
                      electron !== nothing && t_input["implicit_electron_ppar"],
                      t_input["write_after_fixed_step_count"], error_sum_zero,
                      t_input["split_operators"], t_input["steady_state_residual"],
@@ -705,33 +702,11 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     else
         nl_solver_vpa_advection_params = nothing
     end
-    if t_params.implicit_ion_maxwell_diffusion
-        # Implicit solve for ion maxwell diffusion term should be done in serial, as it
-        # will be called within a parallelised s_r_z_vperp loop.
-        nl_solver_maxwell_diffusion_params =
-            setup_nonlinear_solve(input_dict, (vpa=vpa,),
-                                  (composition.n_ion_species, r, z, vperp);
-                                  default_rtol=t_params.rtol / 10.0,
-                                  default_atol=t_params.atol / 10.0,
-                                  serial_solve=true, preconditioner_type="lu")
-    else
-        nl_solver_maxwell_diffusion_params = nothing
-    end
     if nl_solver_ion_advance_params !== nothing &&
             nl_solver_vpa_advection_params !== nothing
         error("Cannot use implicit_ion_advance and implicit_vpa_advection at the same "
               * "time")
     end
-    if nl_solver_ion_advance_params !== nothing &&
-            nl_solver_maxwell_diffusion_params !== nothing
-        error("Cannot use implicit_ion_advance and implicit_ion_maxwell_diffusion at the same "
-              * "time")
-    end
-    if nl_solver_vpa_advection_params !== nothing &&
-            nl_solver_maxwell_diffusion_params !== nothing
-        error("Cannot use implicit_vpa_advection and implicit_ion_maxwell_diffusion at the same "
-              * "time")
-    end
     if nl_solver_electron_advance_params !== nothing && t_params.implicit_electron_ppar
         error("Cannot use implicit_electron_advance and implicit_electron_ppar at the "
               * "same time.")
@@ -739,8 +714,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
     nl_solver_params = (electron_conduction=electron_conduction_nl_solve_parameters,
                         electron_advance=nl_solver_electron_advance_params,
                         ion_advance=nl_solver_ion_advance_params,
-                        vpa_advection=nl_solver_vpa_advection_params,
-                        maxwell_diffusion=nl_solver_maxwell_diffusion_params,)
+                        vpa_advection=nl_solver_vpa_advection_params,)
 
     begin_serial_region()
 
@@ -1222,7 +1196,7 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         if collisions.krook.nuii0 > 0.0
             advance_krook_collisions_ii = !t_params.implicit_ion_advance
         end
-        if collisions.mxwl_diff.D_ii > 0.0 && !t_params.implicit_ion_maxwell_diffusion
+        if collisions.mxwl_diff.D_ii > 0.0
             advance_maxwell_diffusion_ii = true
         end
         if collisions.mxwl_diff.D_nn > 0.0
@@ -1296,9 +1270,9 @@ function setup_advance_flags(moments, composition, t_params, collisions,
         # flag to determine if a d^2/dvpa^2 operator is present
         # When using implicit_vpa_advection, the vpa diffusion is included in the implicit
         # step
-        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1) || collisions.mxwl_diff.D_ii > 0.0)
+        vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1) || advance_maxwell_diffusion_ii)
         vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
-        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0 || collisions.mxwl_diff.D_nn > 0.0)
+        vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0 || advance_maxwell_diffusion_nn)
     end
 
     manufactured_solns_test = manufactured_solns_input.use_for_advance
@@ -1308,8 +1282,8 @@ function setup_advance_flags(moments, composition, t_params, collisions,
                         advance_neutral_vz_advection, advance_ion_cx, advance_neutral_cx,
                         advance_ion_cx_1V, advance_neutral_cx_1V, advance_ion_ionization,
                         advance_neutral_ionization, advance_ion_ionization_1V,
-                        advance_neutral_ionization_1V,
-                        advance_ionization_source, advance_krook_collisions_ii,
+                        advance_neutral_ionization_1V, advance_ionization_source,
+                        advance_krook_collisions_ii,
                         advance_maxwell_diffusion_ii, advance_maxwell_diffusion_nn,
                         explicit_weakform_fp_collisions,
                         advance_external_source, advance_ion_numerical_dissipation,
@@ -1410,8 +1384,6 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     elseif t_params.implicit_vpa_advection
         advance_vpa_advection = true
         advance_ion_numerical_dissipation = true
-    elseif t_params.implicit_ion_maxwell_diffusion && collisions.mxwl_diff.D_ii > 0.0
-        advance_maxwell_diffusion_ii = true
     end
     # *_diffusion flags are set regardless of whether diffusion is included in explicit or
     # implicit part of timestep, because they are used for boundary conditions, not to
@@ -1422,9 +1394,9 @@ function setup_implicit_advance_flags(moments, composition, t_params, collisions
     # flag to determine if a d^2/dvpa^2 operator is present
     # When using implicit_vpa_advection, the vpa diffusion is included in the implicit
     # step
-    vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1)) || (collisions.mxwl_diff.D_ii > 0.0)
+    vpa_diffusion = ((num_diss_params.ion.vpa_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
     vperp_diffusion = ((num_diss_params.ion.vperp_dissipation_coefficient > 0.0) || (collisions.fkpl.nuii > 0.0 && vperp.n > 1))
-    vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0) || (collisions.mxwl_diff.D_nn > 0.0)
+    vz_diffusion = (num_diss_params.neutral.vz_dissipation_coefficient > 0.0)
 
     if t_params.implicit_braginskii_conduction
         # if treating the electrons as a fluid with Braginskii closure, and using an IMEX
@@ -3153,7 +3125,6 @@ function ssp_rk!(pdf, scratch, scratch_implicit, scratch_electron, t_params, vz,
 
     reset_nonlinear_per_stage_counters(nl_solver_params.ion_advance)
     reset_nonlinear_per_stage_counters(nl_solver_params.vpa_advection)
-    reset_nonlinear_per_stage_counters(nl_solver_params.maxwell_diffusion)
 
     if t_params.previous_dt[] > 0.0
         istage = n_rk_stages+1
@@ -3571,13 +3542,6 @@ function backward_euler!(fvec_out, fvec_in, scratch_electron, pdf, fields, momen
                                           nl_solver_params.vpa_advection,
                                           advance.vpa_diffusion, num_diss_params, gyroavs,
                                           scratch_dummy)
-    elseif advance.mxwl_diff_collisions_ii
-        success = implicit_ion_maxwell_diffusion!(fvec_out.pdf, fvec_in, moments,
-                                                  z_advect, vpa, vperp, z, r, dt,
-                                                  r_spectral, vpa_spectral, composition,
-                                                  collisions, geometry,
-                                                  nl_solver_params.maxwell_diffusion,
-                                                  gyroavs, scratch_dummy)
     end
 
     return success

From 7c5f9262d2ad1cd45fbcad0593b311f331de966d Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 14:00:27 +0100
Subject: [PATCH 410/428] Revert "Hacky conversion to Mousseau's 'collision
 operator'"

This reverts commit 8aa7e1e986d94fee81d520644c5a68cac41f8010.
---
 moment_kinetics/src/maxwell_diffusion.jl | 44 ++++++++----------------
 1 file changed, 14 insertions(+), 30 deletions(-)

diff --git a/moment_kinetics/src/maxwell_diffusion.jl b/moment_kinetics/src/maxwell_diffusion.jl
index 69ede0588..03b4daa96 100644
--- a/moment_kinetics/src/maxwell_diffusion.jl
+++ b/moment_kinetics/src/maxwell_diffusion.jl
@@ -17,7 +17,7 @@ export setup_mxwl_diff_collisions_input, ion_vpa_maxwell_diffusion!, neutral_vz_
 
 using ..looping
 using ..input_structs: mxwl_diff_collisions_input, set_defaults_and_check_section!
-using ..calculus: derivative!, second_derivative!
+using ..calculus: second_derivative!
 using ..reference_parameters: get_reference_collision_frequency_ii
 
 """
@@ -100,17 +100,10 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
     # hence no 1/vth^2 term in the exponent.
     if moments.evolve_density && moments.evolve_upar && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
-            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
-            #                exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
-            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            #@views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
-            n = f_in.density[iz,ir,is]
-            upar = f_in.upar[iz,ir,is]
-            vth = moments.ion.vth[iz,ir,is]
-            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
-            @views @. vpa.scratch = (vpa.grid - upar / vth) * f_in.pdf[:,ivperp,iz,ir,is]
-            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vpa.scratch2 + vpa.scratch3)
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - 
+                            exp(-((vpa.grid[:])^2 + (vperp.grid[ivperp])^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
     elseif moments.evolve_density && moments.evolve_upar
         @loop_s_r_z_vperp is ir iz ivperp begin
@@ -120,7 +113,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     elseif moments.evolve_density && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             upar = f_in.upar[iz,ir,is]
@@ -129,7 +121,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     elseif moments.evolve_upar && moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
@@ -138,7 +129,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     elseif moments.evolve_density
         @loop_s_r_z_vperp is ir iz ivperp begin
             vth = moments.ion.vth[iz,ir,is]
@@ -148,7 +138,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     elseif moments.evolve_upar
         @loop_s_r_z_vperp is ir iz ivperp begin
             vth = moments.ion.vth[iz,ir,is]
@@ -158,7 +147,6 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     elseif moments.evolve_ppar
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
@@ -168,25 +156,21 @@ function ion_vpa_maxwell_diffusion!(f_out, f_in, moments, vpa, vperp, spectral::
             second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
             @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
-        error("hack not implemented for this case")
     else
         # Drift kinetic version is the only one that currently can support 2V. 
         @loop_s_r_z_vperp is ir iz ivperp begin
             n = f_in.density[iz,ir,is]
             upar = f_in.upar[iz,ir,is]
             vth = moments.ion.vth[iz,ir,is]
-            #if vperp.n == 1
-            #    vth_prefactor = 1.0 / vth
-            #else
-            #    vth_prefactor = 1.0 / vth^3
-            #end
-            #@views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
-            #                exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
-            #second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
-            @views second_derivative!(vpa.scratch2, f_in.pdf[:,ivperp,iz,ir,is], vpa, spectral)
-            @views @. vpa.scratch = (vpa.grid - upar) * f_in.pdf[:,ivperp,iz,ir,is]
-            derivative!(vpa.scratch3, vpa.scratch, vpa, spectral)
-            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * n / vth^3 * (vth^2 * vpa.scratch2 + vpa.scratch3)
+            if vperp.n == 1
+                vth_prefactor = 1.0 / vth
+            else
+                vth_prefactor = 1.0 / vth^3
+            end
+            @views @. vpa.scratch = f_in.pdf[:,ivperp,iz,ir,is] - n * vth_prefactor * 
+                            exp(-( ((vpa.grid[:] - upar)^2) + (vperp.grid[ivperp])^2)/(vth^2) )
+            second_derivative!(vpa.scratch2, vpa.scratch, vpa, spectral)
+            @views @. f_out[:,ivperp,iz,ir,is] += dt * diffusion_coefficient * vpa.scratch2
         end
     end
     return nothing

From dd643a10acdfd37c69a9e2f6c4903d8e52bbce4e Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 14:02:23 +0100
Subject: [PATCH 411/428] Add scratch8 and scratch9 buffers to `coordinate`
 struct

---
 moment_kinetics/src/coordinates.jl             | 13 +++++++++----
 moment_kinetics/test/nonlinear_solver_tests.jl |  6 ++++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index 8d6591ab8..6f723ffa9 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -86,6 +86,10 @@ struct coordinate{T <: AbstractVector{mk_float}}
     scratch6::Array{mk_float,1}
     # scratch7 is an array used for intermediate calculations requiring n entries
     scratch7::Array{mk_float,1}
+    # scratch8 is an array used for intermediate calculations requiring n entries
+    scratch8::Array{mk_float,1}
+    # scratch9 is an array used for intermediate calculations requiring n entries
+    scratch9::Array{mk_float,1}
     # scratch_shared is a shared-memory array used for intermediate calculations requiring
     # n entries
     scratch_shared::T
@@ -232,10 +236,11 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
         cell_width, igrid, ielement, imin, imax, igrid_full, input.discretization, input.fd_option, input.cheb_option,
         input.bc, wgts, uniform_grid, duniform_dgrid, scratch, copy(scratch),
         copy(scratch), copy(scratch), copy(scratch), copy(scratch), copy(scratch),
-        scratch_shared, scratch_shared2, scratch_2d, copy(scratch_2d), advection,
-        send_buffer, receive_buffer, input.comm, local_io_range, global_io_range,
-        element_scale, element_shift, input.element_spacing_option, element_boundaries,
-        radau_first_element, other_nodes, one_over_denominator)
+        copy(scratch), copy(scratch), scratch_shared, scratch_shared2, scratch_2d,
+        copy(scratch_2d), advection, send_buffer, receive_buffer, input.comm,
+        local_io_range, global_io_range, element_scale, element_shift,
+        input.element_spacing_option, element_boundaries, radau_first_element,
+        other_nodes, one_over_denominator)
 
     if coord.n == 1 && occursin("v", coord.name)
         spectral = null_velocity_dimension_info()
diff --git a/moment_kinetics/test/nonlinear_solver_tests.jl b/moment_kinetics/test/nonlinear_solver_tests.jl
index 5c0b08dfe..23d62acf9 100644
--- a/moment_kinetics/test/nonlinear_solver_tests.jl
+++ b/moment_kinetics/test/nonlinear_solver_tests.jl
@@ -57,7 +57,8 @@ function linear_test()
                                zeros(mk_float, 0, 0), zeros(mk_float, 0),
                                zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
                                zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
-                               zeros(mk_float, 0), zeros(mk_float, 0),
+                               zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
+                               zeros(mk_float, 0),
                                zeros(mk_float, 0, 0), zeros(mk_float, 0, 0),
                                advection_input("", 0.0, 0.0, 0.0), zeros(mk_float, 0),
                                zeros(mk_float, 0), MPI.COMM_NULL, 1:n, 1:n,
@@ -168,7 +169,8 @@ function nonlinear_test()
                                zeros(mk_float, 0, 0), zeros(mk_float, 0),
                                zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
                                zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
-                               zeros(mk_float, 0), zeros(mk_float, 0),
+                               zeros(mk_float, 0), zeros(mk_float, 0), zeros(mk_float, 0),
+                               zeros(mk_float, 0),
                                zeros(mk_float, 0, 0), zeros(mk_float, 0, 0),
                                advection_input("", 0.0, 0.0, 0.0), zeros(mk_float, 0),
                                zeros(mk_float, 0), MPI.COMM_NULL, 1:n, 1:n,

From da8088ce246492877f9fed0d8b9e9c12b6fbe306 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 18:55:29 +0100
Subject: [PATCH 412/428] Limit electron timestep when using `ion_dt` for
 implicit electron_ppar

The electron timestep needs to be less than `ion_dt` to prevent the
source term that makes the solve for `electron_ppar` into a
backward-Euler timestep from being numerically unstable (with explicit
timestepping).

Also add a maybe-improved initial guess for the implicitly-updated
`electron_ppar`, by taking a forward Euler step using `ion_dt` and the
initial values of `moments.electron.qpar`, etc.
---
 .../src/electron_kinetic_equation.jl          | 33 +++++++++++++++----
 moment_kinetics/src/runge_kutta.jl            |  7 ++--
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index dd828ea3a..982befae2 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -163,10 +163,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
 
     begin_r_z_region()
 
-    if ion_dt !== nothing
-        evolve_ppar = true
-    end
-
     # create several (r) dimension dummy arrays for use in taking derivatives
     buffer_r_1 = @view scratch_dummy.buffer_rs_1[:,1]
     buffer_r_2 = @view scratch_dummy.buffer_rs_2[:,1]
@@ -191,6 +187,20 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                            num_diss_params.electron.moment_dissipation_coefficient,
                                            composition.electron_physics)
 
+    if ion_dt !== nothing
+        evolve_ppar = true
+
+        # Use forward-Euler step (with `ion_dt` as the timestep) as initial guess for
+        # updated electron_ppar
+        electron_energy_equation!(scratch[t_params.n_rk_stages+1].electron_ppar,
+                                  moments.electron.ppar, moments.electron.dens,
+                                  moments.electron.upar, moments.ion.dens,
+                                  moments.ion.upar, moments.ion.ppar,
+                                  moments.neutral.dens, moments.neutral.uz,
+                                  moments.neutral.pz, moments.electron, collisions,
+                                  ion_dt, composition, external_source_settings.electron,
+                                  num_diss_params, z)
+    end
 
     # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
     @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
@@ -397,13 +407,21 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
             update_derived_moments_and_derivatives()
 
             if t_params.adaptive && istage == t_params.n_rk_stages
+                if ion_dt === nothing
+                    local_max_dt = Inf
+                else
+                    # Ensure timestep is not too big, so that d(electron_ppar)/dt 'source
+                    # term' is numerically stable.
+                    local_max_dt = 0.5 * ion_dt
+                end
                 electron_adaptive_timestep_update!(scratch, t_params.t[], t_params,
                                                    moments, phi, z_advect, vpa_advect,
                                                    composition, r, z, vperp, vpa,
                                                    vperp_spectral, vpa_spectral,
                                                    external_source_settings,
                                                    num_diss_params;
-                                                   evolve_ppar=evolve_ppar)
+                                                   evolve_ppar=evolve_ppar,
+                                                   local_max_dt=local_max_dt)
                 # Re-do this in case electron_adaptive_timestep_update!() re-arranged the
                 # `scratch` vector
                 new_scratch = scratch[istage+1]
@@ -1446,7 +1464,7 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
                                             vpa_advect, composition, r, z, vperp, vpa,
                                             vperp_spectral, vpa_spectral,
                                             external_source_settings, num_diss_params;
-                                            evolve_ppar=false)
+                                            evolve_ppar=false, local_max_dt=Inf)
     #error_norm_method = "Linf"
     error_norm_method = "L2"
 
@@ -1547,7 +1565,8 @@ function electron_adaptive_timestep_update!(scratch, t, t_params, moments, phi,
 
     adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms, total_points,
                                        current_dt, error_norm_method, "", 0.0,
-                                       composition; electron=true)
+                                       composition; electron=true,
+                                       local_max_dt=local_max_dt)
     if t_params.previous_dt[] == 0.0
         # Timestep failed, so reset  scratch[t_params.n_rk_stages+1] equal to
         # scratch[1] to start the timestep over.
diff --git a/moment_kinetics/src/runge_kutta.jl b/moment_kinetics/src/runge_kutta.jl
index 2ee8d346f..16d0ef600 100644
--- a/moment_kinetics/src/runge_kutta.jl
+++ b/moment_kinetics/src/runge_kutta.jl
@@ -1001,7 +1001,7 @@ Use the calculated `CFL_limits` and `error_norms` to update the timestep in `t_p
 function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
                                             total_points, current_dt, error_norm_method,
                                             success, nl_max_its_fraction, composition;
-                                            electron=false)
+                                            electron=false, local_max_dt::mk_float=Inf)
     # Get global minimum of CFL limits
     CFL_limit = nothing
     this_limit_caused_by = nothing
@@ -1218,8 +1218,9 @@ function adaptive_timestep_update_t_params!(t_params, CFL_limits, error_norms,
                 end
 
                 # Prevent timestep from going above maximum_dt
-                if t_params.dt[] > t_params.maximum_dt
-                    t_params.dt[] = t_params.maximum_dt
+                max_dt = min(t_params.maximum_dt, local_max_dt)
+                if t_params.dt[] > max_dt
+                    t_params.dt[] = max_dt
                     this_limit_caused_by = 4
                 end
 

From 47d83897d00a1401822356dd0c522f060493d360 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 19:24:23 +0100
Subject: [PATCH 413/428] Remove redundant computation of
 `moments.electron.dqpar_dz`

This computation is now taken care of in
`calculate_electron_moment_derivatives!()`.
---
 moment_kinetics/src/electron_kinetic_equation.jl | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 982befae2..039db58ac 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -202,10 +202,6 @@ function update_electron_pdf_with_time_advance!(scratch, pdf, moments, phi, coll
                                   num_diss_params, z)
     end
 
-    # compute the z-derivative of the input electron parallel heat flux, needed for the electron kinetic equation
-    @views derivative_z!(moments.electron.dqpar_dz, moments.electron.qpar, buffer_r_1,
-                         buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
-
     if !evolve_ppar
         # ppar is not updated in the pseudo-timestepping loop below. So that we can read
         # ppar from the scratch structs, copy moments.electron.ppar into all of them.

From 29f630b6d67487616a882f6c2338ddd51c64319c Mon Sep 17 00:00:00 2001
From: Michael Hardman <michael.hardman@tokamakenergy.com>
Date: Thu, 1 Aug 2024 07:59:15 +0000
Subject: [PATCH 414/428] Change default so that dirichlet_bc=false, and
 feature permitting ODE solve with K and L matrices is suppressed (but can be
 revived in a custom script with the dirichlet_bc=true optional argument in
 setup_gausslegendre_pseudospectral()).

---
 moment_kinetics/src/coordinates.jl    | 3 +--
 moment_kinetics/src/gauss_legendre.jl | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/coordinates.jl b/moment_kinetics/src/coordinates.jl
index d034a2dd0..b761ea1fa 100644
--- a/moment_kinetics/src/coordinates.jl
+++ b/moment_kinetics/src/coordinates.jl
@@ -252,8 +252,7 @@ function define_coordinate(input, parallel_io::Bool=false; run_directory=nothing
     elseif input.discretization == "gausslegendre_pseudospectral"
         # create arrays needed for explicit GaussLegendre pseudospectral treatment in this
         # coordinate and create the matrices for differentiation
-        spectral = setup_gausslegendre_pseudospectral(coord, collision_operator_dim=collision_operator_dim,
-                                                      dirichlet_bc=occursin("zero", coord.bc))
+        spectral = setup_gausslegendre_pseudospectral(coord, collision_operator_dim=collision_operator_dim)
         # obtain the local derivatives of the uniform grid with respect to the used grid
         derivative!(coord.duniform_dgrid, coord.uniform_grid, coord, spectral)
     else
diff --git a/moment_kinetics/src/gauss_legendre.jl b/moment_kinetics/src/gauss_legendre.jl
index 02eaf8b80..9a0afb6b3 100644
--- a/moment_kinetics/src/gauss_legendre.jl
+++ b/moment_kinetics/src/gauss_legendre.jl
@@ -100,7 +100,7 @@ struct gausslegendre_info{TSparse, TLU} <: weak_discretization_info
     Qmat::Array{mk_float,2}
 end
 
-function setup_gausslegendre_pseudospectral(coord; collision_operator_dim=true, dirichlet_bc=true)
+function setup_gausslegendre_pseudospectral(coord; collision_operator_dim=true, dirichlet_bc=false)
     lobatto = setup_gausslegendre_pseudospectral_lobatto(coord,collision_operator_dim=collision_operator_dim)
     radau = setup_gausslegendre_pseudospectral_radau(coord,collision_operator_dim=collision_operator_dim)
 

From 2fd62a3bf14e80b948213d46d460a6df5b480392 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 23:21:04 +0100
Subject: [PATCH 415/428] Update some test expected outputs

Mosty due to the changes in "Normalise grid to initial_temperature while
initialising f_neutral" (ac7f0a35). A couple of tests were slightly
affected by a couple of other small changes.
---
 .../test/braginskii_electrons_imex_tests.jl   | 180 +++++++++---------
 moment_kinetics/test/harrisonthompson.jl      |  58 +++---
 .../test/recycling_fraction_tests.jl          | 156 ++++++++-------
 moment_kinetics/test/wall_bc_tests.jl         |  26 +--
 4 files changed, 210 insertions(+), 210 deletions(-)

diff --git a/moment_kinetics/test/braginskii_electrons_imex_tests.jl b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
index 97f0a4f33..f7a67165d 100644
--- a/moment_kinetics/test/braginskii_electrons_imex_tests.jl
+++ b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
@@ -100,8 +100,8 @@ end
 """
 Run a test for a single set of parameters
 """
-function run_test(test_input, expected_p, expected_q, expected_vt; rtol=4.e-14,
-                  atol=1.e-15, args...)
+function run_test(test_input, expected_p, expected_q, expected_vt; rtol=1.e-6,
+                  atol=1.e-8, args...)
     # by passing keyword arguments to run_test, args becomes a Tuple of Pairs which can be
     # used to update the default inputs
 
@@ -175,7 +175,7 @@ function run_test(test_input, expected_p, expected_q, expected_vt; rtol=4.e-14,
             println("data tested would be: ", actual_q)
             @test false
         else
-            @test isapprox(actual_q, expected_q, rtol=rtol, atol=atol)
+            @test isapprox(actual_q, expected_q, rtol=10.0*rtol, atol=atol)
         end
         if expected_vt == nothing
             # Error: no expected input provided
@@ -191,93 +191,93 @@ function runtests()
     # Create a temporary directory for test output
     test_output_directory = get_MPI_tempdir()
 
-    expected_p = [0.44951398349003646, 0.4481071794001275, 0.4446923360697697,
-                  0.44101458684278777, 0.4384371320097695, 0.43732241798810445,
-                  0.4369770330480599, 0.4358502579878266, 0.434769684040917,
-                  0.43442263207284637, 0.43458925629798073, 0.4346831326574436,
-                  0.4350671485602161, 0.4362345316295716, 0.438235619583819,
-                  0.4402734411528, 0.4411447249385275, 0.4420727185406772,
-                  0.44481447378612293, 0.44890303238070045, 0.4530257687705389,
-                  0.4554124231936834, 0.4562726972649091, 0.4597319740656588,
-                  0.46533878727932276, 0.4713773067908409, 0.47559731298259184,
-                  0.4767158034722211, 0.47998294113890483, 0.4863021470825949,
-                  0.4937119852803594, 0.49963221191774776, 0.5018851701590787,
-                  0.5041515792976728, 0.5102381811386025, 0.5181350099985299,
-                  0.525127401588624, 0.52883991125208, 0.530126542118299,
-                  0.5350544857912003, 0.5423139848496283, 0.5492817298572252,
-                  0.5536930071759895, 0.5548035254843448, 0.5579127055178238,
-                  0.5633811620967428, 0.5689229399038384, 0.5726992489776701,
-                  0.5739864539763129, 0.5751981737538447, 0.5780380624660489,
-                  0.5808091879153607, 0.582371244425747, 0.5828442626520689,
-                  0.5829483019641976, 0.5830505324419235, 0.582289662647816,
-                  0.5804174302674363, 0.5785673270305791, 0.5780104294159015,
-                  0.5762363804591635, 0.5722372963746002, 0.5667182906200348,
-                  0.5617560490141181, 0.5597527656925925, 0.5576781862353368,
-                  0.5518325952127137, 0.5437117958155238, 0.5360738744744328,
-                  0.5318632735767702, 0.5303800815638172, 0.52458857332585,
-                  0.5157429166779243, 0.5068932260184091, 0.5010911813300937,
-                  0.4996034253232407, 0.49537353092296554, 0.4876627194305651,
-                  0.47936652714814293, 0.4732825296867743, 0.47108892532041974,
-                  0.46894816235702474, 0.4635185649133237, 0.4571465360817948,
-                  0.452118238600268, 0.4496786524823512]
-    expected_q = [0.6831704519267778, 0.6763596731108003, 0.6550248548395122,
-                  0.621137479330641, 0.5855826766086369, 0.564624348243424,
-                  0.5570221429154738, 0.5262612511440813, 0.4759763560524193,
-                  0.4215209814677685, 0.3833879518808367, 0.3732765291862611,
-                  0.34373639264749195, 0.2865988207461233, 0.21960124384883264,
-                  0.16604253330270563, 0.14564326706615302, 0.1251073951350269,
-                  0.06985203286011953, -0.002189422724091073, -0.06650025444931523,
-                  -0.10092291876877173, -0.11290749491826756, -0.15911660726647084,
-                  -0.22829407549560776, -0.296381042779731, -0.34065325931034834,
-                  -0.35197568475314056, -0.38412394718970916, -0.44264467058363494,
-                  -0.5056139329952417, -0.5518262734055381, -0.5684878570815051,
-                  -0.5847404886752257, -0.6258728785916116, -0.6737205371368551,
-                  -0.7106893697823674, -0.728139347093567, -0.733817023138346,
-                  -0.7537144673017869, -0.7772295129285097, -0.7923028306406594,
-                  -0.7973216863277829, -0.7979485411906018, -0.7981733058556373,
-                  -0.7921566532586084, -0.7749750759560777, -0.7538532886753724,
-                  -0.7441527309138952, -0.7334870232950073, -0.7004061657361658,
-                  -0.6479952272743823, -0.5926982319505321, -0.5599217222341915,
-                  -0.5480028121415452, -0.49964097922347867, -0.4203001696008753,
-                  -0.3344059204499378, -0.2745563922655548, -0.258755893042823,
-                  -0.21281014297392614, -0.1251037616154198, -0.02486081101616467,
-                  0.052711670524221245, 0.08156346618184351, 0.11018860125597137,
-                  0.184956394584215, 0.2769990401549583, 0.3533862259663847,
-                  0.39185061791993997, 0.4048279100244915, 0.4527903930933007,
-                  0.5181604219781556, 0.5745067111528425, 0.6066373424047252,
-                  0.6142607159164151, 0.6345436396519946, 0.6660807860559502,
-                  0.6916551627488451, 0.7043061657863463, 0.7074613359074287,
-                  0.7097605448059333, 0.7117980175789959, 0.7059983443925062,
-                  0.693373867220681, 0.6839039812030164]
-    expected_vt = [60.50323643979975, 60.46152282311776, 60.35195588066853,
-                   60.21533395043499, 60.09982168515988, 60.04066964102238,
-                   60.02053283714272, 59.94518544878992, 59.83955672767372,
-                   59.74465499035528, 59.688176306690224, 59.674434375999695,
-                   59.63705080297302, 59.575593022878486, 59.5200362867479,
-                   59.487302095351524, 59.477403317059846, 59.46882152619191,
-                   59.45242049432979, 59.44514092412334, 59.45163053025197,
-                   59.460024506000224, 59.463744883587445, 59.481935299569656,
-                   59.52060886409094, 59.572294471701106, 59.61340646988501,
-                   59.624904994314015, 59.6598130677176, 59.73238906214771,
-                   59.82488626148378, 59.9037999413929, 59.934901010623335,
-                   59.966754837133514, 60.05500712816102, 60.17514068478408,
-                   60.28669865440711, 60.34792638918577, 60.36947690125584,
-                   60.45364714009672, 60.58266887898263, 60.71293986868975,
-                   60.79930625245654, 60.82160071089043, 60.8853594969313,
-                   61.00319313188073, 61.132712888454186, 61.229832513831234,
-                   61.265349742432086, 61.30028735647212, 61.39021480819434,
-                   61.49843415833545, 61.586182165745534, 61.62960105016664,
-                   61.644121615044114, 61.69715263940207, 61.76740693027959,
-                   61.825107969209625, 61.8560293732873, 61.86304841080026,
-                   61.880900287439, 61.90495043959934, 61.91707450705104,
-                   61.91528849056989, 61.912066851775705, 61.907449920498564,
-                   61.88836178502154, 61.84936547495241, 61.801519377780714,
-                   61.77096806875504, 61.75953059854668, 61.711577658100396,
-                   61.62836639460284, 61.53297073619124, 61.46358406675648,
-                   61.44488197626709, 61.38959988569385, 61.280332330036074,
-                   61.14919211755058, 61.04273591939215, 61.00193140312018,
-                   60.96075236762326, 60.84965248313868, 60.70479794465818,
-                   60.57625670314503, 60.50800809804206]
+    expected_p = [0.4495141154262702, 0.44810731623424865, 0.4446924860674211,
+                  0.44101474518047623, 0.43843729172365287, 0.43732257545500436,
+                  0.4369771896163651, 0.43585040736153463, 0.43476981731494,
+                  0.43442274081030574, 0.43458934725870224, 0.43468321808979,
+                  0.4350672174585142, 0.4362345707586601, 0.4382356226789203,
+                  0.44027341593371017, 0.44114469018060415, 0.4420726748472857,
+                  0.4448144055192906, 0.44890293982902774, 0.45302566125041793,
+                  0.455412309803846, 0.45627258226390244, 0.45973185720339665,
+                  0.46533867318637956, 0.47137720587404724, 0.4755972251834676,
+                  0.4767157199899437, 0.479982869902296, 0.48630210267423224,
+                  0.4937119745209616, 0.499632228117403, 0.5018851956608419,
+                  0.5041516147429022, 0.5102382404520471, 0.5181350935202635,
+                  0.5251274962417639, 0.528840009619943, 0.5301266410196889,
+                  0.5350545811435541, 0.5423140663994042, 0.5492817851159084,
+                  0.5536930382468159, 0.5548035497737663, 0.5579127105543162,
+                  0.5633811248630228, 0.5689228569091898, 0.5726991310447442,
+                  0.5739863254025046, 0.5751980350089889, 0.5780379005475736,
+                  0.5808090106835145, 0.5823710691200007, 0.5828440926355674,
+                  0.5829481350770388, 0.5830503832825662, 0.5822895468763042,
+                  0.5804173591751022, 0.5785672859272517, 0.5780103965926895,
+                  0.5762363699105814, 0.572237324021336, 0.5667183507722322,
+                  0.5617561268531523, 0.559752846313555, 0.5576782694282233,
+                  0.5518326820337154, 0.543711872849146, 0.5360739378273777,
+                  0.5318633271279509, 0.530380132078173, 0.5245886109720006,
+                  0.5157429370998954, 0.5068932335599383, 0.5010911857600961,
+                  0.49960342886504683, 0.4953735336353052, 0.4876627291240921,
+                  0.4793665506564848, 0.4732825693615421, 0.4710889716397071,
+                  0.46894821577973017, 0.4635186380600756, 0.45714663551681506,
+                  0.4521183598666343, 0.4496787838015138]
+    expected_q = [0.6827418582524897, 0.6759389228503324, 0.6546242399054403,
+                  0.6207644558519532, 0.5852347231333254, 0.564287918118347,
+                  0.556689983477969, 0.5259471969259888, 0.4756862865761334,
+                  0.4212550001972507, 0.3831394924057264, 0.37303276222778314,
+                  0.343503487668719, 0.28638988969921425, 0.2194175020090074,
+                  0.1658803337183493, 0.14548865444145112, 0.1249612385669573,
+                  0.06972848884953677, -0.0022814642388705414, -0.06656443212989245,
+                  -0.10097001502687641, -0.11294895521598257, -0.15913617106073358,
+                  -0.22827644490625998, -0.2963224914648622, -0.3405683906933329,
+                  -0.3518831487295314, -0.3840082790031256, -0.44248719673145465,
+                  -0.5054051728451866, -0.551574215204323, -0.5682200554009741,
+                  -0.5844574843383253, -0.6255436186448009, -0.6733338973995404,
+                  -0.7102497022619719, -0.7276721907765155, -0.7333397406079434,
+                  -0.7531999175389356, -0.7766606922910467, -0.7916849353418282,
+                  -0.7966732798546298, -0.7972935019748215, -0.7974998065347064,
+                  -0.7914511420487961, -0.7742462932837018, -0.7531170349745917,
+                  -0.7434154933841022, -0.732749549391789, -0.6996770392831854,
+                  -0.6472925144162879, -0.5920338583314992, -0.5592837663247994,
+                  -0.5473759749629427, -0.49905814442640656, -0.419796532987056,
+                  -0.3339940482100009, -0.27421323483476967, -0.2584305361791927,
+                  -0.21253786937371122, -0.12493437607128674, -0.024809724391701034,
+                  0.05267445297594558, 0.08149394898045487, 0.11008483580643007,
+                  0.18477109572661538, 0.27671771664185124, 0.353033924201842,
+                  0.3914632816914362, 0.4044296067758364, 0.4523545483111222,
+                  0.5176797539993653, 0.5739945138569089, 0.6061138463121605,
+                  0.6137350794737616, 0.6340122577148257, 0.6655474631299673,
+                  0.6911272362100365, 0.7037896304781194, 0.7069493638375064,
+                  0.7092542139170925, 0.7113077755769187, 0.7055330024527745,
+                  0.692930113700812, 0.6834744521342498]
+    expected_vt = [60.503243776596165, 60.46153023170041, 60.35196357745784,
+                   60.21534169442452, 60.099829334123406, 60.040677146084676,
+                   60.02054030308586, 59.945192607070005, 59.83956332270126,
+                   59.74466073286034, 59.6881815170369, 59.67443939896856,
+                   59.63705527455284, 59.5755966280989, 59.52003881563317,
+                   59.487303746230786, 59.47740469832934, 59.46882266803735,
+                   59.45242085810586, 59.44514055782754, 59.45162966582038,
+                   59.460023392631314, 59.4637436896616, 59.48193396567828,
+                   59.520607349859056, 59.57229303072903, 59.6134051525281,
+                   59.62490374169246, 59.65981194145, 59.73238824637004,
+                   59.82488582955655, 59.90379982182214, 59.93490096692693,
+                   59.96675491891175, 60.05500747114224, 60.17514126096532,
+                   60.28669920214769, 60.34792696001425, 60.36947746684328,
+                   60.453647521646644, 60.582668964615706, 60.71293949062846,
+                   60.79930544278526, 60.82159978880084, 60.88535831265987,
+                   61.003191250737984, 61.13271036566074, 61.22982940090206,
+                   61.26534648299963, 61.300283935289386, 61.390210936392684,
+                   61.49842988504546, 61.58617772618817, 61.6295965023691,
+                   61.64411705987245, 61.697148185872855, 61.767402554884846,
+                   61.8251039647873, 61.85602560032788, 61.863044731750065,
+                   61.880896838249456, 61.90494747307992, 61.91707206994147,
+                   61.91528656108256, 61.91206502126355, 61.90744824569871,
+                   61.88836065205279, 61.84936481884471, 61.80151925202637,
+                   61.77096817847411, 61.759530818105134, 61.711578242982426,
+                   61.628367571043206, 61.532972492870755, 61.46358631318027,
+                   61.444884310802, 61.38960250589494, 61.280335702777364,
+                   61.149196249162785, 61.04274069753091, 61.0019364120191,
+                   60.960757611821634, 60.849658312940456, 60.704804495390036,
+                   60.57626380073712, 60.50801542345454]
 
     @testset "Braginskii electron IMEX timestepping" verbose=use_verbose begin
         println("Braginskii electron IMEX timestepping tests")
diff --git a/moment_kinetics/test/harrisonthompson.jl b/moment_kinetics/test/harrisonthompson.jl
index df0579405..06206a263 100644
--- a/moment_kinetics/test/harrisonthompson.jl
+++ b/moment_kinetics/test/harrisonthompson.jl
@@ -223,7 +223,11 @@ function run_test(test_input, analytic_rtol, analytic_atol, expected_phi,
                        rtol=10.0*analytic_rtol, atol=analytic_atol)
 
         # Regression test
-        @test isapprox(phi[:, end], expected_phi, rtol=regression_rtol, atol=regression_atol)
+        if expected_phi === nothing
+            println("values tested would be ", phi[:,end])
+        else
+            @test isapprox(phi[:, end], expected_phi, rtol=regression_rtol, atol=regression_atol)
+        end
     end
 end
 
@@ -242,44 +246,44 @@ function runtests()
         @testset "Chebyshev" begin
             test_input_chebyshev["base_directory"] = test_output_directory
             run_test(test_input_chebyshev, 3.e-2, 3.e-3,
-                     [-0.8270506701954182, -0.6647482038047513, -0.4359510242978734,
-                      -0.2930090318306279, -0.19789542580389763, -0.14560099254974576,
-                      -0.12410802135258239, -0.11657014257474364, -0.11761846656548933,
-                      -0.11657014257474377, -0.12410802135258239, -0.1456009925497464,
-                      -0.19789542580389616, -0.2930090318306262, -0.435951024297872,
-                      -0.66474820380475, -0.8270506701954171], 5.0e-9, 1.e-15)
+                     [-0.8270506736528097, -0.6647482045160528, -0.43595102198197894,
+                      -0.2930090302314022, -0.19789542449264944, -0.14560099229503182,
+                      -0.12410802088624982, -0.11657014266155726, -0.1176184662051167,
+                      -0.11657014266155688, -0.1241080208862487, -0.14560099229503298,
+                      -0.1978954244926481, -0.2930090302313995, -0.4359510219819795,
+                      -0.6647482045160534, -0.8270506736528144], 5.0e-9, 1.e-15)
         end
         @testset "Chebyshev split 1" begin
             test_input_chebyshev_split1["base_directory"] = test_output_directory
             run_test(test_input_chebyshev_split1, 3.e-2, 3.e-3,
-                     [-0.808956646073449, -0.6619131832543625, -0.4308291868843453,
-                      -0.295820339728472, -0.19344190006125275, -0.1492514208442407,
-                      -0.11977511930743077, -0.12060863604650167, -0.11342106824862994,
-                      -0.12060863604649866, -0.11977511930742626, -0.14925142084423915,
-                      -0.1934419000612479, -0.295820339728463, -0.4308291868843545,
-                      -0.6619131832543678, -0.808956646073442], 5.0e-9, 1.e-15)
+                     [-0.8089566460734486, -0.6619131832543634, -0.43082918688434424,
+                      -0.29582033972847016, -0.1934419000612522, -0.14925142084423915,
+                      -0.11977511930743077, -0.12060863604650167, -0.11342106824863019,
+                      -0.1206086360464999, -0.11977511930742751, -0.14925142084423915,
+                      -0.19344190006124898, -0.2958203397284666, -0.43082918688435656,
+                      -0.6619131832543697, -0.808956646073445], 5.0e-9, 1.e-15)
         end
         @testset "Chebyshev split 2" begin
             test_input_chebyshev_split2["base_directory"] = test_output_directory
-            run_test(test_input_chebyshev_split2, 5.e-2, 3.e-3,
-                     [-0.7667804422571606, -0.6128777083267765, -0.39031953439035494,
-                      -0.27326504140885904, -0.15311275955907663, -0.11567486122959246,
-                      -0.09226471519174786, -0.07477085120501512, -0.07206945428218994,
-                      -0.07477085120545898, -0.09226471518828984, -0.11567486123016281,
-                      -0.15311275955613904, -0.273265041412353, -0.3903195344134153,
-                      -0.612877708320375, -0.766780442235556], 5.0e-9, 1.e-15)
+            run_test(test_input_chebyshev_split2, 6.e-2, 3.e-3,
+                     [-0.7798736739831602, -0.661568214314525, -0.409872886370737,
+                      -0.24444487132869974, -0.17244646306807737, -0.11761557291772232,
+                      -0.09113439652298189, -0.09025928800454038, -0.08814925970784306,
+                      -0.09025928800449955, -0.0911343965228694, -0.1176155729185088,
+                      -0.1724464630676158, -0.24444487132881484, -0.40987288637069097,
+                      -0.6615682143148902, -0.7798736739849054], 5.0e-9, 1.e-15)
         end
         # The 'split 3' test is pretty badly resolved, but don't want to increase
         # run-time!
         @testset "Chebyshev split 3" begin
             test_input_chebyshev_split3["base_directory"] = test_output_directory
-            run_test(test_input_chebyshev_split3, 2.1e-1, 3.e-3,
-                     [-0.5535421015240105, -0.502816770781802, -0.3755477646148533,
-                      -0.24212761527100635, -0.15737450156025806, -0.11242832417550296,
-                      -0.09168434722655881, -0.08653015173768085, -0.0858195594227437,
-                      -0.08653015173768933, -0.09168434722650211, -0.11242832417546023,
-                      -0.15737450156026872, -0.24212761527101284, -0.3755477646149367,
-                      -0.5028167707818142, -0.5535421015238932], 5.0e-9, 1.e-15)
+            run_test(test_input_chebyshev_split3, 2.5e-1, 3.e-3,
+                     [-0.5012994554414933, -0.4624277373138882, -0.35356695432752266,
+                      -0.22371207174875177, -0.14096934539193717, -0.10082423314545275,
+                      -0.07938834260378662, -0.07480364283744717, -0.07316256734281283,
+                      -0.07480364283744836, -0.07938834260380849, -0.10082423314551169,
+                      -0.14096934539196504, -0.22371207174878788, -0.35356695432739504,
+                      -0.4624277373114037, -0.5012994554370094], 5.0e-9, 1.e-15)
         end
     end
 
diff --git a/moment_kinetics/test/recycling_fraction_tests.jl b/moment_kinetics/test/recycling_fraction_tests.jl
index 8320d62e5..fed259261 100644
--- a/moment_kinetics/test/recycling_fraction_tests.jl
+++ b/moment_kinetics/test/recycling_fraction_tests.jl
@@ -242,107 +242,110 @@ function runtests()
         @long @testset "Full-f" begin
             test_input["base_directory"] = test_output_directory
             run_test(test_input,
-                     [-0.05499288668923642, -0.017610447066356092, -0.0014497230450292054,
-                      0.0015713106015958053, 0.0021153221201727283, 0.00135154586425295,
-                      0.0013626547300678799, 0.003653592144195716, 0.00672151562009703,
-                      0.014857207950835708, 0.03452385151240508, 0.03591016289984108,
-                      0.02229102871737884, 0.007447997216451657, 0.00505099606227552,
-                      0.0016937650707449176, 0.0013469420674100871, 0.0016965410643657965,
-                      0.002562353505582182, -6.33366212813045e-5, -0.00969571716777773,
-                      -0.048688980279053266])
+                     [-0.05519530428547719, -0.017715187591731293, -0.0014094765667960986,
+                      0.0017408709303110715, 0.002364329303626605, 0.0015912944705669578,
+                      0.0015964146438650865, 0.003860702183595992, 0.0069126570648780075,
+                      0.01502802246799623, 0.034672817945651656, 0.03605806530524313,
+                      0.022451501088277735, 0.007636465002105951, 0.005249626321396431,
+                      0.0019202967839667788, 0.0015870957754252823, 0.0019420461064341924,
+                      0.0027769433764546388, 2.482219694607524e-5, -0.009668963817923988,
+                      -0.04888254078430223])
         end
         @long @testset "Split 1" begin
             test_input_split1["base_directory"] = test_output_directory
             run_test(test_input_split1,
-                     [-0.054793853738618496, -0.017535475032013862,
-                      -0.0014718402826481662, 0.0016368065803215382, 0.002097475822421603,
-                      0.001353447830403315, 0.001356138437924921, 0.0036537497347573,
-                      0.006719973928112565, 0.014855703760316889, 0.03452400419220982,
-                      0.03590889137214591, 0.022290971843531463, 0.007446918804323913,
-                      0.005048816472156039, 0.0016968661957691385, 0.0013266658105610114,
-                      0.0017028442360018413, 0.002534466861251151,
-                      -0.00018703865529355897, -0.009661145065079906,
-                      -0.0484483682752969])
+                     [-0.05499683305021959, -0.017648671323931234, -0.001435044896193346,
+                      0.0018073884147499259, 0.0023450881700708397, 0.0015955143089305631,
+                      0.001589815317774866, 0.003860937118209949, 0.006911057359417227,
+                      0.015026521129765527, 0.03467295602711963, 0.03605680131391841,
+                      0.022451426419324128, 0.007635385849475049, 0.005247408512658831,
+                      0.0019234631747149433, 0.001566853129890402, 0.001949947092934496,
+                      0.0027475042305556622, -0.00010906536252042205,
+                      -0.00962961346763738, -0.04864884428378774])
         end
         @long @testset "Split 2" begin
             test_input_split2["base_directory"] = test_output_directory
             run_test(test_input_split2,
-                     [-0.05555568198447252, -0.020145183717956348, 0.001182118478411508,
-                      0.002193148323751635, 0.0019441188563940751, 0.0011789368818662881,
-                      0.0013514249605048384, 0.003735531583031493, 0.006723696092974834,
-                      0.014826903180374499, 0.03454936277756109, 0.03587040875737859,
-                      0.022277731154827392, 0.007403052912240603, 0.00512153431160143,
-                      0.0017463637584066217, 0.0011452779397062784, 0.0014049872146431029,
-                      0.0022755389057580316, 0.0016780234234311344, -0.008381041468024259,
-                      -0.05005526194222513])
+                     [-0.05584608895693617, -0.020285311011810747, 0.0013246162852513857,
+                      0.002408198079080654, 0.002193404660625343, 0.0014165984310586497,
+                      0.0015838582299817442, 0.003942456292519074, 0.006915806487555444,
+                      0.014996822639406769, 0.034698460163972725, 0.03601812331030096,
+                      0.022438422037486003, 0.007592137358805067, 0.00532114704663586,
+                      0.001973382574270663, 0.0013804707387267182, 0.0016443777257862315,
+                      0.0025134094388913966, 0.0018832456170377893, -0.008404304571565805,
+                      -0.05034925353831177])
         end
         @long @testset "Split 3" begin
             test_input_split3["base_directory"] = test_output_directory
             run_test(test_input_split3,
-                     [-0.036205375991650725, -0.030483334021285433, -0.028961568619094404,
-                      -0.028550383934166465, -0.02551672335720456, -0.021976119708577647,
-                      -0.01982001937014411, -0.01331564927923702, -0.00984100255121529,
-                      -0.005266490060020825, 0.0021494114844098316, 0.004620296275317165,
-                      0.011509404776589328, 0.01757377252325957, 0.02014859036576961,
-                      0.027122126647315926, 0.030505809525427197, 0.034043759795000156,
-                      0.03932240322253646, 0.04089804092628224, 0.04436256082283185,
-                      0.04582461258085377, 0.0457025256980273, 0.04312136181903663,
-                      0.04054653135540802, 0.03787132328029428, 0.03223719811392133,
-                      0.030105212408583878, 0.024827994199332723, 0.018595982530248478,
-                      0.016209527148134187, 0.008754940562653064, 0.0040079860524162405,
-                      3.89264740137833e-5, -0.007642430261913982, -0.010258137085572222,
-                      -0.015541799469166076, -0.021192018291797773, -0.022784703489569562,
-                      -0.026873219344096318, -0.028749404798656616, -0.029220744790456707,
-                      -0.032303083015072])
+                     [-0.03620705983288495, -0.030483526856225397, -0.028960441350906176,
+                      -0.028549218778503995, -0.025515599608030678, -0.021975115062752498,
+                      -0.019818943505064867, -0.013314608790987136, -0.009839994543852062,
+                      -0.005265524793578627, 0.002150328580191541, 0.004621192779726743,
+                      0.011510249894025814, 0.017574569324021832, 0.020149366401796907,
+                      0.027122843852491103, 0.03050649889203747, 0.03404441833358536,
+                      0.039323018405068834, 0.04089864462026069, 0.04436314083820065,
+                      0.04582518382395237, 0.045703097564838854, 0.04312195015009901,
+                      0.04054713854267327, 0.0378719503058148, 0.03223787080558438,
+                      0.030105904373564214, 0.024828730387096765, 0.01859677066598083,
+                      0.01621033424329937, 0.008755805694319756, 0.004008885194932725,
+                      3.98551863685712e-5, -0.007641449438487893, -0.010257112218851392,
+                      -0.01554078023638837, -0.021190999288843403, -0.022783488732253034,
+                      -0.026872207781674047, -0.028748355604856834, -0.02921966520151288,
+                      -0.03230397503173283])
         end
 
+        fullf_expected_output = [-0.04413552960518834, -0.021828861958540703,
+                                 -0.012581752434237016, -0.010579192578141765,
+                                 -0.0068512759586887885, 2.5839426347419376e-5,
+                                 0.006155432086970659, 0.011968120005188723,
+                                 0.019681941313251225, 0.029904657995405693,
+                                 0.03949771002617614, 0.04255717403165738,
+                                 0.031939444000420925, 0.02124792913154375,
+                                 0.015253505295222548, 0.00930078535463162,
+                                 0.0012208189549702839, -0.005751247690841759,
+                                 -0.009217769053947429, -0.011407780388153488,
+                                 -0.020596064125157174, -0.03809484046514718]
         @testset "Adaptive timestep - full-f" begin
             test_input_adaptive["base_directory"] = test_output_directory
             run_test(test_input_adaptive,
-                     [-0.04406302421567803, -0.021804698228788805, -0.012620074187743302,
-                      -0.0106380090360855, -0.00691525568459642, -4.233879648051771e-5,
-                      0.006087337143759688, 0.011901298104136846, 0.019618466419394923,
-                      0.029848931238319366, 0.03944995729484389, 0.0425118861432023,
-                      0.031885423069787075, 0.02118545257675995, 0.015187813894182393,
-                      0.009233281236463076, 0.0011524769858418965, -0.005816337799653475,
-                      -0.009278801135132173, -0.011463100731406845, -0.020577143329947718,
-                      -0.03801740811142216], rtol=6.0e-4, atol=2.0e-12)
+                     fullf_expected_output, rtol=6.0e-4, atol=2.0e-12)
         end
         @testset "Adaptive timestep - split 1" begin
             test_input_adaptive_split1["base_directory"] = test_output_directory
             run_test(test_input_adaptive_split1,
-                     [-0.04404229340935816, -0.021715993643694146, -0.01254532739530461,
-                      -0.010653532207501667, -0.006915890693912259, -4.387439665774011e-5,
-                      0.006085426633745894, 0.01190000270341187, 0.019617011480099204,
-                      0.02984682106369964, 0.03944638388660516, 0.042512679055667614,
-                      0.031887123203488764, 0.02118278590028187, 0.01518621987364583,
-                      0.009231972103916531, 0.001150127914745935, -0.005818837808004554,
-                      -0.009271945665225503, -0.011554150887267762, -0.020405808046592103,
-                      -0.03801003686546914], rtol=6.0e-4, atol=2.0e-12)
+                     [-0.04411241514173686, -0.02173832960900953, -0.012507206800103619,
+                      -0.010594605015685111, -0.006851937496868262, 2.4307749760781248e-5,
+                      0.006153573001701797, 0.011966908336801292, 0.019680600194907028,
+                      0.02990265026906549, 0.039494202889752195, 0.04255801055036077,
+                      0.031941240640856954, 0.02124538117173497, 0.0152520119972139,
+                      0.009299543598122585, 0.001218486350949803, -0.005753814631808573,
+                      -0.009211138614098327, -0.011500666056267622, -0.020424831739003606,
+                      -0.03808534246490289], rtol=6.0e-4, atol=2.0e-12)
         end
         @testset "Adaptive timestep - split 2" begin
             test_input_adaptive_split2["base_directory"] = test_output_directory
             run_test(test_input_adaptive_split2,
-                     [-0.04459179157659319, -0.022746645577738103, -0.012946338116940815,
-                      -0.010871883552302172, -0.006982172600237595, -6.051063734515426e-5,
-                      0.006026142648981889, 0.0119738448881976, 0.01962536485263761,
-                      0.02987382918222712, 0.03946868382103347, 0.0424964960343025,
-                      0.031848821225964274, 0.02115083095729825, 0.015157115049355067,
-                      0.00916355205311941, 0.0011734398171854166, -0.005841676293891676,
-                      -0.009106806670666958, -0.011016868962010862, -0.02091990269514529,
-                      -0.037555410636391916] , rtol=6.0e-4, atol=2.0e-12)
+                     [-0.044658773996679106, -0.022770640350132876, -0.01291279676887995,
+                      -0.010818472056813256, -0.00692137979236985, 8.260915374129437e-6,
+                      0.006095505380954945, 0.012043966021961394, 0.01969312249006842,
+                      0.02993329149668162, 0.03951863202813308, 0.04254329784045647,
+                      0.031905905757383245, 0.0212173464030042, 0.015225863469416798,
+                      0.00923409202948059, 0.0012431576067072942, -0.005777971895837924,
+                      -0.009047333684784373, -0.010964221005155143, -0.020937032832434074,
+                      -0.03762542957657465], rtol=6.0e-4, atol=2.0e-12)
         end
         @testset "Adaptive timestep - split 3" begin
             test_input_adaptive_split3["base_directory"] = test_output_directory
             run_test(test_input_adaptive_split3,
-                     [-0.03462989367404872, -0.032011023958736944, -0.027146091249774885,
-                      -0.020930812236514006, -0.010156481871960516, 0.0027832606290666053,
-                      0.012831753103964484, 0.022090130790633454, 0.03302848792727571,
-                      0.04152538920797896, 0.045375641995415654, 0.046239778021746426,
-                      0.04254552381542057, 0.0348087945161905, 0.02707439699253448,
-                      0.017880292949104492, 0.0047783374339709405, -0.007768102112509674,
-                      -0.016299183329203885, -0.024140030225825285, -0.031567878963645234,
-                      -0.03417555449482003], rtol=6.0e-4, atol=2.0e-12)
+                     [-0.03462849854420647, -0.03201042148267427, -0.027145441768761876,
+                      -0.020930163084065857, -0.010155888649583823, 0.0027838306741690046,
+                      0.012832285048924648, 0.022090624745954916, 0.033028935673319444,
+                      0.041525790733438206, 0.04537601028144605, 0.046240136469284064,
+                      0.042545918106145095, 0.03480923464250141, 0.02707487186844719,
+                      0.017880803893474347, 0.004778898480075153, -0.007767501511791003,
+                      -0.016298557311046468, -0.0241393991660673, -0.03156718517969005,
+                      -0.03417441352897386], rtol=6.0e-4, atol=2.0e-12)
         end
 
         @long @testset "Check other timestep - $type" for
@@ -354,14 +357,7 @@ function runtests()
             timestep_check_input["run_name"] = type
             timestep_check_input["timestepping"]["type"] = type
             run_test(timestep_check_input,
-                     [-0.04406302421567803, -0.021804698228788805, -0.012620074187743302,
-                      -0.0106380090360855, -0.00691525568459642, -4.233879648051771e-5,
-                      0.006087337143759688, 0.011901298104136846, 0.019618466419394923,
-                      0.029848931238319366, 0.03944995729484389, 0.0425118861432023,
-                      0.031885423069787075, 0.02118545257675995, 0.015187813894182393,
-                      0.009233281236463076, 0.0011524769858418965, -0.005816337799653475,
-                      -0.009278801135132173, -0.011463100731406845, -0.020577143329947718,
-                      -0.03801740811142216], rtol=8.e-4, atol=1.e-10)
+                     fullf_expected_output, rtol=8.e-4, atol=1.e-10)
         end
     end
 
diff --git a/moment_kinetics/test/wall_bc_tests.jl b/moment_kinetics/test/wall_bc_tests.jl
index c1c40634d..ab09fec48 100644
--- a/moment_kinetics/test/wall_bc_tests.jl
+++ b/moment_kinetics/test/wall_bc_tests.jl
@@ -246,30 +246,30 @@ function runtests()
         @testset "Chebyshev uniform" begin
             test_input_chebyshev["base_directory"] = test_output_directory
             run_test(test_input_chebyshev,
-                     [-1.1689445031600718, -0.7479504438063098, -0.6947559936893813,
-                      -0.6917252442591313, -0.7180152498764835, -0.9980114095597415],
+                     [-1.168944495073113, -0.747950464799219, -0.6947560093910274,
+                      -0.6917252594440765, -0.7180152693147238, -0.9980114030684668],
                      2.e-3)
         end
         
         @testset "Chebyshev sqrt grid odd" begin
             test_input_chebyshev_sqrt_grid_odd["base_directory"] = test_output_directory
             run_test(test_input_chebyshev_sqrt_grid_odd,
-                     [-1.2047298885671576, -0.9431378294506091, -0.8084332392927167,
-                     -0.7812620422650213, -0.7233303514000929, -0.7003878610612269,
-                     -0.69572751349158, -0.6933148921301019, -0.6992503992521327,
-                     -0.7115787972775218, -0.7596015032228407, -0.795776514029509,
-                     -0.876303297135126, -1.1471244425913258],
+                     [-1.2047298844053338, -0.9431378244038217, -0.8084332486925859,
+                      -0.7812620574297168, -0.7233303715713063, -0.700387877851292,
+                      -0.695727529425101, -0.6933149075958859, -0.6992504158371133,
+                      -0.7115788158947632, -0.7596015227027635, -0.7957765261319207,
+                      -0.876303296785542, -1.1471244373220089],
                      2.e-3)
         end
         @testset "Chebyshev sqrt grid even" begin
             test_input_chebyshev_sqrt_grid_even["base_directory"] = test_output_directory
             run_test(test_input_chebyshev_sqrt_grid_even,
-                     [-1.213617049279473, -1.0054529928344382, -0.871444761913497,
-                     -0.836017699317097, -0.7552110924643832, -0.7264644073096705,
-                     -0.7149147366621806, -0.6950077192395091, -0.6923364889119271,
-                     -0.6950077192395089, -0.7149147366621814, -0.7264644073096692,
-                     -0.7552110924643836, -0.8360176993170979, -0.8714447619134948,
-                     -1.0054529928344376, -1.2136170492794727],
+                     [-1.213617044609117, -1.0054529856551995, -0.8714447622540997,
+                      -0.836017704148175, -0.7552111126205924, -0.7264644278204795,
+                      -0.7149147557607726, -0.6950077350352664, -0.6923365041825125,
+                      -0.6950077350352668, -0.7149147557607729, -0.7264644278204795,
+                      -0.7552111126205917, -0.8360177041481733, -0.8714447622540994,
+                      -1.0054529856552, -1.213617044609118],
                      2.e-3)
         end
     end

From f55623cd55ba4fdc710aac14a8625e1bdf6bfca8 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 31 Jul 2024 23:48:08 +0100
Subject: [PATCH 416/428] Update examples CI job to handle kinetic electron
 inputs

---
 .github/workflows/examples.yml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml
index 405ffcb9e..aaf329ac6 100644
--- a/.github/workflows/examples.yml
+++ b/.github/workflows/examples.yml
@@ -23,4 +23,8 @@ jobs:
           touch Project.toml
           julia -O3 --project -e 'import Pkg; Pkg.develop(path="moment_kinetics/"); Pkg.add("NCDatasets"); Pkg.precompile()'
           # Reduce nstep for each example to 10 to avoid the CI job taking too long
-          julia -O3 --project -e 'using moment_kinetics; for (root, dirs, files) in walkdir("examples") for file in files if endswith(file, ".toml") filename = joinpath(root, file); println(filename); input = moment_kinetics.moment_kinetics_input.read_input_file(filename); t_input = get(input, "timestepping", Dict{String,Any}()); t_input["nstep"] = 10; t_input["dt"] = 1.0e-10; pop!(input, "z_nelement_local", ""); pop!(input, "r_nelement_local", ""); run_moment_kinetics(input) end end end'
+          # Note we skip the example `if (occursin("ARK", get(t_input, "type", "") && Sys.isapple())`
+          # because the way we use MINPACK.jl (needed for nonlinear solvers
+          # used for implicit parts of timestep) doesn't currently work on
+          # macOS.
+          julia -O3 --project -e 'using moment_kinetics; for (root, dirs, files) in walkdir("examples") for file in files if endswith(file, ".toml") filename = joinpath(root, file); println(filename); input = moment_kinetics.moment_kinetics_input.read_input_file(filename); t_input = get(input, "timestepping", Dict{String,Any}()); if (occursin("ARK", get(t_input, "type", "")) && Sys.isapple()) continue end; t_input["nstep"] = 10; t_input["dt"] = 1.0e-12; input["timestepping"] = t_input; pop!(input, "z_nelement_local", ""); pop!(input, "r_nelement_local", ""); electron_t_input = get(input, "electron_timestepping", Dict{String,Any}()); electron_t_input["initialization_residual_value"] = 1.0e8; electron_t_input["converged_residual_value"] = 1.0e8; input["electron_timestepping"] = electron_t_input; nl_solver_input = get(input, "nonlinear_solver", Dict{String,Any}()); nl_solver_input["rtol"] = 1.0e6; nl_solver_input["atol"] = 1.0e6; input["nonlinear_solver"] = nl_solver_input; run_moment_kinetics(input) end end end'

From 4bd1daaaff8c39c164a984d1e458d79082d36af5 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Aug 2024 10:03:28 +0100
Subject: [PATCH 417/428] Remove debug print statements

---
 moment_kinetics/src/initial_conditions.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index acfd10e35..63763147e 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -870,12 +870,10 @@ end
 """
 function init_density!(dens, z, r, spec, n_species)
     for is ∈ 1:n_species
-        println("init_option: ", spec[is].z_IC.initialization_option)
         for ir ∈ 1:r.n
             if spec[is].z_IC.initialization_option == "gaussian"
                 # initial condition is an unshifted Gaussian
                 @. dens[:,ir,is] = spec[is].initial_density + exp(-(z.grid/spec[is].z_IC.width)^2)
-                println("ion_dens: ", dens[1,1,is], " init: ", spec[is].initial_density, " sum_factor: ", exp(-(z.grid[1]/spec[is].z_IC.width)^2))
             elseif spec[is].z_IC.initialization_option == "sinusoid"
                 # initial condition is sinusoid in z
                 @. dens[:,ir,is] =

From 705642bfd8c4f21b1cae331a8eb61d7664e32185 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Aug 2024 11:43:56 +0100
Subject: [PATCH 418/428] Skip Braginskii-IMEX test on macOS because MINPACK.jl
 is broken

---
 .../test/braginskii_electrons_imex_tests.jl   | 29 +++++++++++--------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/moment_kinetics/test/braginskii_electrons_imex_tests.jl b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
index f7a67165d..90a70f505 100644
--- a/moment_kinetics/test/braginskii_electrons_imex_tests.jl
+++ b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
@@ -282,19 +282,24 @@ function runtests()
     @testset "Braginskii electron IMEX timestepping" verbose=use_verbose begin
         println("Braginskii electron IMEX timestepping tests")
 
-        @testset "Split 3" begin
-            test_input["base_directory"] = test_output_directory
-            run_test(test_input, expected_p, expected_q, expected_vt)
-        end
-        @long @testset "Check other timestep - $type" for
-                type ∈ ("KennedyCarpenterARK437",)
+        if Sys.isapple()
+            @testset_skip "MINPACK is broken on macOS (https://github.com/sglyon/MINPACK.jl/issues/18)" "non-linear solvers" begin
+            end
+        else
+            @testset "Split 3" begin
+                test_input["base_directory"] = test_output_directory
+                run_test(test_input, expected_p, expected_q, expected_vt)
+            end
+            @long @testset "Check other timestep - $type" for
+                    type ∈ ("KennedyCarpenterARK437",)
 
-            timestep_check_input = deepcopy(test_input)
-            timestep_check_input["base_directory"] = test_output_directory
-            timestep_check_input["run_name"] = type
-            timestep_check_input["timestepping"]["type"] = type
-            run_test(timestep_check_input, expected_p, expected_q, expected_vt,
-                     rtol=2.e-4, atol=1.e-10)
+                timestep_check_input = deepcopy(test_input)
+                timestep_check_input["base_directory"] = test_output_directory
+                timestep_check_input["run_name"] = type
+                timestep_check_input["timestepping"]["type"] = type
+                run_test(timestep_check_input, expected_p, expected_q, expected_vt,
+                         rtol=2.e-4, atol=1.e-10)
+            end
         end
     end
 

From b7e2b6d3cb47eaccd944a8fb8c91fd8737a5a793 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Aug 2024 12:20:56 +0100
Subject: [PATCH 419/428] Remove incorrect, unused line from
 parallel_delta_x_calc_r_z_vperp_vpa()

---
 moment_kinetics/src/nonlinear_solvers.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/moment_kinetics/src/nonlinear_solvers.jl b/moment_kinetics/src/nonlinear_solvers.jl
index f35f0ade0..365b91295 100644
--- a/moment_kinetics/src/nonlinear_solvers.jl
+++ b/moment_kinetics/src/nonlinear_solvers.jl
@@ -984,7 +984,6 @@ function parallel_delta_x_calc_r_z_vperp_vpa(delta_x::Tuple{AbstractArray{mk_flo
 
     delta_x_ppar, delta_x_pdf = delta_x
     V_ppar, V_pdf = V
-    y_ppar, y_pdf = y
 
     ny = length(y)
 

From 9f460cb053202c5de2fca325b300fda413928f29 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Aug 2024 14:09:43 +0100
Subject: [PATCH 420/428] Fix Braginskii electrons test to work with
 distributed MPI

---
 .../test/braginskii_electrons_imex_tests.jl   | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/moment_kinetics/test/braginskii_electrons_imex_tests.jl b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
index 90a70f505..f784e0cf2 100644
--- a/moment_kinetics/test/braginskii_electrons_imex_tests.jl
+++ b/moment_kinetics/test/braginskii_electrons_imex_tests.jl
@@ -12,8 +12,7 @@ using MPI
 using moment_kinetics.coordinates: define_coordinate
 using moment_kinetics.input_structs: grid_input, advection_input
 using moment_kinetics.interpolation: interpolate_to_grid_z
-using moment_kinetics.load_data: open_readonly_output_file
-using moment_kinetics.load_data: load_electron_moments_data
+using moment_kinetics.load_data: get_run_info_no_setup, close_run_info, get_variable
 
 # default inputs for tests
 test_input = Dict("n_ion_species" => 1,
@@ -143,17 +142,17 @@ function run_test(test_input, expected_p, expected_q, expected_vt; rtol=1.e-6,
             # Load and analyse output
             #########################
 
-            path = joinpath(realpath(input["base_directory"]), name, name)
+            path = joinpath(realpath(input["base_directory"]), name)
 
-            # open the netcdf file and give it the handle 'fid'
-            fid = open_readonly_output_file(path,"moments")
+            # open the output file
+            run_info = get_run_info_no_setup(path)
 
-            # load fields data
-            parallel_pressure_zrt, parallel_heat_flux_zrt, thermal_speed_zrt =
-            load_electron_moments_data(fid)
+            parallel_pressure_zrt = get_variable(run_info, "electron_parallel_pressure")
+            parallel_heat_flux_zrt = get_variable(run_info, "electron_parallel_heat_flux")
+            thermal_speed_zrt = get_variable(run_info, "electron_thermal_speed")
+
+            close_run_info(run_info)
 
-            close(fid)
-            
             p = parallel_pressure_zrt[:,1,:]
             q = parallel_heat_flux_zrt[:,1,:]
             vt = thermal_speed_zrt[:,1,:]

From 7684680b94264797e1d79ae6e20a93f9332efa85 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 1 Aug 2024 15:40:51 +0100
Subject: [PATCH 421/428] Fix loading of diagnostic variables when not using
 parallel_io

---
 moment_kinetics/src/load_data.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 29c6d38c4..d4e60000e 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3943,7 +3943,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
             fid = open_readonly_output_file(run_info.files[1], run_info.ext, iblock=0)
             group = get_group(fid, "dynamic_data")
             result = load_variable(group, variable_name)
-            result = selectdim(result, ndims(result), global_it_start:global_it_end)
+            result = selectdim(result, ndims(result), it)
         elseif nd == 3
             result = allocate_float(run_info.z.n, run_info.r.n, run_info.nt)
             read_distributed_zr_data!(result, variable_name, run_info.files,

From 550257124434eb1750d1f2b3eca7e449b0d09d50 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Sat, 3 Aug 2024 00:22:47 +0100
Subject: [PATCH 422/428] Fix kinetic electron options in precompile run

kinetic electrons have to be moment-kinetic. Currently this requires a
1D1V run.
---
 util/precompile_run.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/util/precompile_run.jl b/util/precompile_run.jl
index 83c60f73c..92b5215d5 100644
--- a/util/precompile_run.jl
+++ b/util/precompile_run.jl
@@ -79,6 +79,14 @@ geo_input1 = merge(wall_bc_cheb_input, Dict("n_neutral_species" => 0,
 kinetic_electron_input = merge(cheb_input, Dict("evolve_moments_density" => true,
                                                 "evolve_moments_parallel_flow" => true,
                                                 "evolve_moments_parallel_pressure" => true,
+                                                "r_ngrid" => 1,
+                                                "r_nelement" => 1,
+                                                "vperp_ngrid" => 1,
+                                                "vperp_nelement" => 1,
+                                                "vzeta_ngrid" => 1,
+                                                "vzeta_nelement" => 1,
+                                                "vr_ngrid" => 1,
+                                                "vr_nelement" => 1,
                                                 "electron_physics" => "kinetic_electrons",
                                                 "electron_timestepping" => Dict{String,Any}("nstep" => 1,
                                                                                             "dt" => 2.0e-11,

From 2d2156eb5ece17ad94e6958325eaac16c6790e6f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 2 Aug 2024 23:18:07 +0100
Subject: [PATCH 423/428] Fix recycling_fraction debug test for new neutral
 initialisation

The normalisation introduced in "Normalise grid to initial_temperature
while initialising f_neutral" (ac7f0a35) resulted in an
interpolation that caused negative pressure (presumably due to the very
low resolution). Tweak `T_wall` to avoid this, by setting it to a value
that means no interpolation is needed.
---
 moment_kinetics/debug_test/recycling_fraction_inputs.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/debug_test/recycling_fraction_inputs.jl b/moment_kinetics/debug_test/recycling_fraction_inputs.jl
index 643b7c954..7f4b24266 100644
--- a/moment_kinetics/debug_test/recycling_fraction_inputs.jl
+++ b/moment_kinetics/debug_test/recycling_fraction_inputs.jl
@@ -12,7 +12,7 @@ test_input = Dict("n_ion_species" => 1,
                   "recycling_fraction" => 0.5,
                   "krook_collisions" => Dict{String,Any}("use_krook" => true),
                   "T_e" => 0.2,
-                  "T_wall" => 0.1,
+                  "T_wall" => 2.0,
                   "initial_density1" => 1.0,
                   "initial_temperature1" => 1.0,
                   "z_IC_option1" => "gaussian",

From df1c97f218044879b9394693348150b90094a10f Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 2 Aug 2024 22:10:04 +0100
Subject: [PATCH 424/428] Compile system image to speed up 'debug_checks' CI
 job

---
 .github/workflows/debug_checks.yml | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/debug_checks.yml b/.github/workflows/debug_checks.yml
index 468547737..26a04ddc7 100644
--- a/.github/workflows/debug_checks.yml
+++ b/.github/workflows/debug_checks.yml
@@ -23,19 +23,17 @@ jobs:
           version: '1.10'
           arch: x64
       - uses: julia-actions/cache@v1
-      - uses: julia-actions/julia-buildpkg@v1
-        with:
-          project: 'moment_kinetics/'
       - name: Debug test
         run: |
-          cd moment_kinetics
-
           # Hard code the debug level so that we can run without using the
           # `--compiled-modules=no` flag, which breaks Symbolics.jl at the
           # moment.
-          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" src/debugging.jl
+          sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" moment_kinetics/src/debugging.jl
 
-          julia --project -e 'using MPIPreferences; MPIPreferences.use_system_binary(); using Pkg; Pkg.precompile()'
+          touch Project.toml
+          julia --project -O3 --check-bounds=yes -e 'using Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_system_binary()'
+          julia --project -O3 --check-bounds=yes -e 'using Pkg; Pkg.develop(path="moment_kinetics/"); Pkg.precompile()'
+          julia --project -O3 --check-bounds=yes precompile.jl --debug 2
 
           # Need to use openmpi so that the following arguments work:
           # * `--mca rmaps_base_oversubscribe 1` allows oversubscription (more processes
@@ -43,5 +41,5 @@ jobs:
           # * `--mca mpi_yield_when_idle 1` changes a setting to prevent excessively
           #   terrible performance when oversubscribing.
           ## Don't use --compiled-modules=no for now, as it currently breaks Symbolics.jl
-          #mpiexec -np 4 --mca rmaps_base_oversubscribe 1 julia --project --check-bounds=yes --compiled-modules=no debug_test/sound_wave_tests.jl --debug 2
-          mpiexec -np 4 --mca rmaps_base_oversubscribe 1 julia --project --check-bounds=yes debug_test/runtests.jl --debug 2
+          #mpiexec -np 4 --mca rmaps_base_oversubscribe 1 julia --project --check-bounds=yes --compiled-modules=no moment_kinetics/debug_test/sound_wave_tests.jl --debug 2
+          mpiexec -np 4 --mca rmaps_base_oversubscribe 1 julia --project -Jmoment_kinetics.so -O3 --check-bounds=yes moment_kinetics/debug_test/runtests.jl --debug 2

From 0ed436918e7a99263f7deda014f53d0b0e7a9c54 Mon Sep 17 00:00:00 2001
From: Michael Hardman <29800382+mrhardman@users.noreply.github.com>
Date: Tue, 6 Aug 2024 12:42:59 +0100
Subject: [PATCH 425/428] Comment out mms_tests.jl from debug_test/runtests.jl
 due to missing Symbolics package.

---
 moment_kinetics/debug_test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/debug_test/runtests.jl b/moment_kinetics/debug_test/runtests.jl
index 94b5dbbc7..7317ee974 100644
--- a/moment_kinetics/debug_test/runtests.jl
+++ b/moment_kinetics/debug_test/runtests.jl
@@ -8,7 +8,7 @@ function runtests()
         include(joinpath(@__DIR__, "fokker_planck_collisions_tests.jl"))
         include(joinpath(@__DIR__, "wall_bc_tests.jl"))
         include(joinpath(@__DIR__, "harrisonthompson.jl"))
-        include(joinpath(@__DIR__, "mms_tests.jl"))
+        #include(joinpath(@__DIR__, "mms_tests.jl"))
         include(joinpath(@__DIR__, "restart_interpolation_tests.jl"))
         include(joinpath(@__DIR__, "recycling_fraction_tests.jl"))
         include(joinpath(@__DIR__, "gyroaverage_tests.jl"))

From 1816d442fe4e7f788ed5bd4cd0f06dcaecafa937 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Aug 2024 15:36:19 +0100
Subject: [PATCH 426/428] Fix loading of electron moment variables from
 'distributed' output files

---
 moment_kinetics/src/load_data.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index d4e60000e..f4c887b84 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -3922,7 +3922,7 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     else
         # Use existing distributed I/O loading functions
         diagnostic_variable = false
-        if variable_name ∈ em_variables
+        if variable_name ∈ tuple(em_variables..., electron_moment_variables...)
             nd = 3
         elseif variable_name ∈ electron_dfn_variables
             nd = 5

From 7ffae2cdf325085d94033a198eb615e54de161d1 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Aug 2024 15:37:31 +0100
Subject: [PATCH 427/428] Install `Symbolics` in 'debug checks' CI job

This addition will allow the MMS test to work properly (once the test
input is updated).
---
 .github/workflows/debug_checks.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/debug_checks.yml b/.github/workflows/debug_checks.yml
index 26a04ddc7..55eef9e63 100644
--- a/.github/workflows/debug_checks.yml
+++ b/.github/workflows/debug_checks.yml
@@ -31,7 +31,7 @@ jobs:
           sed -i -e "s/_debug_level = get_options.*/_debug_level = 2/" moment_kinetics/src/debugging.jl
 
           touch Project.toml
-          julia --project -O3 --check-bounds=yes -e 'using Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler"]); using MPIPreferences; MPIPreferences.use_system_binary()'
+          julia --project -O3 --check-bounds=yes -e 'using Pkg; Pkg.add(["MPI", "MPIPreferences", "PackageCompiler", "Symbolics"]); using MPIPreferences; MPIPreferences.use_system_binary()'
           julia --project -O3 --check-bounds=yes -e 'using Pkg; Pkg.develop(path="moment_kinetics/"); Pkg.precompile()'
           julia --project -O3 --check-bounds=yes precompile.jl --debug 2
 

From a9c7c59ec73bfa62e89e1fb7a2d486d7249a221b Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Wed, 7 Aug 2024 16:49:30 +0100
Subject: [PATCH 428/428] Include kinetic electron tests when running all debug
 checks

Also fix inputs for the kinetic electron debug check.
---
 .../debug_test/kinetic_electron_inputs.jl        | 16 ++++++++--------
 moment_kinetics/debug_test/runtests.jl           |  1 +
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/debug_test/kinetic_electron_inputs.jl b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
index 697d95ee7..684b114ff 100644
--- a/moment_kinetics/debug_test/kinetic_electron_inputs.jl
+++ b/moment_kinetics/debug_test/kinetic_electron_inputs.jl
@@ -10,7 +10,6 @@ test_input = Dict("n_ion_species" => 1,
                   "evolve_moments_parallel_pressure" => true,
                   "evolve_moments_conservation" => true,
                   "recycling_fraction" => 0.5,
-                  "krook_collisions" => true,
                   "T_e" => 0.2,
                   "T_wall" => 0.1,
                   "initial_density1" => 1.0,
@@ -73,20 +72,21 @@ test_input = Dict("n_ion_species" => 1,
                   "z_discretization" => "chebyshev_pseudospectral",
                   "z_element_spacing_option" => "sqrt",
                   "vpa_ngrid" => 3,
-                  "vpa_nelement" => 4,
+                  "vpa_nelement" => 16,
                   "vpa_L" => 6.0,
                   "vpa_bc" => "zero",
                   "vpa_discretization" => "chebyshev_pseudospectral",
                   "vz_ngrid" => 3,
-                  "vz_nelement" => 4,
+                  "vz_nelement" => 6,
                   "vz_L" => 6.0,
                   "vz_bc" => "zero",
                   "vz_discretization" => "chebyshev_pseudospectral",
-                  "ion_source" => Dict("active" => true,
-                                       "z_profile" => "gaussian",
-                                       "z_width" => 0.125,
-                                       "source_strength" => 2.0,
-                                       "source_T" => 2.0),
+                  "ion_source" => Dict{String,Any}("active" => true,
+                                                   "z_profile" => "gaussian",
+                                                   "z_width" => 0.125,
+                                                   "source_strength" => 2.0,
+                                                   "source_T" => 2.0),
+                  "krook_collisions" => Dict{String,Any}("use_krook" => true),
                   "numerical_dissipation" => Dict{String,Any}("force_minimum_pdf_value" => 0.0,
                                                               "vpa_dissipation_coefficient" => 1e-2))
 
diff --git a/moment_kinetics/debug_test/runtests.jl b/moment_kinetics/debug_test/runtests.jl
index 94b5dbbc7..2efa563ad 100644
--- a/moment_kinetics/debug_test/runtests.jl
+++ b/moment_kinetics/debug_test/runtests.jl
@@ -4,6 +4,7 @@ include("setup.jl")
 
 function runtests()
     @testset "moment_kinetics tests" begin
+        include(joinpath(@__DIR__, "kinetic_electron_tests.jl"))
         include(joinpath(@__DIR__, "sound_wave_tests.jl"))
         include(joinpath(@__DIR__, "fokker_planck_collisions_tests.jl"))
         include(joinpath(@__DIR__, "wall_bc_tests.jl"))