From e6676ef07961aa3d45c057d6e69adf9456dccee4 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Wed, 18 Sep 2024 10:11:10 +0100
Subject: [PATCH 01/31] Update naming of update_qpar for ions to be
 update_ion_qpar, in line with update_neutral_qz

---
 moment_kinetics/src/initial_conditions.jl |  6 +++---
 moment_kinetics/src/load_data.jl          |  4 ++--
 moment_kinetics/src/moment_constraints.jl |  2 +-
 moment_kinetics/src/time_advance.jl       |  2 +-
 moment_kinetics/src/velocity_moments.jl   | 18 ++++++++++++------
 5 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index ea9fc5f92..927c73751 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -32,7 +32,7 @@ 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
-using ..velocity_moments: update_qpar!
+using ..velocity_moments: update_ion_qpar!
 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!
@@ -201,7 +201,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
 
         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,
+        update_ion_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)
@@ -1663,7 +1663,7 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
                  vpa, vperp, z, r, composition, moments.evolve_density,
                  moments.evolve_upar)
     update_pperp!(moments.ion.pperp, pdf.ion.norm, vpa, vperp, z, r, composition)
-    update_qpar!(moments.ion.qpar, moments.ion.qpar_updated,
+    update_ion_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,
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index cf9f56883..5243ecad8 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4251,7 +4251,7 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             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).
+            # `update_ion_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
@@ -4266,7 +4266,7 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
             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).
+            # `update_ion_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
diff --git a/moment_kinetics/src/moment_constraints.jl b/moment_kinetics/src/moment_constraints.jl
index f8c0a2274..9169af68c 100644
--- a/moment_kinetics/src/moment_constraints.jl
+++ b/moment_kinetics/src/moment_constraints.jl
@@ -8,7 +8,7 @@ module moment_constraints
 using ..communication: _block_synchronize
 using ..looping
 using ..type_definitions: mk_float
-using ..velocity_moments: integrate_over_vspace, update_qpar!
+using ..velocity_moments: integrate_over_vspace, update_ion_qpar!
 
 export hard_force_moment_constraints!, hard_force_moment_constraints_neutral!
 
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 40dcca66d..13ac7511e 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -20,7 +20,7 @@ using ..initial_conditions: initialize_electrons!
 using ..looping
 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_density!, update_upar!, update_ppar!, update_pperp!, update_ion_qpar!, update_vth!
 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!
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 096dae7b2..e898c5752 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -11,7 +11,7 @@ export update_density!
 export update_upar!
 export update_ppar!
 export update_pperp!
-export update_qpar!
+export update_ion_qpar!
 export update_vth!
 export reset_moments_status!
 export update_neutral_density!
@@ -466,7 +466,7 @@ function update_moments!(moments, ff_in, gyroavs::gyro_operators, vpa, vperp, z,
         end
         @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],
+            @views update_ion_qpar_species!(moments.ion.qpar[:,:,is],
                                         moments.ion.dens[:,:,is],
                                         moments.ion.upar[:,:,is],
                                         moments.ion.vth[:,:,is], ff[:,:,:,:,is], vpa,
@@ -736,7 +736,7 @@ end
 """
 NB: the incoming pdf is the normalized pdf
 """
-function update_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vperp, z, r,
+function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vperp, z, r,
                       composition, evolve_density, evolve_upar, evolve_ppar)
     @boundscheck composition.n_ion_species == size(qpar,3) || throw(BoundsError(qpar))
 
@@ -744,7 +744,7 @@ function update_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vperp, z
 
     @loop_s is begin
         if qpar_updated[is] == false
-            @views update_qpar_species!(qpar[:,:,is], density[:,:,is], upar[:,:,is],
+            @views update_ion_qpar_species!(qpar[:,:,is], density[:,:,is], upar[:,:,is],
                                         vth[:,:,is], pdf[:,:,:,:,is], vpa, vperp, z, r,
                                         evolve_density, evolve_upar, evolve_ppar)
             qpar_updated[is] = true
@@ -755,8 +755,9 @@ end
 """
 calculate the updated parallel heat flux (qpar) for a given species
 """
-function update_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
+function update_ion_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                               evolve_upar, evolve_ppar)
+    calculate_
     @boundscheck r.n == size(ff, 4) || throw(BoundsError(ff))
     @boundscheck z.n == size(ff, 3) || throw(BoundsError(ff))
     @boundscheck vperp.n == size(ff, 2) || throw(BoundsError(ff))
@@ -794,6 +795,11 @@ function update_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, ev
     return nothing
 end
 
+"""
+calculate parallel heat flux if ion composition flag is kinetic ions
+"""
+
+
 """
 runtime diagnostic routine for computing the Chodura ratio
 in a single species plasma with Z = 1
@@ -1621,7 +1627,7 @@ function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composi
         rethrow(e)
     end
     # update the parallel heat flux
-    update_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
+    update_ion_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
                  new_scratch.upar, moments.ion.vth, ff, vpa, vperp, z, r,
                  composition, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)

From 53a673c68fc6614ebfed31ba6f25a6dca876c7b4 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Wed, 18 Sep 2024 13:36:44 +0100
Subject: [PATCH 02/31] Add ion_physics flag to determine drift_kinetic,
 gyrokinetic or braginskii ion physics. Gyrokinetic flag replacement of the
 gyrokinetic_ions case will happen in a later commit, so right now the
 gyrokinetic option for ion_physics does the same as drift_kinetic.

---
 moment_kinetics/src/external_sources.jl |  4 ++-
 moment_kinetics/src/input_structs.jl    | 11 ++++++++
 moment_kinetics/src/species_input.jl    | 10 ++++++++
 moment_kinetics/src/velocity_moments.jl | 34 +++++++++++++++++++------
 4 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 444833d11..0c49c83b4 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -506,8 +506,10 @@ function initialize_external_source_amplitude!(moments, external_source_settings
                 end
             end
         end
+    end
 
-        # now do same for electron sources, which (if present) are mostly mirrors of ion sources
+    # now do same for electron sources, which (if present) are mostly mirrors of ion sources
+    for index ∈ eachindex(electron_source_settings)
         if electron_source_settings[index].active
             if electron_source_settings[index].source_type == "energy"
                 @loop_r_z ir iz begin
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index ec5750e95..1c783efcd 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -185,6 +185,15 @@ export braginskii_fluid
 export kinetic_electrons
 export kinetic_electrons_with_temperature_equation
 
+@enum ion_physics_type begin
+    gyrokinetic_ions
+    drift_kinetic_ions
+    braginskii_ions
+end
+export ion_physics_type
+export gyrokinetic_ions
+export drift_kinetic_ions
+export braginskii_ions
 """
 """
 mutable struct grid_input_mutable
@@ -380,6 +389,8 @@ Base.@kwdef struct species_composition
     #   density is fixed to be Nₑ*(eϕ/T_e) and N_e is calculated using a current
     #   condition at the wall
     electron_physics::electron_physics_type
+    # ion physics can be drift_kinetic_ions, gyrokinetic_ions and braginskii_ions
+    ion_physics::ion_physics_type
     # 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
diff --git a/moment_kinetics/src/species_input.jl b/moment_kinetics/src/species_input.jl
index 35d8755e0..a7a2ff41b 100644
--- a/moment_kinetics/src/species_input.jl
+++ b/moment_kinetics/src/species_input.jl
@@ -11,6 +11,7 @@ using ..input_structs: set_defaults_and_check_section!
 using ..input_structs: species_composition, ion_species_parameters, neutral_species_parameters
 using ..input_structs: spatial_initial_condition_input, velocity_initial_condition_input
 using ..input_structs: boltzmann_electron_response, boltzmann_electron_response_with_simple_sheath
+using ..input_structs: drift_kinetic_ions
 using ..reference_parameters: setup_reference_parameters
 
 function get_species_input(toml_input)
@@ -29,6 +30,15 @@ function get_species_input(toml_input)
         #   electron density is fixed to be N_e*(eϕ/T_e) and N_e is calculated w.r.t a
         #   reference value using J_||e + J_||i = 0 at z = 0
         electron_physics = boltzmann_electron_response,
+        # If ion_physics=drift_kinetic_ions, the ion distribution function is advanced in 
+        # time in the drift kinetic approximation like usual. 
+        # If ion_physics=gyrokinetic_ions, the ion distribution function is
+        # advanced in time using gyroaveraged fields at fixed guiding centre and moments of the
+        # pdf computed at fixed r
+        # If ion_physics=braginskii_ions, there is no need for a shape function to evolve, and the code 
+        # only evolves ions in a fluid sense (i.e. all evolve_moments are set to true), with a 
+        # Braginskii closure for the ion heat flux.
+        ion_physics = drift_kinetic_ions,
         # initial Tₑ = 1
         T_e = 1.0,
         # wall temperature T_wall = Tw/Te
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index e898c5752..5ccbd6c52 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -471,7 +471,7 @@ function update_moments!(moments, ff_in, gyroavs::gyro_operators, vpa, vperp, z,
                                         moments.ion.upar[:,:,is],
                                         moments.ion.vth[:,:,is], ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
-                                        moments.evolve_upar, moments.evolve_ppar)
+                                        moments.evolve_upar, moments.evolve_ppar, composition.ion_physics)
             moments.ion.qpar_updated[is] = true
         end
     end
@@ -746,7 +746,8 @@ function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vper
         if qpar_updated[is] == false
             @views update_ion_qpar_species!(qpar[:,:,is], density[:,:,is], upar[:,:,is],
                                         vth[:,:,is], pdf[:,:,:,:,is], vpa, vperp, z, r,
-                                        evolve_density, evolve_upar, evolve_ppar)
+                                        evolve_density, evolve_upar, evolve_ppar, 
+                                        composition.ion_physics)
             qpar_updated[is] = true
         end
     end
@@ -756,8 +757,24 @@ end
 calculate the updated parallel heat flux (qpar) for a given species
 """
 function update_ion_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
-                              evolve_upar, evolve_ppar)
-    calculate_
+                                  evolve_upar, evolve_ppar, ion_physics)
+    if ion_physics == drift_kinetic_ions || ion_physics == gyrokinetic_ions
+        calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
+                                     evolve_upar, evolve_ppar)
+    elseif ion_physics == braginskii_ions
+        calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
+                                     evolve_upar, evolve_ppar)
+    else
+        throw(ArgumentError("ion model $ion_physics not implemented for qpar calculation"))
+    end
+    return nothing
+end
+
+"""
+calculate parallel heat flux if ion composition flag is kinetic ions
+"""
+function calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density, 
+                                      evolve_upar, evolve_ppar)
     @boundscheck r.n == size(ff, 4) || throw(BoundsError(ff))
     @boundscheck z.n == size(ff, 3) || throw(BoundsError(ff))
     @boundscheck vperp.n == size(ff, 2) || throw(BoundsError(ff))
@@ -794,12 +811,13 @@ function update_ion_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r
     end
     return nothing
 end
-
 """
-calculate parallel heat flux if ion composition flag is kinetic ions
+calculate parallel heat flux if ion composition flag is Braginskii fluid ions
 """
-
-
+function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
+                                     evolve_upar, evolve_ppar)
+    return nothing
+end
 """
 runtime diagnostic routine for computing the Chodura ratio
 in a single species plasma with Z = 1

From 6f6df749e78e13b3e4440c962952a28b1bc3a4be Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Wed, 18 Sep 2024 14:07:04 +0100
Subject: [PATCH 03/31] Update gyrokinetic flag uses throughout the repo, so
 that instead of gyrokinetic_ions = true being the test, we see whether
 ion_physics == gyrokinetic_ions

---
 examples/gk-ions/2D-periodic-gk.toml             | 2 +-
 moment_kinetics/debug_test/gyroaverage_inputs.jl | 2 +-
 moment_kinetics/src/em_fields.jl                 | 2 +-
 moment_kinetics/src/gyroaverages.jl              | 3 ++-
 moment_kinetics/src/input_structs.jl             | 7 +++----
 moment_kinetics/src/species_input.jl             | 5 +----
 moment_kinetics/src/velocity_moments.jl          | 6 +++---
 moment_kinetics/test/gyroaverage_tests.jl        | 2 +-
 test_scripts/gyroaverage_test.jl                 | 4 ++--
 9 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/examples/gk-ions/2D-periodic-gk.toml b/examples/gk-ions/2D-periodic-gk.toml
index 19659aede..77daebb78 100644
--- a/examples/gk-ions/2D-periodic-gk.toml
+++ b/examples/gk-ions/2D-periodic-gk.toml
@@ -32,7 +32,7 @@ vz_discretization = "chebyshev_pseudospectral"
 n_ion_species = 1
 n_neutral_species = 0
 electron_physics = "boltzmann_electron_response"
-gyrokinetic_ions = true
+ion_physics = "gyrokinetic_ions"
 T_e = 1.0
 T_wall = 1.0
 
diff --git a/moment_kinetics/debug_test/gyroaverage_inputs.jl b/moment_kinetics/debug_test/gyroaverage_inputs.jl
index 516d4c674..0343b910f 100644
--- a/moment_kinetics/debug_test/gyroaverage_inputs.jl
+++ b/moment_kinetics/debug_test/gyroaverage_inputs.jl
@@ -6,7 +6,7 @@ test_input = OptionsDict(
     "run_name" => "gyroaverage",
     "composition" => OptionsDict("n_ion_species" => 1,
                           "n_neutral_species" => 0,
-                          "gyrokinetic_ions" => true,
+                          "ion_physics" => "gyrokinetic_ions",
                           "T_e" => 1.0,
                           "T_wall" => 1.0),
     "evolve_moments_density" => false,
diff --git a/moment_kinetics/src/em_fields.jl b/moment_kinetics/src/em_fields.jl
index 9f26f009c..2518f53be 100644
--- a/moment_kinetics/src/em_fields.jl
+++ b/moment_kinetics/src/em_fields.jl
@@ -156,7 +156,7 @@ function update_phi!(fields, fvec, vperp, z, r, composition, collisions, moments
     end
 
     # get gyroaveraged field arrays for distribution function advance
-    gkions = composition.gyrokinetic_ions
+    gkions = composition.ion_physics == gyrokinetic_ions
     if gkions
         gyroaverage_field!(fields.gphi,fields.phi,gyroavs,vperp,z,r,composition)
         gyroaverage_field!(fields.gEz,fields.Ez,gyroavs,vperp,z,r,composition)
diff --git a/moment_kinetics/src/gyroaverages.jl b/moment_kinetics/src/gyroaverages.jl
index b288e0da6..6b6dd26f9 100644
--- a/moment_kinetics/src/gyroaverages.jl
+++ b/moment_kinetics/src/gyroaverages.jl
@@ -13,6 +13,7 @@ using ..type_definitions: mk_float, mk_int
 using ..array_allocation: allocate_float, allocate_shared_float
 using ..array_allocation: allocate_int, allocate_shared_int
 using ..lagrange_polynomials: lagrange_poly
+using ..input_structs: gyrokinetic_ions
 using ..looping
 using ..communication: MPISharedArray, comm_block, _block_synchronize
 
@@ -39,7 +40,7 @@ other nonzero boundary conditions for the z and r domains.
 """
 
 function init_gyro_operators(vperp,z,r,gyrophase,geometry,composition;print_info=false)
-    gkions = composition.gyrokinetic_ions
+    gkions = composition.ion_physics == gyrokinetic_ions
     if !gkions
         gyromatrix =  allocate_shared_float(1,1,1,1,1,1)
         gyroloopsizes = allocate_shared_int(1,1,1,1)
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 1c783efcd..d40c68dd0 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -390,6 +390,9 @@ Base.@kwdef struct species_composition
     #   condition at the wall
     electron_physics::electron_physics_type
     # ion physics can be drift_kinetic_ions, gyrokinetic_ions and braginskii_ions
+    # gyrokinetic_ions (originally gyrokinetic_ions = true) -> use gyroaveraged fields at fixed guiding 
+    # centre and moments of the pdf computed at fixed r
+    # drift_kinetic_ions (originally gyrokinetic_ions = false) -> use drift kinetic approximation
     ion_physics::ion_physics_type
     # 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
@@ -407,10 +410,6 @@ Base.@kwdef struct species_composition
     # 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::mk_float
-    # gyrokinetic_ions is a flag determining if the ion species is gyrokinetic
-    # gyrokinetic_ions = true -> use gyroaveraged fields at fixed guiding centre and moments of the pdf computed at fixed r
-    # gyrokinetic_ions = false -> use drift kinetic approximation
-    gyrokinetic_ions::Bool
     # array of structs of parameters for each ion species
     ion::Vector{ion_species_parameters}
     # array of structs of parameters for each neutral species
diff --git a/moment_kinetics/src/species_input.jl b/moment_kinetics/src/species_input.jl
index a7a2ff41b..abc3c4bb8 100644
--- a/moment_kinetics/src/species_input.jl
+++ b/moment_kinetics/src/species_input.jl
@@ -53,10 +53,7 @@ function get_species_input(toml_input)
         use_test_neutral_wall_pdf = false,
         # 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,
-        # gyrokinetic_ions = True -> use gyroaveraged fields at fixed guiding centre and moments of the pdf computed at fixed r
-        # gyrokinetic_ions = False -> use drift kinetic approximation
-        gyrokinetic_ions = false)
+        recycling_fraction = 1.0)
 
     nspec_ion = composition_section["n_ion_species"]
     nspec_neutral = composition_section["n_neutral_species"]
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 5ccbd6c52..4d8d1e304 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -428,7 +428,7 @@ the function used to update moments at run time is update_derived_moments! in ti
 """
 function update_moments!(moments, ff_in, gyroavs::gyro_operators, vpa, vperp, z, r, composition,
         r_spectral, geometry, scratch_dummy, z_advect)
-    if composition.gyrokinetic_ions
+    if composition.ion_physics == gyrokinetic_ions
         ff = scratch_dummy.buffer_vpavperpzrs_1 # the buffer array for the ion pdf -> make sure not to reuse this array below
         # fill buffer with ring-averaged F (gyroaverage at fixed position)
         gyroaverage_pdf!(ff,ff_in,gyroavs,vpa,vperp,z,r,composition)
@@ -758,7 +758,7 @@ calculate the updated parallel heat flux (qpar) for a given species
 """
 function update_ion_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                   evolve_upar, evolve_ppar, ion_physics)
-    if ion_physics == drift_kinetic_ions || ion_physics == gyrokinetic_ions
+    if ion_physics ∈ (drift_kinetic_ions, gyrokinetic_ions)
         calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                      evolve_upar, evolve_ppar)
     elseif ion_physics == braginskii_ions
@@ -1599,7 +1599,7 @@ update velocity moments that are calculable from the evolved ion pdf
 function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
     r_spectral, geometry, gyroavs, scratch_dummy, z_advect, diagnostic_moments)
 
-    if composition.gyrokinetic_ions
+    if composition.ion_physics == gyrokinetic_ions
         ff = scratch_dummy.buffer_vpavperpzrs_1
         # fill buffer with ring-averaged F (gyroaverage at fixed position)
         gyroaverage_pdf!(ff,new_scratch.pdf,gyroavs,vpa,vperp,z,r,composition)
diff --git a/moment_kinetics/test/gyroaverage_tests.jl b/moment_kinetics/test/gyroaverage_tests.jl
index c4f2e5a2a..429a5f022 100644
--- a/moment_kinetics/test/gyroaverage_tests.jl
+++ b/moment_kinetics/test/gyroaverage_tests.jl
@@ -230,7 +230,7 @@ function gyroaverage_test(absolute_error; rhostar=0.1, pitch=0.5, ngrid=5, kr=2,
 end
 
 function create_test_composition()
-    input_dict = OptionsDict("composition" => OptionsDict("n_ion_species" => 1, "n_neutral_species" => 0, "gyrokinetic_ions" => true ) )
+    input_dict = OptionsDict("composition" => OptionsDict("n_ion_species" => 1, "n_neutral_species" => 0, "ion_physics" => "gyrokinetic_ions") )
     #println(input_dict)
     return get_species_input(input_dict)
 end
diff --git a/test_scripts/gyroaverage_test.jl b/test_scripts/gyroaverage_test.jl
index eb1c3bcb6..94105998c 100644
--- a/test_scripts/gyroaverage_test.jl
+++ b/test_scripts/gyroaverage_test.jl
@@ -213,10 +213,10 @@ function create_test_composition()
     # 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
-    gyrokinetic_ions = true
+    ion_physics = gyrokinetic_ions
     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,
-            mn_over_mi, me_over_mi, recycling_fraction, gyrokinetic_ions, allocate_float(n_species))
+            mn_over_mi, me_over_mi, recycling_fraction, ion_physics, allocate_float(n_species))
 end
 
 function fill_test_arrays!(phi,gphi,vperp,z,r,geometry,kz,kr,phasez,phaser)

From e92dc5cac52f99d436539ef3e5d03338b1a63aef Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 19 Sep 2024 09:45:38 +0100
Subject: [PATCH 04/31] Some of the Braginskii functionality (unfinished)

---
 moment_kinetics/src/velocity_moments.jl | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 4d8d1e304..2f7130adf 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -814,8 +814,24 @@ end
 """
 calculate parallel heat flux if ion composition flag is Braginskii fluid ions
 """
-function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
-                                     evolve_upar, evolve_ppar)
+function calculate_ion_qpar_from_braginskii!(qpar, density, vth, ff, vpa, vperp, z, r, collisions)
+    # Note that this is a braginskii heat flux for ions using the krook operator. The full Fokker-Planck operator
+    # Braginskii heat flux is different! This also assumes one ion species, and so no friction between ions.
+    @boundscheck r.n == size(ff, 4) || throw(BoundsError(ff))
+    @boundscheck z.n == size(ff, 3) || throw(BoundsError(ff))
+    @boundscheck vperp.n == size(ff, 2) || throw(BoundsError(ff))
+    @boundscheck vpa.n == size(ff, 1) || throw(BoundsError(ff))
+    @boundscheck r.n == size(qpar, 2) || throw(BoundsError(qpar))
+    @boundscheck z.n == size(qpar, 1) || throw(BoundsError(qpar))
+    # For now, I'll do the dT_dz calculation here, because it is only used for the Braginskii so should
+    # not clutter up the rest of the code.
+    dT_dz = 
+    begin_r_z_region()
+    @loop_r_z ir iz begin
+        nu_ii = get_collision_frequency_ii(collisions, density[iz,ir], vth[iz,ir])
+        qpar[iz,ir] = -(1/2) * 5/4 * density[iz,ir] * vth[iz,ir]^2 /nu_ii * dT_dz[iz,ir]
+    end
+
     return nothing
 end
 """

From 1dbf6612ae91456091b7837eb7e7017e13e75789 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 19 Sep 2024 11:34:31 +0100
Subject: [PATCH 05/31] Make initialisation for braginskii ion heat flux be
 calculated from unnormalised pdf (then subsequent heat flux updates will be
 with the braginskii closure)

---
 moment_kinetics/src/initial_conditions.jl | 7 ++++---
 moment_kinetics/src/velocity_moments.jl   | 6 +++---
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index fdec6d36b..524d46dbc 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -200,10 +200,11 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                         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
+        # calculate the initial parallel heat flux from the initial un-normalised pdf. Even if Braginskii fluid is being
+        # advanced, initialised ion_qpar uses the pdf 
         update_ion_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,
+                     pdf.ion.norm, vpa, vperp, z, r, composition, drift_kinetic_ions,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
 
         begin_serial_region()
@@ -1668,7 +1669,7 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     update_ion_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,
+                 composition, drift_kinetic_ions, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
     update_vth!(moments.ion.vth, moments.ion.ppar, moments.ion.pperp, moments.ion.dens, vperp, z, r, composition)
 
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 2f7130adf..618c66c92 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -737,7 +737,7 @@ end
 NB: the incoming pdf is the normalized pdf
 """
 function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vperp, z, r,
-                      composition, evolve_density, evolve_upar, evolve_ppar)
+                          composition, ion_physics, evolve_density, evolve_upar, evolve_ppar)
     @boundscheck composition.n_ion_species == size(qpar,3) || throw(BoundsError(qpar))
 
     begin_s_r_z_region()
@@ -747,7 +747,7 @@ function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vper
             @views update_ion_qpar_species!(qpar[:,:,is], density[:,:,is], upar[:,:,is],
                                         vth[:,:,is], pdf[:,:,:,:,is], vpa, vperp, z, r,
                                         evolve_density, evolve_upar, evolve_ppar, 
-                                        composition.ion_physics)
+                                        ion_physics)
             qpar_updated[is] = true
         end
     end
@@ -1663,7 +1663,7 @@ function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composi
     # update the parallel heat flux
     update_ion_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
                  new_scratch.upar, moments.ion.vth, ff, vpa, vperp, z, r,
-                 composition, moments.evolve_density, moments.evolve_upar,
+                 composition, composition.ion_physics, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
     # add further moments to be computed here
 

From 3b9bc66d8e4eb805dfc62bbf21f86329be22f06c Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 20 Sep 2024 10:36:26 +0100
Subject: [PATCH 06/31] Add braginskii heat flux formula by adding an ion dT_dz
 component in the struct, as well as an ion temp

---
 moment_kinetics/src/initial_conditions.jl     | 13 ++--
 .../src/moment_kinetics_structs.jl            |  4 ++
 moment_kinetics/src/time_advance.jl           |  6 +-
 moment_kinetics/src/velocity_moments.jl       | 70 +++++++++++--------
 moment_kinetics/src/vpa_advection.jl          |  2 +-
 5 files changed, 56 insertions(+), 39 deletions(-)

diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index 524d46dbc..a142a65b8 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -141,7 +141,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                                              r, composition.n_ion_species,
                                              composition.n_neutral_species,
                                              geometry.input, composition, species,
-                                             manufactured_solns_input)
+                                             manufactured_solns_input, collisions)
     else
         n_ion_species = composition.n_ion_species
         n_neutral_species = composition.n_neutral_species
@@ -203,8 +203,8 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
         # calculate the initial parallel heat flux from the initial un-normalised pdf. Even if Braginskii fluid is being
         # advanced, initialised ion_qpar uses the pdf 
         update_ion_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, drift_kinetic_ions,
+                     moments.ion.dens, moments.ion.upar, moments.ion.vth, moments.ion.dT_dz,
+                     pdf.ion.norm, vpa, vperp, z, r, composition, drift_kinetic_ions, collisions,
                      moments.evolve_density, moments.evolve_upar, moments.evolve_ppar)
 
         begin_serial_region()
@@ -1639,7 +1639,8 @@ function init_electron_pdf_over_density_and_boundary_phi!(pdf, phi, density, upa
     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)
+function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa, vperp, z, r, n_ion_species, n_neutral_species, 
+                                              geometry, composition, species, manufactured_solns_input, collisions)
     manufactured_solns_list = manufactured_solutions(r.L,z.L,r.bc,z.bc,geometry,composition,r.n)
     dfni_func = manufactured_solns_list.dfni_func
     densi_func = manufactured_solns_list.densi_func
@@ -1668,8 +1669,8 @@ function init_pdf_moments_manufactured_solns!(pdf, moments, vz, vr, vzeta, vpa,
     update_pperp!(moments.ion.pperp, pdf.ion.norm, vpa, vperp, z, r, composition)
     update_ion_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, drift_kinetic_ions, moments.evolve_density, moments.evolve_upar,
+                 moments.ion.vth, moments.ion.dT_dz, pdf.ion.norm, vpa, vperp, z, r,
+                 composition, drift_kinetic_ions, collisions, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
     update_vth!(moments.ion.vth, moments.ion.ppar, moments.ion.pperp, moments.ion.dens, vperp, z, r, composition)
 
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index a1ef580e2..8568eade6 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -95,6 +95,8 @@ struct moments_ion_substruct
     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}
+    # this is the temperature, calculated from 2ppar/dens (the comment below may be out of date?)
+    temp::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}
@@ -125,6 +127,8 @@ struct moments_ion_substruct
     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 z-derivative of the temperature
+    dT_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 (third index is for different sources)
diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 0fc9f4fd4..4814117c1 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -942,7 +942,7 @@ function setup_time_advance!(pdf, fields, vz, vr, vzeta, vpa, vperp, z, r, gyrop
         # constraints to the pdf
         reset_moments_status!(moments)
         update_moments!(moments, pdf.ion.norm, gyroavs, vpa, vperp, z, r, composition,
-           r_spectral,geometry,scratch_dummy,z_advect)
+           r_spectral,geometry,scratch_dummy,z_advect, collisions)
         # 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
@@ -2351,7 +2351,7 @@ function apply_all_bcs_constraints_update_moments!(
     # calculated before that is applied. Also may be needed to calculate advection speeds
     # for for CFL stability limit calculations in adaptive_timestep_update!().
     update_derived_moments!(this_scratch, moments, vpa, vperp, z, r, composition,
-        r_spectral, geometry, gyroavs, scratch_dummy, z_advect, diagnostic_moments)
+        r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, diagnostic_moments)
 
     calculate_ion_moment_derivatives!(moments, this_scratch, scratch_dummy, z, z_spectral,
                                       num_diss_params.ion.moment_dissipation_coefficient)
@@ -3727,7 +3727,7 @@ function implicit_ion_advance!(fvec_out, fvec_in, pdf, fields, moments, advect_o
         # Ensure moments are consistent with f_new
         update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
                                 r_spectral, geometry, gyroavs, scratch_dummy, z_advect,
-                                false)
+                                collisions, false)
         calculate_ion_moment_derivatives!(moments, new_scratch, scratch_dummy, z,
                                           z_spectral,
                                           num_diss_params.ion.moment_dissipation_coefficient)
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 618c66c92..1cfe5dc77 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -41,10 +41,12 @@ using ..derivatives: derivative_z!, second_derivative_z!
 using ..derivatives: derivative_r!, second_derivative_r!
 using ..looping
 using ..gyroaverages: gyro_operators, gyroaverage_pdf!
+using ..krook_collisions: get_collision_frequency_ii
 using ..input_structs
 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)
@@ -71,6 +73,8 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
     parallel_pressure_updated .= false
     # allocate array used for the perpendicular pressure
     perpendicular_pressure = allocate_shared_float(nz, nr, n_species)
+    # allocate array used for the temperature
+    temperature = allocate_shared_float(nz, nr, n_species)
     # allocate array used for the parallel flow
     parallel_heat_flux = allocate_shared_float(nz, nr, n_species)
     # allocate array of Bools that indicate if the parallel flow is updated for each species
@@ -129,11 +133,13 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
         d2ppar_dz2 = allocate_shared_float(nz, nr, n_species)
         dqpar_dz = allocate_shared_float(nz, nr, n_species)
         dvth_dz = allocate_shared_float(nz, nr, n_species)
+        dT_dz = allocate_shared_float(nz, nr, n_species)
     else
         dppar_dz_upwind = nothing
         d2ppar_dz2 = nothing
         dqpar_dz = nothing
         dvth_dz = nothing
+        dT_dz = nothing
     end
 
     entropy_production = allocate_shared_float(nz, nr, n_species)
@@ -187,10 +193,10 @@ function create_moments_ion(nz, nr, n_species, evolve_density, evolve_upar,
     # 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,
-        parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, 
+        parallel_heat_flux, parallel_heat_flux_updated, thermal_speed, temperature, 
         chodura_integral_lower, chodura_integral_upper, v_norm_fac,
         ddens_dz, ddens_dz_upwind, d2dens_dz2, dupar_dz, dupar_dz_upwind, d2upar_dz2,
-        dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz, entropy_production,
+        dppar_dz, dppar_dz_upwind, d2ppar_dz2, dqpar_dz, dvth_dz, dT_dz, entropy_production,
         external_source_amplitude, external_source_density_amplitude,
         external_source_momentum_amplitude, external_source_pressure_amplitude,
         external_source_controller_integral, constraints_A_coefficient,
@@ -427,7 +433,7 @@ this function is only used once after initialisation
 the function used to update moments at run time is update_derived_moments! in time_advance.jl
 """
 function update_moments!(moments, ff_in, gyroavs::gyro_operators, vpa, vperp, z, r, composition,
-        r_spectral, geometry, scratch_dummy, z_advect)
+        r_spectral, geometry, scratch_dummy, z_advect, collisions)
     if composition.ion_physics == gyrokinetic_ions
         ff = scratch_dummy.buffer_vpavperpzrs_1 # the buffer array for the ion pdf -> make sure not to reuse this array below
         # fill buffer with ring-averaged F (gyroaverage at fixed position)
@@ -469,9 +475,9 @@ function update_moments!(moments, ff_in, gyroavs::gyro_operators, vpa, vperp, z,
             @views update_ion_qpar_species!(moments.ion.qpar[:,:,is],
                                         moments.ion.dens[:,:,is],
                                         moments.ion.upar[:,:,is],
-                                        moments.ion.vth[:,:,is], ff[:,:,:,:,is], vpa,
+                                        moments.ion.vth[:,:,is],  moments.ion.dT_dz, ff[:,:,:,:,is], vpa,
                                         vperp, z, r, moments.evolve_density,
-                                        moments.evolve_upar, moments.evolve_ppar, composition.ion_physics)
+                                        moments.evolve_upar, moments.evolve_ppar, composition.ion_physics, collisions)
             moments.ion.qpar_updated[is] = true
         end
     end
@@ -736,8 +742,8 @@ end
 """
 NB: the incoming pdf is the normalized pdf
 """
-function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vperp, z, r,
-                          composition, ion_physics, evolve_density, evolve_upar, evolve_ppar)
+function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, dT_dz, pdf, vpa, vperp, z, r,
+                          composition, ion_physics, collisions, evolve_density, evolve_upar, evolve_ppar)
     @boundscheck composition.n_ion_species == size(qpar,3) || throw(BoundsError(qpar))
 
     begin_s_r_z_region()
@@ -745,9 +751,9 @@ function update_ion_qpar!(qpar, qpar_updated, density, upar, vth, pdf, vpa, vper
     @loop_s is begin
         if qpar_updated[is] == false
             @views update_ion_qpar_species!(qpar[:,:,is], density[:,:,is], upar[:,:,is],
-                                        vth[:,:,is], pdf[:,:,:,:,is], vpa, vperp, z, r,
+                                        vth[:,:,is], dT_dz, pdf[:,:,:,:,is], vpa, vperp, z, r,
                                         evolve_density, evolve_upar, evolve_ppar, 
-                                        ion_physics)
+                                        ion_physics, collisions)
             qpar_updated[is] = true
         end
     end
@@ -756,14 +762,14 @@ end
 """
 calculate the updated parallel heat flux (qpar) for a given species
 """
-function update_ion_qpar_species!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
-                                  evolve_upar, evolve_ppar, ion_physics)
+function update_ion_qpar_species!(qpar, density, upar, vth, dT_dz, ff, vpa, vperp, z, r, evolve_density,
+                                  evolve_upar, evolve_ppar, ion_physics, collisions)
     if ion_physics ∈ (drift_kinetic_ions, gyrokinetic_ions)
         calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                      evolve_upar, evolve_ppar)
     elseif ion_physics == braginskii_ions
-        calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
-                                     evolve_upar, evolve_ppar)
+        calculate_ion_qpar_from_braginskii!(qpar, density, vth, dT_dz, z, r, collisions, evolve_density, 
+                                            evolve_upar, evolve_ppar)
     else
         throw(ArgumentError("ion model $ion_physics not implemented for qpar calculation"))
     end
@@ -814,24 +820,22 @@ end
 """
 calculate parallel heat flux if ion composition flag is Braginskii fluid ions
 """
-function calculate_ion_qpar_from_braginskii!(qpar, density, vth, ff, vpa, vperp, z, r, collisions)
+function calculate_ion_qpar_from_braginskii!(qpar, density, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
     # Note that this is a braginskii heat flux for ions using the krook operator. The full Fokker-Planck operator
     # Braginskii heat flux is different! This also assumes one ion species, and so no friction between ions.
-    @boundscheck r.n == size(ff, 4) || throw(BoundsError(ff))
-    @boundscheck z.n == size(ff, 3) || throw(BoundsError(ff))
-    @boundscheck vperp.n == size(ff, 2) || throw(BoundsError(ff))
-    @boundscheck vpa.n == size(ff, 1) || throw(BoundsError(ff))
     @boundscheck r.n == size(qpar, 2) || throw(BoundsError(qpar))
     @boundscheck z.n == size(qpar, 1) || throw(BoundsError(qpar))
-    # For now, I'll do the dT_dz calculation here, because it is only used for the Braginskii so should
-    # not clutter up the rest of the code.
-    dT_dz = 
-    begin_r_z_region()
-    @loop_r_z ir iz begin
-        nu_ii = get_collision_frequency_ii(collisions, density[iz,ir], vth[iz,ir])
-        qpar[iz,ir] = -(1/2) * 5/4 * density[iz,ir] * vth[iz,ir]^2 /nu_ii * dT_dz[iz,ir]
-    end
 
+    # calculate braginskii heat flux. Currently only works for one ion species! (hence the 1 in dT_dz[iz,ir,1])
+    if evolve_density && evolve_upar && evolve_ppar
+        begin_r_z_region()
+        @loop_r_z ir iz begin
+            nu_ii = get_collision_frequency_ii(collisions, density[iz,ir], vth[iz,ir])
+            qpar[iz,ir] = -(1/2) * 5/4 * density[iz,ir] * vth[iz,ir]^2 /nu_ii * dT_dz[iz,ir,1]
+        end
+    else
+        throw(ArgumentError("Braginskii heat flux simulation requires evolve_density, evolve_upar and evolve_ppar to be true, since it is a purely fluid simulation"))
+    end
     return nothing
 end
 """
@@ -999,6 +1003,14 @@ function calculate_ion_moment_derivatives!(moments, scratch, scratch_dummy, z, z
                              buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
         @views derivative_z!(moments.ion.dvth_dz, vth, buffer_r_1,
                              buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
+
+        # calculate the z derivative of the ion temperature
+        @loop_s_r_z is ir iz begin
+            # store the temperature in dummy_zrs
+            dummy_zrs[iz,ir,is] = 2*ppar[iz,ir,is]/density[iz,ir,is]
+        end
+        @views derivative_z!(moments.ion.dT_dz, dummy_zrs, buffer_r_1,
+                            buffer_r_2, buffer_r_3, buffer_r_4, z_spectral, z)
     end
 end
 
@@ -1613,7 +1625,7 @@ end
 update velocity moments that are calculable from the evolved ion pdf
 """
 function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composition,
-    r_spectral, geometry, gyroavs, scratch_dummy, z_advect, diagnostic_moments)
+    r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, diagnostic_moments)
 
     if composition.ion_physics == gyrokinetic_ions
         ff = scratch_dummy.buffer_vpavperpzrs_1
@@ -1662,8 +1674,8 @@ function update_derived_moments!(new_scratch, moments, vpa, vperp, z, r, composi
     end
     # update the parallel heat flux
     update_ion_qpar!(moments.ion.qpar, moments.ion.qpar_updated, new_scratch.density,
-                 new_scratch.upar, moments.ion.vth, ff, vpa, vperp, z, r,
-                 composition, composition.ion_physics, moments.evolve_density, moments.evolve_upar,
+                 new_scratch.upar, moments.ion.vth, moments.ion.dT_dz, ff, vpa, vperp, z, r,
+                 composition, composition.ion_physics, collisions, moments.evolve_density, moments.evolve_upar,
                  moments.evolve_ppar)
     # add further moments to be computed here
 
diff --git a/moment_kinetics/src/vpa_advection.jl b/moment_kinetics/src/vpa_advection.jl
index 74033245d..edfb57b5b 100644
--- a/moment_kinetics/src/vpa_advection.jl
+++ b/moment_kinetics/src/vpa_advection.jl
@@ -63,7 +63,7 @@ function implicit_vpa_advection!(f_out, fvec_in, fields, moments, z_advect, vpa_
                               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)
+                            r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, false)
     calculate_ion_moment_derivatives!(moments, new_scratch, scratch_dummy, z,
                                       z_spectral,
                                       num_diss_params.ion.moment_dissipation_coefficient)

From 1d8b5592889a7a345977e8147e8a6c7063401555 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 20 Sep 2024 10:53:51 +0100
Subject: [PATCH 07/31] add loop to allow n_neutral_species = 0 and CFL plots
 to still be made for ions

---
 moment_kinetics/src/load_data.jl | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index b1a146944..9a60c0d3d 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4343,9 +4343,6 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         density = get_variable(run_info, "density")
         upar = get_variable(run_info, "parallel_flow")
         ppar = get_variable(run_info, "parallel_pressure")
-        density_neutral = get_variable(run_info, "density_neutral")
-        uz_neutral = get_variable(run_info, "uz_neutral")
-        pz_neutral = get_variable(run_info, "pz_neutral")
         vth = get_variable(run_info, "thermal_speed")
         dupar_dz = get_z_derivative(run_info, "parallel_flow")
         dppar_dz = get_z_derivative(run_info, "parallel_pressure")
@@ -4395,6 +4392,15 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         setup_loop_ranges!(0, 1; s=nspecies, 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)
+        
+        # Use neutrals for fvec calculation in moment_kinetic version only when 
+        # n_neutrals != 0
+        if run_info.n_neutral_species != 0
+            density_neutral = get_variable(run_info, "density_neutral")
+            uz_neutral = get_variable(run_info, "uz_neutral")
+            pz_neutral = get_variable(run_info, "pz_neutral")
+        end
+
         for it ∈ 1:nt
             begin_serial_region()
             # Only need some struct with a 'speed' variable
@@ -4413,12 +4419,18 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
                              evolve_density=run_info.evolve_density,
                              evolve_upar=run_info.evolve_upar,
                              evolve_ppar=run_info.evolve_ppar)
-            @views fvec = (density=density[:,:,:,it],
-                           upar=upar[:,:,:,it],
-                           ppar=ppar[:,:,:,it],
-                           density_neutral=density_neutral[:,:,:,it],
-                           uz_neutral=uz_neutral[:,:,:,it],
-                           pz_neutral=pz_neutral[:,:,:,it])
+            if run_info.n_neutral_species != 0
+                @views fvec = (density=density[:,:,:,it],
+                            upar=upar[:,:,:,it],
+                            ppar=ppar[:,:,:,it],
+                            density_neutral=density_neutral[:,:,:,it],
+                            uz_neutral=uz_neutral[:,:,:,it],
+                            pz_neutral=pz_neutral[:,:,:,it])
+            else
+                @views fvec = (density=density[:,:,:,it],
+                            upar=upar[:,:,:,it],
+                            ppar=ppar[:,:,:,it])
+            end
             @views update_speed_vpa!(advect, fields, fvec, moments, run_info.vpa,
                                      run_info.vperp, run_info.z, run_info.r,
                                      run_info.composition, run_info.collisions,

From 77cce1ada766a0d87071504cf1396e577fd4beda Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 20 Sep 2024 11:22:40 +0100
Subject: [PATCH 08/31] Move order of krook_collisions.jl inclusion to allow
 velocity_moments.jl to use the reference collision frequency to calculate
 braginskii heat flux

---
 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 86810a66f..33f4796c6 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -37,6 +37,7 @@ include("nonlinear_solvers.jl")
 include("file_io.jl")
 include("geo.jl")
 include("gyroaverages.jl")
+include("krook_collisions.jl")
 include("velocity_moments.jl")
 include("velocity_grid_transforms.jl")
 include("electron_fluid_equations.jl")
@@ -61,7 +62,6 @@ include("neutral_z_advection.jl")
 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")

From 18cd33e9a3cd3e4028603435516c7618562d4c37 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Sat, 21 Sep 2024 18:59:59 +0100
Subject: [PATCH 09/31] Add plotting for collisionality (comparing gradient
 scale lengths to mean free path) and braginskii heat flux boundary condition
 at walls, following Stangeby (though I haven't checked the subtracted part
 for conductive heat flux yet).

---
 .../src/makie_post_processing.jl              | 216 ++++++++++++++++++
 moment_kinetics/src/load_data.jl              |  76 ++++++
 moment_kinetics/src/velocity_moments.jl       |  49 +++-
 3 files changed, 338 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 653b28792..9c092286c 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
@@ -346,6 +346,8 @@ function makie_post_process(run_dir::Union{String,Tuple},
 
     sound_wave_plots(run_info; plot_prefix=plot_prefix)
 
+    collisionality_plots(run_info, plot_prefix)
+
     if all(ri === nothing for ri ∈ run_info_dfns)
         nvperp = nothing
     else
@@ -807,6 +809,24 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         plot_steady_state_residual=false,
         animate_steady_state_residual=false,
        )
+    
+    set_defaults_and_check_section!(
+        this_input_dict, "collisionality_plots";
+        plot=true,
+        plot_dT_dz_vs_z=false,
+        animate_dT_dz_vs_z=false,
+        plot_mfp_vs_z=false,
+        animate_mfp_vs_z=false,
+        plot_nu_ii_vth_mfp_vs_z = false,
+        plot_LT_mfp_vs_z = false,
+        animate_LT_mfp_vs_z = false,
+        plot_LT_dT_dz_temp_vs_z = false,
+        plot_Ln_mfp_vs_z = false,
+        animate_Ln_mfp_vs_z = false,
+        plot_Lupar_mfp_vs_z = false,
+        animate_Lupar_mfp_vs_z = false,
+        animation_ext = "gif"
+       )
 
     # We allow top-level options in the post-processing input file
     check_sections!(this_input_dict; check_no_top_level_options=false)
@@ -8130,6 +8150,202 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
     return steps_fig, dt_fig, CFL_fig
 end
 
+
+
+function collisionality_plots(run_info, plot_prefix=nothing)
+    # plot the mean free path of the ions at every point in z, and on the same graph
+    # plot the lengthscales of the gradients of all moments at each point in z
+    #println("run_info is ", run_info)
+    println("Making plots for collisionality")
+    if !isa(run_info, Tuple)
+        run_info = (run_info,)
+    end
+
+    input = Dict_to_NamedTuple(input_dict["collisionality_plots"])
+    temperature_input = Dict_to_NamedTuple(input_dict["temperature"])
+
+    mfp = get_variable(run_info, "mfp")
+    L_T = get_variable(run_info, "L_T")
+    L_n = get_variable(run_info, "L_n")
+    L_upar = get_variable(run_info, "L_upar")
+
+    if input.plot_mfp_vs_z
+        variable_name = "mean_free_path"
+        variable = mfp
+        variable_prefix = plot_prefix * variable_name
+        plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                  outfile=variable_prefix * "_vs_z.pdf")
+    end
+
+    if input.animate_mfp_vs_z
+        variable_name = "mean_free_path"
+        variable = mfp
+        variable_prefix = plot_prefix * variable_name
+        animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                        outfile=variable_prefix * "_vs_z." * input.animation_ext)
+    end
+    # TASKS: plot mean free path vs z, plot collision frequency vs z, make sure it agrees
+    # with original plot for it, and plot thermal speed and make sure that agrees, and then
+    # make sure that the mean free path then makes sense.
+    # Then plot the temperature and temperature gradient on the same plot
+    # Eventually plot the upar and dupar_dz and density and ddens_dz on plots too. Again,
+    # make sure they agree with the original plots.
+    if input.plot_dT_dz_vs_z
+        variable_name = "dT_dz"
+        variable = nothing
+        try
+            variable = get_variable(run_info, variable_name)
+        catch e
+            return makie_post_processing_error_handler(
+                    e,
+                    "collisionality_plots () failed for $variable_name - could not load data.")
+        end
+        #println("variable is ", variable)
+        variable_prefix = plot_prefix * variable_name
+        plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                  outfile=variable_prefix * "_vs_z.pdf")
+    end
+
+    if input.animate_dT_dz_vs_z
+        variable_name = "dT_dz"
+        variable = nothing
+        try
+            variable = get_variable(run_info, variable_name)
+        catch e
+            return makie_post_processing_error_handler(
+                    e,
+                    "collisionality_plots () failed for $variable_name - could not load data.")
+        end
+        #println("variable is ", variable)
+        variable_prefix = plot_prefix * variable_name
+        animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                        outfile=variable_prefix * "_vs_z." * input.animation_ext)
+    end
+
+    if input.plot_nu_ii_vth_mfp_vs_z
+        vth = get_variable(run_info, "thermal_speed")
+        nu_ii = get_variable(run_info, "collision_frequency_ii")
+        variable_prefix = plot_prefix * "checking_mfp_vth_and_nu_ii"
+        #println("vth[1] is ", vth[1])
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        plot_1d(run_info[1].z.grid, vth[1][:,1,1,21], xlabel="z",
+                ylabel="values", label="vth", ax=ax[1], title = "checking_mfp_vth")
+        plot_1d(run_info[1].z.grid, nu_ii[1][:,1,1,21], label="nu_ii", ax=ax[1])
+        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        #plot_1d(z.grid, select_slice(error, :z; input=input), xlabel=L"z",
+        #        ylabel=norm_label, ax=ax[2])
+        outfile = variable_prefix * "_vs_z.pdf"
+        save(outfile, fig)
+    end
+
+    if input.plot_LT_dT_dz_temp_vs_z
+        dT_dz = get_variable(run_info, "dT_dz")
+        temp = get_variable(run_info, "temperature")
+        variable_prefix = plot_prefix * "LT_dT_dz_temp"
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        plot_1d(run_info[1].z.grid, L_T[1][:,1,1,21], xlabel="z",
+                ylabel="values", label="L_T", ax=ax[1], title = "LT_dT_dz_temp")
+        plot_1d(run_info[1].z.grid, dT_dz[1][:,1,1,21], label="dT_dz", ax=ax[1])
+        plot_1d(run_info[1].z.grid, temp[1][:,1,1,21], label="temp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z.pdf"
+        save(outfile, fig)
+    end
+
+    if input.plot_LT_mfp_vs_z
+        variable_prefix = plot_prefix * "LT_mfp_first_comparison"
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        plot_1d(run_info[1].z.grid, L_T[1][:,1,1,21], xlabel="z",
+                ylabel="values", label="L_T", ax=ax[1], title = "L_T and mean free path comparison")
+        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z.pdf"
+        save(outfile, fig)
+    end
+
+    if input.animate_LT_mfp_vs_z
+        nt = length(mfp[1][1,1,1,:])
+        variable_prefix = plot_prefix * "LT_mfp_first_comparison"
+
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        frame_index = Observable(1)
+        animate_1d(run_info[1].z.grid, L_T[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="L_T", ax=ax[1])
+        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z." * input.animation_ext
+        save_animation(fig, frame_index, nt, outfile)
+    end
+
+    if input.plot_Ln_mfp_vs_z
+        variable_prefix = plot_prefix * "Ln_mfp_first_comparison"
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        plot_1d(run_info[1].z.grid, L_n[1][:,1,1,21], xlabel="z",
+                ylabel="values", label="L_n", ax=ax[1], title = "Ln and mean free path comparison")
+        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z.pdf"
+        save(outfile, fig)
+    end
+
+    if input.animate_Ln_mfp_vs_z
+        nt = length(mfp[1][1,1,1,:])
+        variable_prefix = plot_prefix * "Ln_mfp_first_comparison"
+
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        frame_index = Observable(1)
+        animate_1d(run_info[1].z.grid, L_n[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="L_n", ax=ax[1])
+        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z." * input.animation_ext
+        save_animation(fig, frame_index, nt, outfile)
+    end
+
+    if input.plot_Lupar_mfp_vs_z
+        variable_prefix = plot_prefix * "Lupar_mfp_first_comparison"
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        plot_1d(run_info[1].z.grid, L_upar[1][:,1,1,21], xlabel="z",
+                ylabel="values", label="L_upar", ax=ax[1], title = "Lupar and mean free path comparison")
+        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z.pdf"
+        save(outfile, fig)
+    end
+
+    if input.animate_Lupar_mfp_vs_z
+        nt = length(mfp[1][1,1,1,:])
+        variable_prefix = plot_prefix * "Lupar_mfp_first_comparison"
+
+        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+        frame_index = Observable(1)
+        animate_1d(run_info[1].z.grid, L_upar[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="L_upar", ax=ax[1])
+        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
+                   frame_index=frame_index, xlabel="z", ylabel="values",
+                   label="mfp", ax=ax[1])
+        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+               orientation=:horizontal)
+        outfile = variable_prefix * "_vs_z." * input.animation_ext
+        save_animation(fig, frame_index, nt, outfile)
+    end
+end
+
 # Utility functions
 ###################
 #
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index d2a1a60e1..6b91ff72f 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4145,6 +4145,82 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
     if variable_name == "temperature"
         vth = get_variable(run_info, "thermal_speed"; kwargs...)
         variable = vth.^2
+    elseif variable_name == "dT_dz"
+        T = get_variable(run_info, "temperature"; kwargs...)
+        variable = similar(T)
+        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, 3)
+                @views derivative!(variable[:,:,it], T[:,:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 4), ir ∈ 1:run_info.r.n
+                @views derivative!(variable[:,:,ir,it], T[:,:,ir,it], run_info.z, run_info.z_spectral)
+            end
+        end
+    elseif variable_name == "dn_dz"
+        n = get_variable(run_info, "density"; kwargs...)
+        variable = similar(n)
+        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, 3)
+                @views derivative!(variable[:,:,it], n[:,:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 4), ir ∈ 1:run_info.r.n
+                @views derivative!(variable[:,:,ir,it], n[:,:,ir,it], run_info.z, run_info.z_spectral)
+            end
+        end
+    elseif variable_name == "dupar_dz"
+        upar = get_variable(run_info, "parallel_flow"; kwargs...)
+        variable = similar(upar)
+        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, 3)
+                @views derivative!(variable[:,:,it], upar[:,:,it], run_info.z, run_info.z_spectral)
+            end
+        else
+            for it ∈ 1:size(variable, 4), 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 == "mfp"
+        vth = get_variable(run_info, "thermal_speed"; kwargs...)
+        nu_ii = get_variable(run_info, "collision_frequency_ii"; kwargs...)
+        variable = vth ./ nu_ii
+    elseif variable_name == "L_T"
+        dT_dz = get_variable(run_info, "dT_dz"; kwargs...)
+        temp = get_variable(run_info, "temperature"; kwargs...)
+        # We define gradient lengthscale of T as LT^-1 = dln(T)/dz (ignore negative sign
+        # tokamak convention as we're only concerned with comparing magnitudes)
+        variable = abs.(temp .* dT_dz.^(-1))
+        # flat points in temperature have diverging LT, so ignore those with NaN
+        # using a hard coded 10.0 tolerance for now
+        variable[variable .> 10.0] .= NaN
+    elseif variable_name == "L_n"
+        dn_dz = get_variable(run_info, "dn_dz"; kwargs...)
+        n = get_variable(run_info, "density"; kwargs...)
+        # We define gradient lengthscale of n as Ln^-1 = dln(n)/dz (ignore negative sign
+        # tokamak convention as we're only concerned with comparing magnitudes)
+        variable = abs.(n .* dn_dz.^(-1))
+        # flat points in temperature have diverging Ln, so ignore those with NaN
+        # using a hard coded 10.0 tolerance for now
+        variable[variable .> 10.0] .= NaN
+    elseif variable_name == "L_upar"
+        dupar_dz = get_variable(run_info, "dupar_dz"; kwargs...)
+        upar = get_variable(run_info, "parallel_flow"; kwargs...)
+        # We define gradient lengthscale of upar as Lupar^-1 = dln(upar)/dz (ignore negative sign
+        # tokamak convention as we're only concerned with comparing magnitudes)
+        variable = abs.(upar .* dupar_dz.^(-1))
+        # flat points in temperature have diverging Lupar, so ignore those with NaN
+        # using a hard coded 10.0 tolerance for now
+        variable[variable .> 10.0] .= NaN
     elseif variable_name == "collision_frequency_ii"
         n = get_variable(run_info, "density"; kwargs...)
         vth = get_variable(run_info, "thermal_speed"; kwargs...)
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 1cfe5dc77..783a02549 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -768,7 +768,7 @@ function update_ion_qpar_species!(qpar, density, upar, vth, dT_dz, ff, vpa, vper
         calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                      evolve_upar, evolve_ppar)
     elseif ion_physics == braginskii_ions
-        calculate_ion_qpar_from_braginskii!(qpar, density, vth, dT_dz, z, r, collisions, evolve_density, 
+        calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, 
                                             evolve_upar, evolve_ppar)
     else
         throw(ArgumentError("ion model $ion_physics not implemented for qpar calculation"))
@@ -820,7 +820,7 @@ end
 """
 calculate parallel heat flux if ion composition flag is Braginskii fluid ions
 """
-function calculate_ion_qpar_from_braginskii!(qpar, density, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
+function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
     # Note that this is a braginskii heat flux for ions using the krook operator. The full Fokker-Planck operator
     # Braginskii heat flux is different! This also assumes one ion species, and so no friction between ions.
     @boundscheck r.n == size(qpar, 2) || throw(BoundsError(qpar))
@@ -834,7 +834,50 @@ function calculate_ion_qpar_from_braginskii!(qpar, density, vth, dT_dz, z, r, co
             qpar[iz,ir] = -(1/2) * 5/4 * density[iz,ir] * vth[iz,ir]^2 /nu_ii * dT_dz[iz,ir,1]
         end
     else
-        throw(ArgumentError("Braginskii heat flux simulation requires evolve_density, evolve_upar and evolve_ppar to be true, since it is a purely fluid simulation"))
+        throw(ArgumentError("Braginskii heat flux simulation requires evolve_density, 
+              evolve_upar and evolve_ppar to be true, since it is a purely fluid simulation"))
+    end
+
+    # add boundary condition to the heat flux, since now there is no distribution function 
+    # (in this case shape function) whose cutoff boundary condition can hold the parallel heat
+    # flux in check. See Stangeby textbook, equations (2.92) and (2.93), and the paragraph between.
+
+    if z.bc == "periodic"
+        # There's no wall boundary condition here, do nothing (qpar can be what it wants)
+        return nothing
+    end
+
+    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
+            this_ppar = vth[iz,ir]^2 * density[iz,ir]/2.0
+            this_upar = upar[iz,ir]
+            this_dens = density[iz,ir]
+            particle_flux = this_dens * this_upar
+            T_i = vth[iz,ir]^2
+
+            # Stangeby paragraph after (2.92)
+            gamma_i = 2.0
+
+            # Stangeby (2.92)
+            total_heat_flux = gamma_i * T_i * particle_flux
+
+            # E.g. Helander&Sigmar (2.14) ??????? what is it for ions?
+            conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar
+
+            qpar[iz,ir] = conductive_heat_flux
+        end
     end
     return nothing
 end

From 9a1e30a90dca422230eb5398d79e9de7437da4b2 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 23 Sep 2024 11:41:58 +0100
Subject: [PATCH 10/31] Add comparison plots functionality for
 collisionality_plots, and add small correction to ion braginskii heat flux
 boundary condition

---
 .../src/makie_post_processing.jl              | 444 +++++++++++-------
 moment_kinetics/src/velocity_moments.jl       |   8 +-
 2 files changed, 279 insertions(+), 173 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 9c092286c..28b1d687f 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
@@ -825,6 +825,8 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         animate_Ln_mfp_vs_z = false,
         plot_Lupar_mfp_vs_z = false,
         animate_Lupar_mfp_vs_z = false,
+        plot_Lupar_Ln_LT_mfp_vs_z = false,
+        animate_Lupar_Ln_LT_mfp_vs_z = false,
         animation_ext = "gif"
        )
 
@@ -8156,193 +8158,293 @@ function collisionality_plots(run_info, plot_prefix=nothing)
     # plot the mean free path of the ions at every point in z, and on the same graph
     # plot the lengthscales of the gradients of all moments at each point in z
     #println("run_info is ", run_info)
-    println("Making plots for collisionality")
+
     if !isa(run_info, Tuple)
         run_info = (run_info,)
     end
-
     input = Dict_to_NamedTuple(input_dict["collisionality_plots"])
-    temperature_input = Dict_to_NamedTuple(input_dict["temperature"])
-
-    mfp = get_variable(run_info, "mfp")
-    L_T = get_variable(run_info, "L_T")
-    L_n = get_variable(run_info, "L_n")
-    L_upar = get_variable(run_info, "L_upar")
-
-    if input.plot_mfp_vs_z
-        variable_name = "mean_free_path"
-        variable = mfp
-        variable_prefix = plot_prefix * variable_name
-        plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
-                  outfile=variable_prefix * "_vs_z.pdf")
-    end
-
-    if input.animate_mfp_vs_z
-        variable_name = "mean_free_path"
-        variable = mfp
-        variable_prefix = plot_prefix * variable_name
-        animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
-                        outfile=variable_prefix * "_vs_z." * input.animation_ext)
-    end
-    # TASKS: plot mean free path vs z, plot collision frequency vs z, make sure it agrees
-    # with original plot for it, and plot thermal speed and make sure that agrees, and then
-    # make sure that the mean free path then makes sense.
-    # Then plot the temperature and temperature gradient on the same plot
-    # Eventually plot the upar and dupar_dz and density and ddens_dz on plots too. Again,
-    # make sure they agree with the original plots.
-    if input.plot_dT_dz_vs_z
-        variable_name = "dT_dz"
-        variable = nothing
-        try
-            variable = get_variable(run_info, variable_name)
-        catch e
-            return makie_post_processing_error_handler(
-                    e,
-                    "collisionality_plots () failed for $variable_name - could not load data.")
+
+    if input.plot
+        println("Making plots for collisionality")
+
+        temperature_input = Dict_to_NamedTuple(input_dict["temperature"])
+        mfp = get_variable(run_info, "mfp")
+        L_T = get_variable(run_info, "L_T")
+        L_n = get_variable(run_info, "L_n")
+        L_upar = get_variable(run_info, "L_upar")
+
+        if input.plot_mfp_vs_z
+            variable_name = "mean_free_path"
+            variable = mfp
+            variable_prefix = plot_prefix * variable_name
+            plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                    outfile=variable_prefix * "_vs_z.pdf")
         end
-        #println("variable is ", variable)
-        variable_prefix = plot_prefix * variable_name
-        plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
-                  outfile=variable_prefix * "_vs_z.pdf")
-    end
 
-    if input.animate_dT_dz_vs_z
-        variable_name = "dT_dz"
-        variable = nothing
-        try
-            variable = get_variable(run_info, variable_name)
-        catch e
-            return makie_post_processing_error_handler(
-                    e,
-                    "collisionality_plots () failed for $variable_name - could not load data.")
-        end
-        #println("variable is ", variable)
-        variable_prefix = plot_prefix * variable_name
-        animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
-                        outfile=variable_prefix * "_vs_z." * input.animation_ext)
-    end
-
-    if input.plot_nu_ii_vth_mfp_vs_z
-        vth = get_variable(run_info, "thermal_speed")
-        nu_ii = get_variable(run_info, "collision_frequency_ii")
-        variable_prefix = plot_prefix * "checking_mfp_vth_and_nu_ii"
-        #println("vth[1] is ", vth[1])
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        plot_1d(run_info[1].z.grid, vth[1][:,1,1,21], xlabel="z",
-                ylabel="values", label="vth", ax=ax[1], title = "checking_mfp_vth")
-        plot_1d(run_info[1].z.grid, nu_ii[1][:,1,1,21], label="nu_ii", ax=ax[1])
-        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        #plot_1d(z.grid, select_slice(error, :z; input=input), xlabel=L"z",
-        #        ylabel=norm_label, ax=ax[2])
-        outfile = variable_prefix * "_vs_z.pdf"
-        save(outfile, fig)
-    end
+        if input.animate_mfp_vs_z
+            variable_name = "mean_free_path"
+            variable = mfp
+            variable_prefix = plot_prefix * variable_name
+            animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                            outfile=variable_prefix * "_vs_z." * input.animation_ext)
+        end
 
-    if input.plot_LT_dT_dz_temp_vs_z
-        dT_dz = get_variable(run_info, "dT_dz")
-        temp = get_variable(run_info, "temperature")
-        variable_prefix = plot_prefix * "LT_dT_dz_temp"
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        plot_1d(run_info[1].z.grid, L_T[1][:,1,1,21], xlabel="z",
-                ylabel="values", label="L_T", ax=ax[1], title = "LT_dT_dz_temp")
-        plot_1d(run_info[1].z.grid, dT_dz[1][:,1,1,21], label="dT_dz", ax=ax[1])
-        plot_1d(run_info[1].z.grid, temp[1][:,1,1,21], label="temp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z.pdf"
-        save(outfile, fig)
-    end
+        if input.plot_dT_dz_vs_z
+            variable_name = "dT_dz"
+            variable = nothing
+            try
+                variable = get_variable(run_info, variable_name)
+            catch e
+                return makie_post_processing_error_handler(
+                        e,
+                        "collisionality_plots () failed for $variable_name - could not load data.")
+            end
 
-    if input.plot_LT_mfp_vs_z
-        variable_prefix = plot_prefix * "LT_mfp_first_comparison"
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        plot_1d(run_info[1].z.grid, L_T[1][:,1,1,21], xlabel="z",
-                ylabel="values", label="L_T", ax=ax[1], title = "L_T and mean free path comparison")
-        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z.pdf"
-        save(outfile, fig)
-    end
+            variable_prefix = plot_prefix * variable_name
+            plot_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                    outfile=variable_prefix * "_vs_z.pdf")
+        end
 
-    if input.animate_LT_mfp_vs_z
-        nt = length(mfp[1][1,1,1,:])
-        variable_prefix = plot_prefix * "LT_mfp_first_comparison"
+        if input.animate_dT_dz_vs_z
+            variable_name = "dT_dz"
+            variable = nothing
+            try
+                variable = get_variable(run_info, variable_name)
+            catch e
+                return makie_post_processing_error_handler(
+                        e,
+                        "collisionality_plots () failed for $variable_name - could not load data.")
+            end
 
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        frame_index = Observable(1)
-        animate_1d(run_info[1].z.grid, L_T[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="L_T", ax=ax[1])
-        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z." * input.animation_ext
-        save_animation(fig, frame_index, nt, outfile)
-    end
+            variable_prefix = plot_prefix * variable_name
+            animate_vs_z(run_info, variable_name, is=1, data=variable, input=temperature_input,
+                            outfile=variable_prefix * "_vs_z." * input.animation_ext)
+        end
 
-    if input.plot_Ln_mfp_vs_z
-        variable_prefix = plot_prefix * "Ln_mfp_first_comparison"
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        plot_1d(run_info[1].z.grid, L_n[1][:,1,1,21], xlabel="z",
-                ylabel="values", label="L_n", ax=ax[1], title = "Ln and mean free path comparison")
-        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z.pdf"
-        save(outfile, fig)
-    end
+        if input.plot_nu_ii_vth_mfp_vs_z
+            vth = get_variable(run_info, "thermal_speed")
+            nu_ii = get_variable(run_info, "collision_frequency_ii")
+            variable_prefix = plot_prefix * "checking_mfp_vth_and_nu_ii"
 
-    if input.animate_Ln_mfp_vs_z
-        nt = length(mfp[1][1,1,1,:])
-        variable_prefix = plot_prefix * "Ln_mfp_first_comparison"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, vth[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"vth", ax=ax[1], title = "checking_mfp_vth")
+                plot_1d(run_info[ri].z.grid, nu_ii[ri][:,1,1,21], label=run_label*"nu_ii", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
 
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        frame_index = Observable(1)
-        animate_1d(run_info[1].z.grid, L_n[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="L_n", ax=ax[1])
-        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z." * input.animation_ext
-        save_animation(fig, frame_index, nt, outfile)
-    end
+        if input.plot_LT_dT_dz_temp_vs_z
+            dT_dz = get_variable(run_info, "dT_dz")
+            temp = get_variable(run_info, "temperature")
+            variable_prefix = plot_prefix * "LT_dT_dz_temp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"L_T", ax=ax[1], title = "LT_dT_dz_temp")
+                plot_1d(run_info[ri].z.grid, dT_dz[ri][:,1,1,21], label=run_label*"dT_dz", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, temp[ri][:,1,1,21], label=run_label*"temp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
 
-    if input.plot_Lupar_mfp_vs_z
-        variable_prefix = plot_prefix * "Lupar_mfp_first_comparison"
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        plot_1d(run_info[1].z.grid, L_upar[1][:,1,1,21], xlabel="z",
-                ylabel="values", label="L_upar", ax=ax[1], title = "Lupar and mean free path comparison")
-        plot_1d(run_info[1].z.grid, mfp[1][:,1,1,21], label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z.pdf"
-        save(outfile, fig)
-    end
+        if input.plot_LT_mfp_vs_z
+            variable_prefix = plot_prefix * "LT_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"L_T", ax=ax[1], title = "L_T and mean free path comparison")
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
+
+        if input.animate_LT_mfp_vs_z
+            variable_prefix = plot_prefix * "LT_mfp"
+            nt = length(mfp[1][1,1,1,:])
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            frame_index = Observable(1)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                animate_1d(run_info[ri].z.grid, L_T[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_T", ax=ax[1], title = "L_T and mean free path comparison")
+                animate_1d(run_info[ri].z.grid, mfp[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z." * input.animation_ext
+            save_animation(fig, frame_index, nt, outfile)
+        end
 
-    if input.animate_Lupar_mfp_vs_z
-        nt = length(mfp[1][1,1,1,:])
-        variable_prefix = plot_prefix * "Lupar_mfp_first_comparison"
+        if input.plot_Ln_mfp_vs_z
+            variable_prefix = plot_prefix * "Ln_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"L_n", ax=ax[1], title = "Ln and mean free path comparison")
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
 
-        fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-        frame_index = Observable(1)
-        animate_1d(run_info[1].z.grid, L_upar[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="L_upar", ax=ax[1])
-        animate_1d(run_info[1].z.grid, mfp[1][:,1,1,:],
-                   frame_index=frame_index, xlabel="z", ylabel="values",
-                   label="mfp", ax=ax[1])
-        Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-               orientation=:horizontal)
-        outfile = variable_prefix * "_vs_z." * input.animation_ext
-        save_animation(fig, frame_index, nt, outfile)
+        if input.animate_Ln_mfp_vs_z
+            variable_prefix = plot_prefix * "Ln_mfp"
+            nt = length(mfp[1][1,1,1,:])
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            frame_index = Observable(1)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                animate_1d(run_info[ri].z.grid, L_n[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_n", ax=ax[1], title = "L_n and mean free path comparison")
+                animate_1d(run_info[ri].z.grid, mfp[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z." * input.animation_ext
+            save_animation(fig, frame_index, nt, outfile)
+        end
+
+        if input.plot_Lupar_mfp_vs_z
+            variable_prefix = plot_prefix * "Lupar_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"L_upar", ax=ax[1], title = "Lupar and mean free path comparison")
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
+
+        if input.animate_Lupar_mfp_vs_z
+            variable_prefix = plot_prefix * "Lupar_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            nt = length(mfp[1][1,1,1,:])
+            frame_index = Observable(1)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                animate_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_upar", ax=ax[1], title = "L_upar and mean free path comparison")
+                animate_1d(run_info[ri].z.grid, mfp[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z." * input.animation_ext
+            save_animation(fig, frame_index, nt, outfile)
+        end
+
+        if input.plot_Lupar_Ln_LT_mfp_vs_z
+            variable_prefix = plot_prefix * "Lupar_Ln_LT_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,21], xlabel="z",
+                        ylabel="values", label=run_label*"L_upar", ax=ax[1], title = "Lupar Ln LT and mean free path comparison")
+                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,21], label=run_label*"L_n", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], label=run_label*"L_T", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
+
+        if input.animate_Lupar_Ln_LT_mfp_vs_z
+            nt = length(mfp[1][1,1,1,:])
+            variable_prefix = plot_prefix * "Lupar_Ln_LT_mfp"
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            frame_index = Observable(1)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                animate_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_upar", ax=ax[1], title = "Lupar Ln LT and mean free path comparison")
+                animate_1d(run_info[ri].z.grid, L_n[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_n", ax=ax[1])
+                animate_1d(run_info[ri].z.grid, L_T[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"L_T", ax=ax[1])
+                animate_1d(run_info[ri].z.grid, mfp[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"mfp", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z." * input.animation_ext
+            save_animation(fig, frame_index, nt, outfile)
+        end
     end
 end
 
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 783a02549..e47e115b8 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -873,8 +873,12 @@ function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z,
             # Stangeby (2.92)
             total_heat_flux = gamma_i * T_i * particle_flux
 
-            # E.g. Helander&Sigmar (2.14) ??????? what is it for ions?
-            conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar
+            # E.g. Helander&Sigmar (2.14) ??????? what is it for ions? From back
+            # of the envelope calculation, ignoring viscosity means what we ignored 
+            # from the electrons was their mean flow contribution, but here it does matter
+            # (I don't have access to the book right now)
+            conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar - 
+                                                     0.5 * this_dens * this_upar^3
 
             qpar[iz,ir] = conductive_heat_flux
         end

From 0e28693bb523edea6d1d3fc6c49119a4d43995ed Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 23 Sep 2024 15:00:10 +0100
Subject: [PATCH 11/31] Add comments to boundary condition of braginskii heat
 flux, and add docstring for collisionality_plots function in
 makie_post_processing

---
 .../src/makie_post_processing.jl                     | 12 +++++++-----
 moment_kinetics/src/velocity_moments.jl              |  2 ++
 2 files changed, 9 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 28b1d687f..73f9ff4c0 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
@@ -8152,13 +8152,15 @@ function timestep_diagnostics(run_info, run_info_dfns; plot_prefix=nothing, it=n
     return steps_fig, dt_fig, CFL_fig
 end
 
+"""
+A function to plot collisionalities. The mean free path is plotted (or animated) 
+along with the lengthscales of the gradients of density, parallel flow and temperature.
 
-
+There are also functions to check the calculations of the mean free path and the 
+comparison of temperature, L_T and dT_dz. They would only be for making sure
+lengthscales and mean free path calculations are sensible.
+"""
 function collisionality_plots(run_info, plot_prefix=nothing)
-    # plot the mean free path of the ions at every point in z, and on the same graph
-    # plot the lengthscales of the gradients of all moments at each point in z
-    #println("run_info is ", run_info)
-
     if !isa(run_info, Tuple)
         run_info = (run_info,)
     end
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index e47e115b8..c1b077b1a 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -877,6 +877,8 @@ function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z,
             # of the envelope calculation, ignoring viscosity means what we ignored 
             # from the electrons was their mean flow contribution, but here it does matter
             # (I don't have access to the book right now)
+            # I can now see in the book that I was pretty much right here, though
+            # I need to consider viscosity (in 1D, is it 0?)
             conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar - 
                                                      0.5 * this_dens * this_upar^3
 

From 5ef19bd10290fedab3182f264e881f8babe2437b Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 23 Sep 2024 16:03:16 +0100
Subject: [PATCH 12/31] Fix bug during restarts to allow for multiple sources
 while external_source_controller_integral is still unused

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

diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 6b91ff72f..c2004e3f7 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -779,7 +779,8 @@ function reload_evolving_fields!(pdf, moments, fields, boundary_distributions,
                     length(moments.ion.external_source_controller_integral) == 1
                 moments.ion.external_source_controller_integral .=
                     load_slice(dynamic, "external_source_controller_integral", time_index)
-            elseif length(moments.ion.external_source_controller_integral) > 1
+            elseif size(moments.ion.external_source_controller_integral)[1] > 1 ||
+                    size(moments.ion.external_source_controller_integral)[2] > 1 
                 moments.ion.external_source_controller_integral .=
                     reload_moment("external_source_controller_integral", dynamic,
                                   time_index, r, z, r_range, z_range, restart_r,

From 4213c524727b7d78f5876f96aa35047dec204e1b Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 23 Sep 2024 16:06:14 +0100
Subject: [PATCH 13/31] Make collisionality_plots plot the last timestep value
 (which was originally hard coded)

---
 .../src/makie_post_processing.jl              | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 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 73f9ff4c0..3c41d31e6 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
@@ -8235,10 +8235,10 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, vth[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, vth[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"vth", ax=ax[1], title = "checking_mfp_vth")
-                plot_1d(run_info[ri].z.grid, nu_ii[ri][:,1,1,21], label=run_label*"nu_ii", ax=ax[1])
-                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, nu_ii[ri][:,1,1,end], label=run_label*"nu_ii", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)
@@ -8257,10 +8257,10 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"L_T", ax=ax[1], title = "LT_dT_dz_temp")
-                plot_1d(run_info[ri].z.grid, dT_dz[ri][:,1,1,21], label=run_label*"dT_dz", ax=ax[1])
-                plot_1d(run_info[ri].z.grid, temp[ri][:,1,1,21], label=run_label*"temp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, dT_dz[ri][:,1,1,end], label=run_label*"dT_dz", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, temp[ri][:,1,1,end], label=run_label*"temp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)
@@ -8277,9 +8277,9 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"L_T", ax=ax[1], title = "L_T and mean free path comparison")
-                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)
@@ -8320,9 +8320,9 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"L_n", ax=ax[1], title = "Ln and mean free path comparison")
-                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)
@@ -8363,9 +8363,9 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"L_upar", ax=ax[1], title = "Lupar and mean free path comparison")
-                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)
@@ -8406,11 +8406,11 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,21], xlabel="z",
+                plot_1d(run_info[ri].z.grid, L_upar[ri][:,1,1,end], xlabel="z",
                         ylabel="values", label=run_label*"L_upar", ax=ax[1], title = "Lupar Ln LT and mean free path comparison")
-                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,21], label=run_label*"L_n", ax=ax[1])
-                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,21], label=run_label*"L_T", ax=ax[1])
-                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,21], label=run_label*"mfp", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, L_n[ri][:,1,1,end], label=run_label*"L_n", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, L_T[ri][:,1,1,end], label=run_label*"L_T", ax=ax[1])
+                plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
                 orientation=:vertical)

From 701ab843b740e194c3e7dee1920a491eb4b5ebd0 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 27 Sep 2024 12:00:41 +0100
Subject: [PATCH 14/31] Add option of overlaying what the Braginskii heat flux
 would be for that run (for when heat flux in the simulation is calculated
 using the shape function - having this option for braginskii_ions would just
 plot the same thing twice)

---
 .../src/makie_post_processing.jl              | 49 +++++++++++++++++++
 moment_kinetics/src/load_data.jl              |  6 +++
 2 files changed, 55 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 6c542662f..d97e8e5fc 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
@@ -830,6 +830,8 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         animate_Lupar_mfp_vs_z = false,
         plot_Lupar_Ln_LT_mfp_vs_z = false,
         animate_Lupar_Ln_LT_mfp_vs_z = false,
+        plot_overlay_braginskii_heat_flux = false,
+        animate_overlay_braginskii_heat_flux = false,
         animation_ext = "gif"
        )
 
@@ -8628,6 +8630,53 @@ function collisionality_plots(run_info, plot_prefix=nothing)
             outfile = variable_prefix * "_vs_z." * input.animation_ext
             save_animation(fig, frame_index, nt, outfile)
         end
+
+        if input.plot_overlay_braginskii_heat_flux
+            variable_prefix = plot_prefix * "braginskii_vs_original_heat_flux"
+            braginskii_q = get_variable(run_info, "braginskii_heat_flux")
+            original_q = get_variable(run_info, "parallel_heat_flux")
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                plot_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,end], xlabel="z",
+                        ylabel="values", label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                plot_1d(run_info[ri].z.grid, original_q[ri][:,1,1,end], label=run_label*"original_q", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z.pdf"
+            save(outfile, fig)
+        end
+
+        if input.animate_overlay_braginskii_heat_flux
+            nt = length(mfp[1][1,1,1,:])
+            variable_prefix = plot_prefix * "braginskii_vs_original_heat_flux"
+            braginskii_q = get_variable(run_info, "braginskii_heat_flux")
+            original_q = get_variable(run_info, "parallel_heat_flux")
+            fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
+            frame_index = Observable(1)
+            for ri ∈ eachindex(run_info)
+                if length(run_info) > 1
+                    run_label = run_info[ri].run_name * " "
+                else
+                    run_label = " "
+                end
+                animate_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                animate_1d(run_info[ri].z.grid, original_q[ri][:,1,1,:],
+                        frame_index=frame_index, xlabel="z", ylabel="values",
+                        label=run_label*"original_q", ax=ax[1])
+            end
+            Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
+                orientation=:vertical)
+            outfile = variable_prefix * "_vs_z." * input.animation_ext
+            save_animation(fig, frame_index, nt, outfile)
+        end
     end
 end
 
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index c2004e3f7..db117ccd1 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4222,6 +4222,12 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         # flat points in temperature have diverging Lupar, so ignore those with NaN
         # using a hard coded 10.0 tolerance for now
         variable[variable .> 10.0] .= NaN
+    elseif variable_name == "braginskii_heat_flux"
+        n = get_variable(run_info, "density"; kwargs...)
+        vth = get_variable(run_info, "thermal_speed"; kwargs...)
+        dT_dz = get_variable(run_info, "dT_dz"; kwargs...)
+        nu_ii = get_variable(run_info, "collision_frequency_ii"; kwargs...)
+        variable = @. -(1/2) * 5/4 * n * vth^2 * dT_dz / nu_ii
     elseif variable_name == "collision_frequency_ii"
         n = get_variable(run_info, "density"; kwargs...)
         vth = get_variable(run_info, "thermal_speed"; kwargs...)

From 0e6d395b5190d8e6d248b90d8cfbf4b215ccd888 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 27 Sep 2024 12:13:12 +0100
Subject: [PATCH 15/31] Add option for super Gaussian with decay of 4th power
 instead of second power in upstream external source. Might add higher powers
 later.

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

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index e2551e6d7..8e5e71e2e 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -409,6 +409,9 @@ function get_source_profile(profile_type, width, relative_minimum, coord)
         x = coord.grid
         return @. (1.0 - relative_minimum) * exp(-(x-x[1]) / width) + relative_minimum +
                   (1.0 - relative_minimum) * exp(-(x[end]-x) / width) + relative_minimum
+    elseif profile_type == "super_gaussian_4"
+        x = coord.grid
+        return @. (1.0 - relative_minimum) * exp(-(x / width)^4) + relative_minimum
     else
         error("Unrecognised source profile type '$profile_type'.")
     end

From 9d92ca23748148d8ab94c794f64cb0c0554b0996 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Sat, 28 Sep 2024 15:47:20 +0100
Subject: [PATCH 16/31] Fix typo when energy source used

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

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index c345491ca..55c22710f 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -471,7 +471,7 @@ function initialize_external_source_amplitude!(moments, external_source_settings
                 if moments.evolve_ppar
                     @loop_r_z ir iz begin
                         moments.ion.external_source_pressure_amplitude[iz,ir,index] =
-                            (0.5 * ion_source.source_T +
+                            (0.5 * ion_source_settings[index].source_T +
                             moments.ion.upar[iz,ir]^2 - moments.ion.ppar[iz,ir]) *
                             ion_source_settings[index].source_strength *
                             ion_source_settings[index].r_amplitude[ir] *

From fab3f8a78d602d6c70a142235787ffd2c53bbd26 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 30 Sep 2024 10:31:38 +0100
Subject: [PATCH 17/31] Don't plot braginskii overlay if the original
 simulation itself was a braginskii one (which just overplots the old one,
 apart from the boundaries)

---
 .../src/makie_post_processing.jl               | 18 +++++++++++-------
 1 file changed, 11 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 d97e8e5fc..0d3cb658d 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
@@ -8596,7 +8596,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 plot_1d(run_info[ri].z.grid, mfp[ri][:,1,1,end], label=run_label*"mfp", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
-                orientation=:vertical)
+                orientation=:horizontal)
             outfile = variable_prefix * "_vs_z.pdf"
             save(outfile, fig)
         end
@@ -8642,8 +8642,10 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                plot_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,end], xlabel="z",
-                        ylabel="values", label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
+                    plot_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,end], xlabel="z",
+                            ylabel="values", label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                end
                 plot_1d(run_info[ri].z.grid, original_q[ri][:,1,1,end], label=run_label*"original_q", ax=ax[1])
             end
             Legend(legend_place[1], ax[1]; tellheight=true, tellwidth=false,
@@ -8652,7 +8654,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
             save(outfile, fig)
         end
 
-        if input.animate_overlay_braginskii_heat_flux
+        if input.animate_overlay_braginskii_heat_flux 
             nt = length(mfp[1][1,1,1,:])
             variable_prefix = plot_prefix * "braginskii_vs_original_heat_flux"
             braginskii_q = get_variable(run_info, "braginskii_heat_flux")
@@ -8665,9 +8667,11 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                animate_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,:],
-                        frame_index=frame_index, xlabel="z", ylabel="values",
-                        label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
+                    animate_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,:],
+                            frame_index=frame_index, xlabel="z", ylabel="values",
+                            label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                end
                 animate_1d(run_info[ri].z.grid, original_q[ri][:,1,1,:],
                         frame_index=frame_index, xlabel="z", ylabel="values",
                         label=run_label*"original_q", ax=ax[1])

From 8e5ecd9be11e06da2d7eb32cd4f0c9e59b40005e Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Tue, 1 Oct 2024 16:45:39 +0100
Subject: [PATCH 18/31] Add string to braginskii_overlay system to tell you its
 an overlay, which is much clearer than just writing braginskii everywhere.
 Note that I will soon change the naming system to collisional_krook

---
 .../makie_post_processing/src/makie_post_processing.jl       | 5 +++--
 1 file changed, 3 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 0d3cb658d..fabee921e 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
@@ -8358,6 +8358,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
         L_n = get_variable(run_info, "L_n")
         L_upar = get_variable(run_info, "L_upar")
 
+        # write function to check that mfp[ri][1, 1, 1, :] is the same length (i.e. nt) for all ri 
         if input.plot_mfp_vs_z
             variable_name = "mean_free_path"
             variable = mfp
@@ -8644,7 +8645,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 end
                 if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
                     plot_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,end], xlabel="z",
-                            ylabel="values", label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                            ylabel="values", label=run_label*"braginskii_q_overlay", ax=ax[1], title = "Braginskii heat flux overlay")
                 end
                 plot_1d(run_info[ri].z.grid, original_q[ri][:,1,1,end], label=run_label*"original_q", ax=ax[1])
             end
@@ -8670,7 +8671,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
                     animate_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,:],
                             frame_index=frame_index, xlabel="z", ylabel="values",
-                            label=run_label*"braginskii_q", ax=ax[1], title = "Braginskii heat flux overlay")
+                            label=run_label*"braginskii_q_overlay", ax=ax[1], title = "Braginskii heat flux overlay")
                 end
                 animate_1d(run_info[ri].z.grid, original_q[ri][:,1,1,:],
                         frame_index=frame_index, xlabel="z", ylabel="values",

From 8309e996e766039f96ec9a5944420d03850efa40 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 3 Oct 2024 10:59:25 +0100
Subject: [PATCH 19/31] Fix PI_controller file io for  option

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

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 9b0fafd54..d336a2713 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -1425,8 +1425,10 @@ function define_dynamic_ion_moment_variables!(fid, n_ion_species, r::coordinate,
                     parallel_io=parallel_io,
                     description="Integral term for the PID controller of the external source for ions")
             else
+                r_midpoint = (name="midpoint_controller_r", n=1)
+                z_midpoint = (name="midpoint_controller_z", n=1)
                 external_source_controller_integral = create_dynamic_variable!(
-                    dynamic, "external_source_controller_integral", mk_float;
+                    dynamic, "external_source_controller_integral", mk_float, r_midpoint, z_midpoint, n_sources;
                     parallel_io=parallel_io,
                     description="Integral term for the PID controller of the external source for ions")
             end
@@ -2581,10 +2583,12 @@ function write_ion_moments_data_to_binary(scratch, moments, n_ion_species, t_par
             end
         end
         if io_moments.external_source_controller_integral !== nothing
+            n_sources = size(moments.ion.external_source_amplitude)[3]
             if size(moments.ion.external_source_controller_integral) == (1,1, n_sources)
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
                                       moments.ion.external_source_controller_integral,
-                                      t_idx, parallel_io)
+                                      t_idx, parallel_io, 1, 1, n_sources)
+                println("moments.ion.external_source_controller_integral = ", moments.ion.external_source_controller_integral)
             else
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
                                       moments.ion.external_source_controller_integral,

From 0062ca8f44e072e2212f19bf7d7b744c4461fd0b Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 3 Oct 2024 11:02:19 +0100
Subject: [PATCH 20/31] Edit PI_controller for density_midpoint_control option
 to begin the source with an input guess, which is the source_strength option
 in the input file. This means when you restart a run from another that had a
 constant external source, that choice of starting point for the PI_controller
 source should mean the PI_controller doesn't have to do too much work (as
 long as the target amplitude is not very far from the current midpoint
 amplitude).

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

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 55c22710f..fa295cc68 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -719,7 +719,7 @@ function external_ion_source!(pdf, fvec, moments, ion_source, index, vperp, vpa,
     end
     vpa_grid = vpa.grid
     vperp_grid = vperp.grid
-    if source_type in ("Maxwellian","energy")
+    if source_type in ("Maxwellian","energy","density_midpoint_control","density_profile_control")
         begin_s_r_z_vperp_region()
         if moments.evolve_ppar && moments.evolve_upar && moments.evolve_density
             vth = moments.ion.vth
@@ -1180,7 +1180,7 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
                     dt * ion_source_settings.PI_density_controller_I * n_error
 
                 # Only want a source, so never allow amplitude to be negative
-                amplitude = max(
+                amplitude = max(ion_source_settings.source_strength +
                     ion_source_settings.PI_density_controller_P * n_error +
                     ion_moments.external_source_controller_integral[1,1,index],
                     0)

From 139462670f0ae150e49bfeb00bac967a5e96d659 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 3 Oct 2024 11:25:08 +0100
Subject: [PATCH 21/31] remove printing function that was used for debugging

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

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index d336a2713..757c91a73 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -2588,7 +2588,6 @@ function write_ion_moments_data_to_binary(scratch, moments, n_ion_species, t_par
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
                                       moments.ion.external_source_controller_integral,
                                       t_idx, parallel_io, 1, 1, n_sources)
-                println("moments.ion.external_source_controller_integral = ", moments.ion.external_source_controller_integral)
             else
                 append_to_dynamic_var(io_moments.external_source_controller_integral,
                                       moments.ion.external_source_controller_integral,

From 22f91bdfd4b01c3fe84b7a570d156fc474f79e49 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 3 Oct 2024 14:21:23 +0100
Subject: [PATCH 22/31] Change naming from Braginskii_ions to coll_krook_ions,
 and change instances of the name braginskii for ion functions to coll_krook.
 The electron braginskii fluid naming is unchanged for now.

---
 .../src/makie_post_processing.jl              | 28 +++++++++----------
 moment_kinetics/src/initial_conditions.jl     |  2 +-
 moment_kinetics/src/input_structs.jl          |  6 ++--
 moment_kinetics/src/load_data.jl              |  2 +-
 moment_kinetics/src/species_input.jl          |  4 +--
 moment_kinetics/src/velocity_moments.jl       | 12 ++++----
 6 files changed, 27 insertions(+), 27 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 fabee921e..53902ed2c 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
@@ -830,8 +830,8 @@ function _setup_single_input!(this_input_dict::OrderedDict{String,Any},
         animate_Lupar_mfp_vs_z = false,
         plot_Lupar_Ln_LT_mfp_vs_z = false,
         animate_Lupar_Ln_LT_mfp_vs_z = false,
-        plot_overlay_braginskii_heat_flux = false,
-        animate_overlay_braginskii_heat_flux = false,
+        plot_overlay_coll_krook_heat_flux = false,
+        animate_overlay_coll_krook_heat_flux = false,
         animation_ext = "gif"
        )
 
@@ -8632,9 +8632,9 @@ function collisionality_plots(run_info, plot_prefix=nothing)
             save_animation(fig, frame_index, nt, outfile)
         end
 
-        if input.plot_overlay_braginskii_heat_flux
-            variable_prefix = plot_prefix * "braginskii_vs_original_heat_flux"
-            braginskii_q = get_variable(run_info, "braginskii_heat_flux")
+        if input.plot_overlay_coll_krook_heat_flux
+            variable_prefix = plot_prefix * "coll_krook_vs_original_heat_flux"
+            coll_krook_q = get_variable(run_info, "coll_krook_heat_flux")
             original_q = get_variable(run_info, "parallel_heat_flux")
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
             for ri ∈ eachindex(run_info)
@@ -8643,9 +8643,9 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
-                    plot_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,end], xlabel="z",
-                            ylabel="values", label=run_label*"braginskii_q_overlay", ax=ax[1], title = "Braginskii heat flux overlay")
+                if get(run_info[ri].input["composition"], "ion_physics", "") !== coll_krook_ions
+                    plot_1d(run_info[ri].z.grid, coll_krook_q[ri][:,1,1,end], xlabel="z",
+                            ylabel="values", label=run_label*"coll_krook_q_overlay", ax=ax[1], title = "coll_krook heat flux overlay")
                 end
                 plot_1d(run_info[ri].z.grid, original_q[ri][:,1,1,end], label=run_label*"original_q", ax=ax[1])
             end
@@ -8655,10 +8655,10 @@ function collisionality_plots(run_info, plot_prefix=nothing)
             save(outfile, fig)
         end
 
-        if input.animate_overlay_braginskii_heat_flux 
+        if input.animate_overlay_coll_krook_heat_flux 
             nt = length(mfp[1][1,1,1,:])
-            variable_prefix = plot_prefix * "braginskii_vs_original_heat_flux"
-            braginskii_q = get_variable(run_info, "braginskii_heat_flux")
+            variable_prefix = plot_prefix * "coll_krook_vs_original_heat_flux"
+            coll_krook_q = get_variable(run_info, "coll_krook_heat_flux")
             original_q = get_variable(run_info, "parallel_heat_flux")
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
             frame_index = Observable(1)
@@ -8668,10 +8668,10 @@ function collisionality_plots(run_info, plot_prefix=nothing)
                 else
                     run_label = " "
                 end
-                if get(run_info[ri].input["composition"], "ion_physics", "") !== braginskii_ions
-                    animate_1d(run_info[ri].z.grid, braginskii_q[ri][:,1,1,:],
+                if get(run_info[ri].input["composition"], "ion_physics", "") !== coll_krook_ions
+                    animate_1d(run_info[ri].z.grid, coll_krook_q[ri][:,1,1,:],
                             frame_index=frame_index, xlabel="z", ylabel="values",
-                            label=run_label*"braginskii_q_overlay", ax=ax[1], title = "Braginskii heat flux overlay")
+                            label=run_label*"coll_krook_q_overlay", ax=ax[1], title = "coll_krook heat flux overlay")
                 end
                 animate_1d(run_info[ri].z.grid, original_q[ri][:,1,1,:],
                         frame_index=frame_index, xlabel="z", ylabel="values",
diff --git a/moment_kinetics/src/initial_conditions.jl b/moment_kinetics/src/initial_conditions.jl
index a142a65b8..3d357232b 100644
--- a/moment_kinetics/src/initial_conditions.jl
+++ b/moment_kinetics/src/initial_conditions.jl
@@ -200,7 +200,7 @@ function init_pdf_and_moments!(pdf, moments, fields, boundary_distributions, geo
                         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. Even if Braginskii fluid is being
+        # calculate the initial parallel heat flux from the initial un-normalised pdf. Even if coll_krook fluid is being
         # advanced, initialised ion_qpar uses the pdf 
         update_ion_qpar!(moments.ion.qpar, moments.ion.qpar_updated,
                      moments.ion.dens, moments.ion.upar, moments.ion.vth, moments.ion.dT_dz,
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index a30dffa0e..d898908c7 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -163,12 +163,12 @@ export kinetic_electrons_with_temperature_equation
 @enum ion_physics_type begin
     gyrokinetic_ions
     drift_kinetic_ions
-    braginskii_ions
+    coll_krook_ions
 end
 export ion_physics_type
 export gyrokinetic_ions
 export drift_kinetic_ions
-export braginskii_ions
+export coll_krook_ions
 """
 """
 Base.@kwdef struct spatial_initial_condition_input
@@ -268,7 +268,7 @@ Base.@kwdef struct species_composition
     #   density is fixed to be Nₑ*(eϕ/T_e) and N_e is calculated using a current
     #   condition at the wall
     electron_physics::electron_physics_type
-    # ion physics can be drift_kinetic_ions, gyrokinetic_ions and braginskii_ions
+    # ion physics can be drift_kinetic_ions, gyrokinetic_ions and coll_krook_ions
     # gyrokinetic_ions (originally gyrokinetic_ions = true) -> use gyroaveraged fields at fixed guiding 
     # centre and moments of the pdf computed at fixed r
     # drift_kinetic_ions (originally gyrokinetic_ions = false) -> use drift kinetic approximation
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index db117ccd1..6f61f1649 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -4222,7 +4222,7 @@ function get_variable(run_info, variable_name; normalize_advection_speed_shape=t
         # flat points in temperature have diverging Lupar, so ignore those with NaN
         # using a hard coded 10.0 tolerance for now
         variable[variable .> 10.0] .= NaN
-    elseif variable_name == "braginskii_heat_flux"
+    elseif variable_name == "coll_krook_heat_flux"
         n = get_variable(run_info, "density"; kwargs...)
         vth = get_variable(run_info, "thermal_speed"; kwargs...)
         dT_dz = get_variable(run_info, "dT_dz"; kwargs...)
diff --git a/moment_kinetics/src/species_input.jl b/moment_kinetics/src/species_input.jl
index d2c8fb415..a4b1ef641 100644
--- a/moment_kinetics/src/species_input.jl
+++ b/moment_kinetics/src/species_input.jl
@@ -35,9 +35,9 @@ function get_species_input(toml_input)
         # If ion_physics=gyrokinetic_ions, the ion distribution function is
         # advanced in time using gyroaveraged fields at fixed guiding centre and moments of the
         # pdf computed at fixed r
-        # If ion_physics=braginskii_ions, there is no need for a shape function to evolve, and the code 
+        # If ion_physics=coll_krook_ions, there is no need for a shape function to evolve, and the code 
         # only evolves ions in a fluid sense (i.e. all evolve_moments are set to true), with a 
-        # Braginskii closure for the ion heat flux.
+        # coll_krook closure for the ion heat flux.
         ion_physics = drift_kinetic_ions,
         # initial Tₑ = 1
         T_e = 1.0,
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index c1b077b1a..a8416af16 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -767,8 +767,8 @@ function update_ion_qpar_species!(qpar, density, upar, vth, dT_dz, ff, vpa, vper
     if ion_physics ∈ (drift_kinetic_ions, gyrokinetic_ions)
         calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                      evolve_upar, evolve_ppar)
-    elseif ion_physics == braginskii_ions
-        calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, 
+    elseif ion_physics == coll_krook_ions
+        calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, 
                                             evolve_upar, evolve_ppar)
     else
         throw(ArgumentError("ion model $ion_physics not implemented for qpar calculation"))
@@ -818,15 +818,15 @@ function calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp,
     return nothing
 end
 """
-calculate parallel heat flux if ion composition flag is Braginskii fluid ions
+calculate parallel heat flux if ion composition flag is coll_krook fluid ions
 """
-function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
+function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
     # Note that this is a braginskii heat flux for ions using the krook operator. The full Fokker-Planck operator
     # Braginskii heat flux is different! This also assumes one ion species, and so no friction between ions.
     @boundscheck r.n == size(qpar, 2) || throw(BoundsError(qpar))
     @boundscheck z.n == size(qpar, 1) || throw(BoundsError(qpar))
 
-    # calculate braginskii heat flux. Currently only works for one ion species! (hence the 1 in dT_dz[iz,ir,1])
+    # calculate coll_krook heat flux. Currently only works for one ion species! (hence the 1 in dT_dz[iz,ir,1])
     if evolve_density && evolve_upar && evolve_ppar
         begin_r_z_region()
         @loop_r_z ir iz begin
@@ -834,7 +834,7 @@ function calculate_ion_qpar_from_braginskii!(qpar, density, upar, vth, dT_dz, z,
             qpar[iz,ir] = -(1/2) * 5/4 * density[iz,ir] * vth[iz,ir]^2 /nu_ii * dT_dz[iz,ir,1]
         end
     else
-        throw(ArgumentError("Braginskii heat flux simulation requires evolve_density, 
+        throw(ArgumentError("coll_krook heat flux simulation requires evolve_density, 
               evolve_upar and evolve_ppar to be true, since it is a purely fluid simulation"))
     end
 

From 0462f32d7b0682f1935fd19c85c6326ea6e58e8d Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Thu, 3 Oct 2024 17:58:42 +0100
Subject: [PATCH 23/31] Don't evolve shape function if heat flux closure comes
 from coll_krook coefficient calculation.

---
 moment_kinetics/src/time_advance.jl | 302 ++++++++++++++--------------
 1 file changed, 152 insertions(+), 150 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index fa9f73bca..09f56c170 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -3224,167 +3224,169 @@ function euler_time_advance!(fvec_out, fvec_in, pdf, fields, moments,
                                             external_source_settings.neutral, r, z, dt)
     end
 
-    if advance.vpa_advection
-        vpa_advection!(fvec_out.pdf, fvec_in, fields, moments, vpa_advect, vpa, vperp, z, r, dt, t,
-            vpa_spectral, composition, collisions, external_source_settings.ion, geometry)
-    end
-
-    # z_advection! advances 1D advection equation in z
-    # 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,
-                     dt, t, z_spectral, composition, geometry, scratch_dummy)
-    end
+    if composition.ion_physics ∈ (drift_kinetic_ions, gyrokinetic_ions)
+        if advance.vpa_advection
+            vpa_advection!(fvec_out.pdf, fvec_in, fields, moments, vpa_advect, vpa, vperp, z, r, dt, t,
+                vpa_spectral, composition, collisions, external_source_settings.ion, geometry)
+        end
 
-    # r advection relies on derivatives in z to get ExB
-    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
-    # vperp_advection requires information about z and r advection
-    # 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,
-                      moments, fields, t)
-    end
+        # z_advection! advances 1D advection equation in z
+        # apply z-advection operation to ion species
 
-    if advance.source_terms
-        source_terms!(fvec_out.pdf, fvec_in, moments, vpa, z, r, dt, z_spectral,
-                      composition, collisions, external_source_settings.ion)
-    end
+        if advance.z_advection
+            z_advection!(fvec_out.pdf, fvec_in, moments, fields, z_advect, z, vpa, vperp, r,
+                        dt, t, z_spectral, composition, geometry, scratch_dummy)
+        end
 
-    if advance.neutral_z_advection
-        neutral_advection_z!(fvec_out.pdf_neutral, fvec_in, moments, neutral_z_advect,
-            r, z, vzeta, vr, vz, dt, t, z_spectral, composition, scratch_dummy)
-    end
+        # r advection relies on derivatives in z to get ExB
+        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
+        # vperp_advection requires information about z and r advection
+        # 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,
+                        moments, fields, t)
+        end
 
-    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
+        if advance.source_terms
+            source_terms!(fvec_out.pdf, fvec_in, moments, vpa, z, r, dt, z_spectral,
+                        composition, collisions, external_source_settings.ion)
+        end
 
-    if advance.neutral_vz_advection
-        neutral_advection_vz!(fvec_out.pdf_neutral, fvec_in, fields, moments,
-                              neutral_vz_advect, vz, vr, vzeta, z, r, dt, vz_spectral,
-                              composition, collisions, external_source_settings.neutral)
-    end
+        if advance.neutral_z_advection
+            neutral_advection_z!(fvec_out.pdf_neutral, fvec_in, moments, neutral_z_advect,
+                r, z, vzeta, vr, vz, dt, t, z_spectral, composition, scratch_dummy)
+        end
 
-    if advance.neutral_source_terms
-        source_terms_neutral!(fvec_out.pdf_neutral, fvec_in, moments, vpa, z, r, dt, z_spectral,
-                      composition, collisions, external_source_settings.neutral)
-    end
+        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
 
-    if advance.manufactured_solns_test
-        source_terms_manufactured!(fvec_out.pdf, fvec_out.pdf_neutral, vz, vr, vzeta, vpa, vperp, z, r, t, dt, composition, manufactured_source_list)
-    end
+        if advance.neutral_vz_advection
+            neutral_advection_vz!(fvec_out.pdf_neutral, fvec_in, fields, moments,
+                                neutral_vz_advect, vz, vr, vzeta, z, r, dt, vz_spectral,
+                                composition, collisions, external_source_settings.neutral)
+        end
 
-    if advance.ion_cx_collisions || advance.ion_ionization_collisions
-        # 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)
-    end
-    if advance.neutral_cx_collisions || advance.neutral_ionization_collisions
-        # 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
+        if advance.neutral_source_terms
+            source_terms_neutral!(fvec_out.pdf_neutral, fvec_in, moments, vpa, z, r, dt, z_spectral,
+                        composition, collisions, external_source_settings.neutral)
+        end
 
-    # account for charge exchange collisions between ions and neutrals
-    if advance.ion_cx_collisions_1V
-        ion_charge_exchange_collisions_1V!(fvec_out.pdf, fvec_in, moments, composition,
-                                           vpa, vz, collisions.reactions.charge_exchange_frequency,
-                                           vpa_spectral, vz_spectral, dt)
-    elseif advance.ion_cx_collisions
-        ion_charge_exchange_collisions_3V!(fvec_out.pdf, pdf.ion.buffer, fvec_in,
-                                           composition, vz, vr, vzeta, vpa, vperp, z, r,
-                                           collisions.reactions.charge_exchange_frequency, dt)
-    end
-    if advance.neutral_cx_collisions_1V
-        neutral_charge_exchange_collisions_1V!(fvec_out.pdf_neutral, fvec_in, moments,
-                                               composition, vpa, vz,
-                                               collisions.reactions.charge_exchange_frequency, vpa_spectral,
-                                               vz_spectral, dt)
-    elseif advance.neutral_cx_collisions
-        neutral_charge_exchange_collisions_3V!(fvec_out.pdf_neutral, pdf.neutral.buffer,
-                                               fvec_in, composition, vz, vr, vzeta, vpa,
-                                               vperp, z, r, collisions.reactions.charge_exchange_frequency,
-                                               dt)
-    end
-    # account for ionization collisions between ions and neutrals
-    if advance.ion_ionization_collisions_1V
-        ion_ionization_collisions_1V!(fvec_out.pdf, fvec_in, vz, vpa, vperp, z, r,
-                                      vz_spectral, moments, composition, collisions, dt)
-    elseif advance.ion_ionization_collisions
-        ion_ionization_collisions_3V!(fvec_out.pdf, pdf.ion.buffer, fvec_in, composition,
-                                      vz, vr, vzeta, vpa, vperp, z, r, collisions, dt)
-    end
-    if advance.neutral_ionization_collisions_1V
-        neutral_ionization_collisions_1V!(fvec_out.pdf_neutral, fvec_in, vz, vpa, vperp,
-                                          z, r, vz_spectral, moments, composition,
-                                          collisions, dt)
-    elseif advance.neutral_ionization_collisions
-        neutral_ionization_collisions_3V!(fvec_out.pdf_neutral, fvec_in, composition, vz,
-                                          vr, vzeta, vpa, vperp, z, r, collisions, dt)
-    end
+        if advance.manufactured_solns_test
+            source_terms_manufactured!(fvec_out.pdf, fvec_out.pdf_neutral, vz, vr, vzeta, vpa, vperp, z, r, t, dt, composition, manufactured_source_list)
+        end
 
-    # Add Krook collision operator for ions
-    if advance.krook_collisions_ii
-        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.ion_cx_collisions || advance.ion_ionization_collisions
+            # 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)
+        end
+        if advance.neutral_cx_collisions || advance.neutral_ionization_collisions
+            # 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
 
-    if advance.external_source
-        total_external_ion_sources!(fvec_out.pdf, fvec_in, moments, external_source_settings.ion,
-                            vperp, vpa, dt, scratch_dummy)
-    end
-    if advance.neutral_external_source
-        total_external_neutral_sources!(fvec_out.pdf_neutral, fvec_in, moments,
-                                external_source_settings.neutral, vzeta, vr, vz, dt)
-    end
-
-    # add numerical dissipation
-    if advance.ion_numerical_dissipation
-        vpa_dissipation!(fvec_out.pdf, fvec_in.pdf, vpa, vpa_spectral, dt,
-                         num_diss_params.ion.vpa_dissipation_coefficient)
-        vperp_dissipation!(fvec_out.pdf, fvec_in.pdf, vperp, vperp_spectral, dt,
-                         num_diss_params.ion.vperp_dissipation_coefficient)
-        z_dissipation!(fvec_out.pdf, fvec_in.pdf, z, z_spectral, dt,
-                       num_diss_params.ion.z_dissipation_coefficient, scratch_dummy)
-        r_dissipation!(fvec_out.pdf, fvec_in.pdf, r, r_spectral, dt,
-                       num_diss_params.ion.r_dissipation_coefficient, scratch_dummy)
-    end
-    if advance.neutral_numerical_dissipation
-        vz_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, vz,
-                                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.neutral.z_dissipation_coefficient, scratch_dummy)
-        r_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, r, r_spectral,
-                               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
-        update_entropy_diagnostic = (istage == 1)
-        if collisions.fkpl.self_collisions
-            # self collisions for each species
-            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
-        if collisions.fkpl.slowing_down_test
-        # include cross-collsions with fixed Maxwellian backgrounds
-            explicit_fp_collisions_weak_form_Maxwellian_cross_species!(fvec_out.pdf,fvec_in.pdf,moments.ion.dSdt,
-                             composition,collisions,dt,fp_arrays,r,z,vperp,vpa,vperp_spectral,vpa_spectral,
-                                             diagnose_entropy_production = update_entropy_diagnostic)
+        # account for charge exchange collisions between ions and neutrals
+        if advance.ion_cx_collisions_1V
+            ion_charge_exchange_collisions_1V!(fvec_out.pdf, fvec_in, moments, composition,
+                                            vpa, vz, collisions.reactions.charge_exchange_frequency,
+                                            vpa_spectral, vz_spectral, dt)
+        elseif advance.ion_cx_collisions
+            ion_charge_exchange_collisions_3V!(fvec_out.pdf, pdf.ion.buffer, fvec_in,
+                                            composition, vz, vr, vzeta, vpa, vperp, z, r,
+                                            collisions.reactions.charge_exchange_frequency, dt)
+        end
+        if advance.neutral_cx_collisions_1V
+            neutral_charge_exchange_collisions_1V!(fvec_out.pdf_neutral, fvec_in, moments,
+                                                composition, vpa, vz,
+                                                collisions.reactions.charge_exchange_frequency, vpa_spectral,
+                                                vz_spectral, dt)
+        elseif advance.neutral_cx_collisions
+            neutral_charge_exchange_collisions_3V!(fvec_out.pdf_neutral, pdf.neutral.buffer,
+                                                fvec_in, composition, vz, vr, vzeta, vpa,
+                                                vperp, z, r, collisions.reactions.charge_exchange_frequency,
+                                                dt)
+        end
+        # account for ionization collisions between ions and neutrals
+        if advance.ion_ionization_collisions_1V
+            ion_ionization_collisions_1V!(fvec_out.pdf, fvec_in, vz, vpa, vperp, z, r,
+                                        vz_spectral, moments, composition, collisions, dt)
+        elseif advance.ion_ionization_collisions
+            ion_ionization_collisions_3V!(fvec_out.pdf, pdf.ion.buffer, fvec_in, composition,
+                                        vz, vr, vzeta, vpa, vperp, z, r, collisions, dt)
+        end
+        if advance.neutral_ionization_collisions_1V
+            neutral_ionization_collisions_1V!(fvec_out.pdf_neutral, fvec_in, vz, vpa, vperp,
+                                            z, r, vz_spectral, moments, composition,
+                                            collisions, dt)
+        elseif advance.neutral_ionization_collisions
+            neutral_ionization_collisions_3V!(fvec_out.pdf_neutral, fvec_in, composition, vz,
+                                            vr, vzeta, vpa, vperp, z, r, collisions, dt)
+        end
+
+        # Add Krook collision operator for ions
+        if advance.krook_collisions_ii
+            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
+            total_external_ion_sources!(fvec_out.pdf, fvec_in, moments, external_source_settings.ion,
+                                vperp, vpa, dt, scratch_dummy)
+        end
+        if advance.neutral_external_source
+            total_external_neutral_sources!(fvec_out.pdf_neutral, fvec_in, moments,
+                                    external_source_settings.neutral, vzeta, vr, vz, dt)
+        end
+
+        # add numerical dissipation
+        if advance.ion_numerical_dissipation
+            vpa_dissipation!(fvec_out.pdf, fvec_in.pdf, vpa, vpa_spectral, dt,
+                            num_diss_params.ion.vpa_dissipation_coefficient)
+            vperp_dissipation!(fvec_out.pdf, fvec_in.pdf, vperp, vperp_spectral, dt,
+                            num_diss_params.ion.vperp_dissipation_coefficient)
+            z_dissipation!(fvec_out.pdf, fvec_in.pdf, z, z_spectral, dt,
+                        num_diss_params.ion.z_dissipation_coefficient, scratch_dummy)
+            r_dissipation!(fvec_out.pdf, fvec_in.pdf, r, r_spectral, dt,
+                        num_diss_params.ion.r_dissipation_coefficient, scratch_dummy)
+        end
+        if advance.neutral_numerical_dissipation
+            vz_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, vz,
+                                    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.neutral.z_dissipation_coefficient, scratch_dummy)
+            r_dissipation_neutral!(fvec_out.pdf_neutral, fvec_in.pdf_neutral, r, r_spectral,
+                                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
+            update_entropy_diagnostic = (istage == 1)
+            if collisions.fkpl.self_collisions
+                # self collisions for each species
+                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
+            if collisions.fkpl.slowing_down_test
+            # include cross-collsions with fixed Maxwellian backgrounds
+                explicit_fp_collisions_weak_form_Maxwellian_cross_species!(fvec_out.pdf,fvec_in.pdf,moments.ion.dSdt,
+                                composition,collisions,dt,fp_arrays,r,z,vperp,vpa,vperp_spectral,vpa_spectral,
+                                                diagnose_entropy_production = update_entropy_diagnostic)
+            end
         end
     end
     

From 73488093920f967ba0204ed12b093f89fe611d79 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 4 Oct 2024 10:56:17 +0100
Subject: [PATCH 24/31] Allow collisionality_plots to animate comparison runs
 with different numbers of timesteps

---
 .../src/makie_post_processing.jl                    | 13 +++++++------
 1 file changed, 7 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 53902ed2c..f69d76b9e 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
@@ -8357,6 +8357,12 @@ function collisionality_plots(run_info, plot_prefix=nothing)
         L_T = get_variable(run_info, "L_T")
         L_n = get_variable(run_info, "L_n")
         L_upar = get_variable(run_info, "L_upar")
+        nt = minimum(length(mfp[ri][1,1,1,:]) for ri in eachindex(run_info))
+        # print warning if the lengths of all the mfp[ri][1,1,1,:] are not the same
+        if any(length(mfp[ri][1,1,1,:]) != nt for ri in eachindex(run_info))
+            println("Warning: The number of timesteps of some simulations are different, " *
+                    "so only the first common timesteps will be animated.")
+        end
 
         # write function to check that mfp[ri][1, 1, 1, :] is the same length (i.e. nt) for all ri 
         if input.plot_mfp_vs_z
@@ -8473,7 +8479,6 @@ function collisionality_plots(run_info, plot_prefix=nothing)
 
         if input.animate_LT_mfp_vs_z
             variable_prefix = plot_prefix * "LT_mfp"
-            nt = length(mfp[1][1,1,1,:])
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
             frame_index = Observable(1)
             for ri ∈ eachindex(run_info)
@@ -8516,7 +8521,6 @@ function collisionality_plots(run_info, plot_prefix=nothing)
 
         if input.animate_Ln_mfp_vs_z
             variable_prefix = plot_prefix * "Ln_mfp"
-            nt = length(mfp[1][1,1,1,:])
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
             frame_index = Observable(1)
             for ri ∈ eachindex(run_info)
@@ -8560,7 +8564,6 @@ function collisionality_plots(run_info, plot_prefix=nothing)
         if input.animate_Lupar_mfp_vs_z
             variable_prefix = plot_prefix * "Lupar_mfp"
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
-            nt = length(mfp[1][1,1,1,:])
             frame_index = Observable(1)
             for ri ∈ eachindex(run_info)
                 if length(run_info) > 1
@@ -8603,7 +8606,6 @@ function collisionality_plots(run_info, plot_prefix=nothing)
         end
 
         if input.animate_Lupar_Ln_LT_mfp_vs_z
-            nt = length(mfp[1][1,1,1,:])
             variable_prefix = plot_prefix * "Lupar_Ln_LT_mfp"
             fig, ax, legend_place = get_1d_ax(1; get_legend_place=:below)
             frame_index = Observable(1)
@@ -8655,8 +8657,7 @@ function collisionality_plots(run_info, plot_prefix=nothing)
             save(outfile, fig)
         end
 
-        if input.animate_overlay_coll_krook_heat_flux 
-            nt = length(mfp[1][1,1,1,:])
+        if input.animate_overlay_coll_krook_heat_flux
             variable_prefix = plot_prefix * "coll_krook_vs_original_heat_flux"
             coll_krook_q = get_variable(run_info, "coll_krook_heat_flux")
             original_q = get_variable(run_info, "parallel_heat_flux")

From dc627cf81fbab4f0b3f0e5bc0e7fbae512a32778 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Fri, 4 Oct 2024 12:34:49 +0100
Subject: [PATCH 25/31] add neutral controller_integral file_io.jl bug fix (ion
 version was fixed in previous commit)

---
 moment_kinetics/src/file_io.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/file_io.jl b/moment_kinetics/src/file_io.jl
index 757c91a73..9a493c340 100644
--- a/moment_kinetics/src/file_io.jl
+++ b/moment_kinetics/src/file_io.jl
@@ -2796,14 +2796,15 @@ function write_neutral_moments_data_to_binary(scratch, moments, n_neutral_specie
             end
         end
         if io_moments.external_source_neutral_controller_integral !== nothing
-            if size(moments.neutral.external_source_neutral_controller_integral) == (1,1)
+            n_sources = size(moments.neutral.external_source_amplitude)[3]
+            if size(moments.neutral.external_source_neutral_controller_integral) == (1,1, n_sources)
                 append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
                                       moments.neutral.external_source_controller_integral,
-                                      t_idx, parallel_io)
+                                      t_idx, parallel_io, 1, 1, n_sources)
             else
                 append_to_dynamic_var(io_moments.external_source_neutral_controller_integral,
                                       moments.neutral.external_source_controller_integral,
-                                      t_idx, parallel_io, z, r)
+                                      t_idx, parallel_io, z, r, n_sources)
             end
         end
         if moments.evolve_density || moments.evolve_upar || moments.evolve_ppar

From 44ee25a945758059079c771db018593dda30f492 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Sun, 6 Oct 2024 20:39:17 +0100
Subject: [PATCH 26/31] skip chodura condition check in coll_krook fluid ion
 case, so that ngrid and nelement can both be 1 (to speed up process easily).
 Also change gamma_i to 4.0 for stability of more dense cases in
 collisionality scans.

---
 moment_kinetics/src/time_advance.jl     | 9 +++++++--
 moment_kinetics/src/velocity_moments.jl | 7 +++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/moment_kinetics/src/time_advance.jl b/moment_kinetics/src/time_advance.jl
index 09f56c170..37b7db23a 100644
--- a/moment_kinetics/src/time_advance.jl
+++ b/moment_kinetics/src/time_advance.jl
@@ -2350,8 +2350,13 @@ function apply_all_bcs_constraints_update_moments!(
     # Note these may be needed for the boundary condition on the neutrals, so must be
     # calculated before that is applied. Also may be needed to calculate advection speeds
     # for for CFL stability limit calculations in adaptive_timestep_update!().
-    update_derived_moments!(this_scratch, moments, vpa, vperp, z, r, composition,
-        r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, diagnostic_moments)
+    if composition.ion_physics ∈ (drift_kinetic_ions, gyrokinetic_ions)
+        update_derived_moments!(this_scratch, moments, vpa, vperp, z, r, composition,
+            r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, diagnostic_moments)
+    else
+        update_derived_moments!(this_scratch, moments, vpa, vperp, z, r, composition,
+            r_spectral, geometry, gyroavs, scratch_dummy, z_advect, collisions, false)
+    end
 
     calculate_ion_moment_derivatives!(moments, this_scratch, scratch_dummy, z, z_spectral,
                                       num_diss_params.ion.moment_dissipation_coefficient)
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index a8416af16..d2a23c81e 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -867,8 +867,11 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
             particle_flux = this_dens * this_upar
             T_i = vth[iz,ir]^2
 
-            # Stangeby paragraph after (2.92)
-            gamma_i = 2.0
+            # Stangeby (2.92) suggests a factor of 5.0 for this gamma_i value, but
+            # then suggests in the next paragraph that a factor of 2.0 is possibly 
+            # more accurate. 2.0 gives unstable results at high densities, so for
+            # now I'm using 4.0.
+            gamma_i = 4.0
 
             # Stangeby (2.92)
             total_heat_flux = gamma_i * T_i * particle_flux

From bc7401d233783b023f492652f47504d06f25c225 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 7 Oct 2024 11:02:44 +0100
Subject: [PATCH 27/31] change gamma back to 2, and change the boundary
 condition for heat flux of coll_krook calculation to reflect the 1D pressure
 rather than 3D isotropic pressure used. Also comment out line where
 boundary_conditions.jl is used in krook_collisions.jl since I need to use
 krook_collisions.jl in the velocity_moments.jl module. THIS IS NOT A
 WORKAROUND AND NEEDS TO BE FIXED AT SOME POINT.

---
 moment_kinetics/src/krook_collisions.jl | 2 +-
 moment_kinetics/src/velocity_moments.jl | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index 766fe7d39..c89ccf2ec 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -7,7 +7,7 @@ export setup_krook_collisions_input, get_collision_frequency_ii, get_collision_f
        add_electron_krook_collisions_to_Jacobian!
 
 using ..looping
-using ..boundary_conditions: skip_f_electron_bc_points_in_Jacobian
+#using ..boundary_conditions: skip_f_electron_bc_points_in_Jacobian
 using ..input_structs: krook_collisions_input, set_defaults_and_check_section!
 using ..reference_parameters: get_reference_collision_frequency_ii,
                               get_reference_collision_frequency_ee,
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index de71273c4..5b18731e2 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -871,7 +871,7 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
             # then suggests in the next paragraph that a factor of 2.0 is possibly 
             # more accurate. 2.0 gives unstable results at high densities, so for
             # now I'm using 4.0.
-            gamma_i = 4.0
+            gamma_i = 2.0
 
             # Stangeby (2.92)
             total_heat_flux = gamma_i * T_i * particle_flux
@@ -882,7 +882,7 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
             # (I don't have access to the book right now)
             # I can now see in the book that I was pretty much right here, though
             # I need to consider viscosity (in 1D, is it 0?)
-            conductive_heat_flux = total_heat_flux - 2.5 * this_ppar * this_upar - 
+            conductive_heat_flux = total_heat_flux - 1.5 * this_ppar * this_upar - 
                                                      0.5 * this_dens * this_upar^3
 
             qpar[iz,ir] = conductive_heat_flux

From 8155690d4fc806e99db0d9982eca69e65fd1c441 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 7 Oct 2024 13:08:46 +0100
Subject: [PATCH 28/31] Create new module for collision_frequency calculations,
 so that krook collisions module can be defined later in moment_kinetics.jl.
 This fixes the issue highlighted in the previous commit.

---
 .../src/makie_post_processing.jl              |   6 +
 moment_kinetics/src/collision_frequencies.jl  | 118 +++++++++++++++++
 .../src/electron_kinetic_equation.jl          |   5 +-
 moment_kinetics/src/krook_collisions.jl       | 119 +-----------------
 moment_kinetics/src/load_data.jl              |   2 +
 moment_kinetics/src/moment_kinetics.jl        |   3 +-
 moment_kinetics/src/velocity_moments.jl       |   2 +-
 7 files changed, 136 insertions(+), 119 deletions(-)
 create mode 100644 moment_kinetics/src/collision_frequencies.jl

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 bf376f465..8c2e72316 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
@@ -258,6 +258,12 @@ function makie_post_process(run_dir::Union{String,Tuple},
             mkpath(comparison_plot_dir)
             plot_prefix = joinpath(comparison_plot_dir, "compare_")
         end
+    else
+        if length(run_info) != 1
+            comparison_plot_dir = "comparison_plots_$plot_prefix"
+            mkpath(comparison_plot_dir)
+            plot_prefix = joinpath(comparison_plot_dir, "compare_")
+        end
     end
 
     timestep_diagnostics(run_info, run_info_dfns; plot_prefix=plot_prefix)
diff --git a/moment_kinetics/src/collision_frequencies.jl b/moment_kinetics/src/collision_frequencies.jl
new file mode 100644
index 000000000..3da8aa5b4
--- /dev/null
+++ b/moment_kinetics/src/collision_frequencies.jl
@@ -0,0 +1,118 @@
+module collision_frequencies
+
+export get_collision_frequency_ii, get_collision_frequency_ee, get_collision_frequency_ei
+
+using ..reference_parameters: get_reference_collision_frequency_ii,
+                              get_reference_collision_frequency_ee,
+                              get_reference_collision_frequency_ei
+using ..reference_parameters: setup_reference_parameters
+
+"""
+    get_collision_frequency_ii(collisions, n, 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_ii(collisions, n, vth)
+    # extract krook options from collisions struct
+    colk = collisions.krook
+    nuii0 = colk.nuii0
+    frequency_option = colk.frequency_option
+    if frequency_option ∈ ("reference_parameters", "collisionality_scan")
+        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 @. 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
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. 0.0 * n
+    else
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_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)
+    # 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 @. 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
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. 0.0 * n
+    else
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_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)
+    # extract krook options from collisions struct
+    colk = collisions.krook
+    nuei0 = colk.nuei0
+    frequency_option = colk.frequency_option
+    if frequency_option == "reference_parameters"
+        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 @. 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
+        # returning a scalar from this branch but an array from the "reference_parameters"
+        # branch).
+        return @. 0.0 * n
+    else
+        error("Unrecognised option [krook_collisions] "
+              * "frequency_option=$(frequency_option)")
+    end
+end
+
+end
\ No newline at end of file
diff --git a/moment_kinetics/src/electron_kinetic_equation.jl b/moment_kinetics/src/electron_kinetic_equation.jl
index 07d8229df..552c037c9 100644
--- a/moment_kinetics/src/electron_kinetic_equation.jl
+++ b/moment_kinetics/src/electron_kinetic_equation.jl
@@ -38,8 +38,9 @@ using ..em_fields: update_phi!
 using ..external_sources: total_external_electron_sources!,
                           add_total_external_electron_source_to_Jacobian!
 using ..file_io: get_electron_io_info, write_electron_state, finish_electron_io
-using ..krook_collisions: electron_krook_collisions!, get_collision_frequency_ee,
-                          get_collision_frequency_ei,
+using ..collision_frequencies: get_collision_frequency_ee,
+                                 get_collision_frequency_ei
+using ..krook_collisions: electron_krook_collisions!,
                           add_electron_krook_collisions_to_Jacobian!
 using ..moment_constraints: hard_force_moment_constraints!,
                             moment_constraints_on_residual!,
diff --git a/moment_kinetics/src/krook_collisions.jl b/moment_kinetics/src/krook_collisions.jl
index c89ccf2ec..8557d3b7b 100644
--- a/moment_kinetics/src/krook_collisions.jl
+++ b/moment_kinetics/src/krook_collisions.jl
@@ -2,18 +2,14 @@
 """
 module krook_collisions
 
-export setup_krook_collisions_input, get_collision_frequency_ii, get_collision_frequency_ee,
-       get_collision_frequency_ei, krook_collisions!, electron_krook_collisions!,
+export setup_krook_collisions_input, krook_collisions!, electron_krook_collisions!,
        add_electron_krook_collisions_to_Jacobian!
 
 using ..looping
-#using ..boundary_conditions: skip_f_electron_bc_points_in_Jacobian
+using ..boundary_conditions: skip_f_electron_bc_points_in_Jacobian
 using ..input_structs: krook_collisions_input, set_defaults_and_check_section!
-using ..reference_parameters: get_reference_collision_frequency_ii,
-                              get_reference_collision_frequency_ee,
-                              get_reference_collision_frequency_ei
-using ..reference_parameters: setup_reference_parameters
-
+using ..collision_frequencies
+using ..reference_parameters
 
 """
 Function for reading Krook collision operator input parameters. 
@@ -68,113 +64,6 @@ function setup_krook_collisions_input(toml_input::Dict)
     return krook_collisions_input(; input...)
 end
 
-"""
-    get_collision_frequency_ii(collisions, n, 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_ii(collisions, n, vth)
-    # extract krook options from collisions struct
-    colk = collisions.krook
-    nuii0 = colk.nuii0
-    frequency_option = colk.frequency_option
-    if frequency_option ∈ ("reference_parameters", "collisionality_scan")
-        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 @. 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
-        # returning a scalar from this branch but an array from the "reference_parameters"
-        # branch).
-        return @. 0.0 * n
-    else
-        error("Unrecognised option [krook_collisions] "
-              * "frequency_option=$(frequency_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)
-    # 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 @. 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
-        # returning a scalar from this branch but an array from the "reference_parameters"
-        # branch).
-        return @. 0.0 * n
-    else
-        error("Unrecognised option [krook_collisions] "
-              * "frequency_option=$(frequency_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)
-    # extract krook options from collisions struct
-    colk = collisions.krook
-    nuei0 = colk.nuei0
-    frequency_option = colk.frequency_option
-    if frequency_option == "reference_parameters"
-        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 @. 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
-        # returning a scalar from this branch but an array from the "reference_parameters"
-        # branch).
-        return @. 0.0 * n
-    else
-        error("Unrecognised option [krook_collisions] "
-              * "frequency_option=$(frequency_option)")
-    end
-end
 
 """
 Add collision operator
diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl
index 22e5442c8..8575cb9d5 100644
--- a/moment_kinetics/src/load_data.jl
+++ b/moment_kinetics/src/load_data.jl
@@ -26,6 +26,8 @@ using ..file_io: check_io_implementation, get_group, get_subgroup_keys, get_vari
 using ..input_structs
 using ..interpolation: interpolate_to_grid_1d!
 using ..krook_collisions
+using ..collision_frequencies: get_collision_frequency_ii, get_collision_frequency_ee,
+                                get_collision_frequency_ei
 using ..looping
 using ..moment_kinetics_input: mk_input
 using ..neutral_vz_advection: update_speed_neutral_vz!
diff --git a/moment_kinetics/src/moment_kinetics.jl b/moment_kinetics/src/moment_kinetics.jl
index 4e14d1588..21f87e3c8 100644
--- a/moment_kinetics/src/moment_kinetics.jl
+++ b/moment_kinetics/src/moment_kinetics.jl
@@ -37,10 +37,11 @@ include("nonlinear_solvers.jl")
 include("file_io.jl")
 include("geo.jl")
 include("gyroaverages.jl")
-include("krook_collisions.jl")
+include("collision_frequencies.jl")
 include("velocity_moments.jl")
 include("velocity_grid_transforms.jl")
 include("boundary_conditions.jl")
+include("krook_collisions.jl")
 include("electron_fluid_equations.jl")
 include("em_fields.jl")
 include("bgk.jl")
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 5b18731e2..97b899c09 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 ..krook_collisions: get_collision_frequency_ii
+using ..collision_frequencies: get_collision_frequency_ii
 using ..input_structs
 using ..moment_kinetics_structs: moments_ion_substruct, moments_electron_substruct,
                                  moments_neutral_substruct

From 66cbcc1b67474c0235c2bc6dab1ebe136105d334 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Mon, 7 Oct 2024 18:23:34 +0100
Subject: [PATCH 29/31] Change gamma_i to 2.5, from Stangeby (25.2)

---
 moment_kinetics/src/velocity_moments.jl | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 97b899c09..6d19e409c 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -867,21 +867,14 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
             particle_flux = this_dens * this_upar
             T_i = vth[iz,ir]^2
 
-            # Stangeby (2.92) suggests a factor of 5.0 for this gamma_i value, but
-            # then suggests in the next paragraph that a factor of 2.0 is possibly 
-            # more accurate. 2.0 gives unstable results at high densities, so for
-            # now I'm using 4.0.
-            gamma_i = 2.0
+            # Stangeby (25.2)
+            gamma_i = 2.5
 
             # Stangeby (2.92)
             total_heat_flux = gamma_i * T_i * particle_flux
 
-            # E.g. Helander&Sigmar (2.14) ??????? what is it for ions? From back
-            # of the envelope calculation, ignoring viscosity means what we ignored 
-            # from the electrons was their mean flow contribution, but here it does matter
-            # (I don't have access to the book right now)
-            # I can now see in the book that I was pretty much right here, though
-            # I need to consider viscosity (in 1D, is it 0?)
+            # E.g. Helander&Sigmar (2.14), but in 1D we have no viscosity and only 3/2
+            # rather than 5/2.
             conductive_heat_flux = total_heat_flux - 1.5 * this_ppar * this_upar - 
                                                      0.5 * this_dens * this_upar^3
 

From eb018728f902a15510598d33c89bfdf16a2261c2 Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Wed, 9 Oct 2024 13:11:32 +0100
Subject: [PATCH 30/31] Add temperature_midpoint_control source option. This
 will change the amplitude of an energy source so that the temperature
 midpoint is kept at a desired level. A few notes on this for now: 1. The
 temperature of the source has to be higher than the desired temperature,
 otherwise it can never be made hot enough, no matter how high the input
 amplitude. 2. The temperature of the source actually needs to be at least
 DOUBLE the desired value, which I think is coming from a factor of 2 error
 somewhere - this should be fixed when I comb through all the stuff I've
 written to iron out these missing factors of 2 one day. 3. This currently
 only really works for the ion fluid case - something has gone wrong with the
 part that actually adds to the pdf. I'll try to fix this soon.

---
 moment_kinetics/src/external_sources.jl | 141 ++++++++++++++++++++++--
 moment_kinetics/src/input_structs.jl    |   7 ++
 2 files changed, 140 insertions(+), 8 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 06b043036..4626cf4e8 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -73,6 +73,9 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
                      PI_density_target_z_profile="constant",
                      PI_density_target_z_width=1.0,
                      PI_density_target_z_relative_minimum=0.0,
+                     PI_temperature_controller_P=0.0,
+                     PI_temperature_controller_I=0.0,
+                     PI_temperature_target_amplitude=1.0,
                      recycling_controller_fraction=0.0,
                     )
 
@@ -103,6 +106,10 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
             PI_density_target_ir = nothing
             PI_density_target_iz = nothing
             PI_density_target_rank = nothing
+            PI_temperature_target = nothing
+            PI_temperature_target_ir = nothing
+            PI_temperature_target_iz = nothing
+            PI_temperature_target_rank = nothing
         elseif input["source_type"] == "density_midpoint_control"
             PI_density_target = input["PI_density_target_amplitude"]
 
@@ -142,6 +149,53 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
             else
                 PI_density_target_rank = nothing
             end
+            PI_temperature_target = nothing
+            PI_temperature_target_ir = nothing
+            PI_temperature_target_iz = nothing
+            PI_temperature_target_rank = nothing
+        elseif input["source_type"] == "temperature_midpoint_control"
+            PI_temperature_target = input["PI_temperature_target_amplitude"]
+            PI_density_target = nothing
+            PI_density_target_ir = nothing
+            PI_density_target_iz = nothing
+            PI_density_target_rank = nothing
+
+            if comm_block[] != MPI.COMM_NULL
+                PI_controller_amplitude = allocate_shared_float(1)
+                controller_source_profile = allocate_shared_float(z.n, r.n)
+            else
+                PI_controller_amplitude = allocate_float(1)
+                controller_source_profile = allocate_float(z.n, r.n)
+            end
+            for ir ∈ 1:r.n, iz ∈ 1:z.n
+                controller_source_profile[iz,ir] = r_amplitude[ir] * z_amplitude[iz]
+            end
+
+            # Find the indices, and process rank of the point at r=0, z=0.
+            # The result of findfirst() will be `nothing` if the point was not found.
+            PI_temperature_target_ir = findfirst(x->abs(x)<1.e-14, r.grid)
+            PI_temperature_target_iz = findfirst(x->abs(x)<1.e-14, z.grid)
+            if block_rank[] == 0
+                # Only need to do communications from the root process of each
+                # shared-memory block
+                if PI_temperature_target_ir !== nothing && PI_temperature_target_iz !== nothing
+                    PI_temperature_target_rank = iblock_index[]
+                else
+                    PI_temperature_target_rank = 0
+                end
+                if comm_inter_block[] != MPI.COMM_NULL
+                    PI_temperature_target_rank = MPI.Allreduce(PI_temperature_target_rank, +,
+                                                           comm_inter_block[])
+                end
+                if PI_temperature_target_rank == 0 && iblock_index[] == 0 &&
+                        (PI_temperature_target_ir === nothing ||
+                         PI_temperature_target_iz === nothing)
+                    error("No grid point with r=0 and z=0 was found for the "
+                          * "'temperature_midpoint' controller.")
+                end
+            else
+                PI_temperature_target_rank = nothing
+            end
         elseif input["source_type"] ∈ ("Maxwellian", "energy", "alphas", "alphas-with-losses", "beam", "beam-with-losses")
             PI_density_target = nothing
             PI_controller_amplitude = nothing
@@ -149,16 +203,22 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
             PI_density_target_ir = nothing
             PI_density_target_iz = nothing
             PI_density_target_rank = nothing
+            PI_temperature_target = nothing
+            PI_temperature_target_ir = nothing
+            PI_temperature_target_iz = nothing
+            PI_temperature_target_rank = nothing
         else
             error("Unrecognised ion source_type=$(input["source_type"])."
                   * "Possible values are: Maxwellian, density_profile_control, "
-                  * "density_midpoint_control, energy, alphas, alphas-with-losses, "
-                  * "beam, beam-with-losses")
+                  * "density_midpoint_control, temperature_midpoint_control, energy, "
+                  * "alphas, alphas-with-losses, beam, beam-with-losses")
         end
         return ion_source_data(; Dict(Symbol(k)=>v for (k,v) ∈ input)..., r_amplitude,
                 z_amplitude=z_amplitude, PI_density_target=PI_density_target,
                 PI_controller_amplitude, controller_source_profile,
-                PI_density_target_ir, PI_density_target_iz, PI_density_target_rank)
+                PI_density_target_ir, PI_density_target_iz, PI_density_target_rank,
+                PI_temperature_target, PI_temperature_target_ir, PI_temperature_target_iz,
+                PI_temperature_target_rank)
     end
 
     function get_settings_neutrals(source_index, active_flag)
@@ -327,7 +387,7 @@ function setup_external_sources!(input_dict, r, z, electron_physics)
                      source_strength=ion_settings.source_strength,
                      source_T=ion_settings.source_T,
                     )
-        if ion_settings.source_type != "energy"
+        if ion_settings.source_type != "energy" && ion_settings.source_type != "temperature_midpoint_control"
             # Need to keep same amplitude for ions and electrons so there is no charge
             # source.
             if input["source_strength"] != ion_settings.source_strength
@@ -449,7 +509,7 @@ function initialize_external_source_amplitude!(moments, external_source_settings
 
     for index ∈ eachindex(ion_source_settings)
         if ion_source_settings[index].active
-            if ion_source_settings[index].source_type == "energy"
+            if ion_source_settings[index].source_type ∈ ("energy", "temperature_midpoint_control")
                 @loop_r_z ir iz begin
                     moments.ion.external_source_amplitude[iz,ir,index] =
                         ion_source_settings[index].source_strength *
@@ -667,7 +727,7 @@ function initialize_external_source_controller_integral!(
         for index ∈ eachindex(ion_source_settings)
             if ion_source_settings[index].active && 
                ion_source_settings[index].PI_density_controller_I != 0.0 && 
-               ion_source_settings[index].source_type ∈ ("density_profile_control", "density_midpoint_control")
+               ion_source_settings[index].source_type ∈ ("density_profile_control", "density_midpoint_control", "temperature_midpoint_control")
                 moments.ion.external_source_controller_integral[:, :, index] .= 0.0
             end
         end
@@ -721,7 +781,7 @@ function external_ion_source!(pdf, fvec, moments, ion_source, index, vperp, vpa,
     end
     vpa_grid = vpa.grid
     vperp_grid = vperp.grid
-    if source_type in ("Maxwellian","energy","density_midpoint_control","density_profile_control")
+    if source_type in ("Maxwellian","energy","density_midpoint_control","density_profile_control","temperature_midpoint_control")
         begin_s_r_z_vperp_region()
         if moments.evolve_ppar && moments.evolve_upar && moments.evolve_density
             vth = moments.ion.vth
@@ -785,7 +845,7 @@ function external_ion_source!(pdf, fvec, moments, ion_source, index, vperp, vpa,
                   * "evolve_upar=$(moments.evolve_upar), evolve_ppar=$(moments.evolve_ppar)")
         end
 
-        if source_type == "energy"
+        if source_type ∈ ("energy", "temperature_midpoint_control")
             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
@@ -1284,6 +1344,71 @@ function external_ion_source_controller!(fvec_in, moments, ion_source_settings,
                     amplitude * ion_source_settings.controller_source_profile[iz,ir]
             end
         end
+    elseif ion_source_settings.source_type == "temperature_midpoint_control"
+        begin_serial_region()
+        temperature = 2 * ppar ./ density
+        # controller_amplitude error is a shared memory Vector of length 1
+        controller_amplitude = ion_source_settings.PI_controller_amplitude
+        @serial_region begin
+            if ion_source_settings.PI_temperature_target_ir !== nothing &&
+                    ion_source_settings.PI_temperature_target_iz !== nothing
+                # This process has the target point
+
+                T_mid = temperature[ion_source_settings.PI_temperature_target_iz,
+                                ion_source_settings.PI_temperature_target_ir, is]
+                T_error = ion_source_settings.PI_temperature_target - T_mid
+
+                ion_moments.external_source_controller_integral[1,1,index] +=
+                    dt * ion_source_settings.PI_temperature_controller_I * T_error
+
+                # Only want a source, so never allow amplitude to be negative
+                amplitude = max(ion_source_settings.source_strength +
+                    ion_source_settings.PI_temperature_controller_P * T_error +
+                    ion_moments.external_source_controller_integral[1,1,index],
+                    0)
+            else
+                amplitude = nothing
+            end
+            controller_amplitude[1] =
+                MPI.Bcast(amplitude, ion_source_settings.PI_temperature_target_rank,
+                          comm_inter_block[])
+        end
+
+        begin_r_z_region()
+
+        amplitude = controller_amplitude[1]
+        #println("amplitude is $amplitude")
+        #@loop_r_z ir iz begin
+        #    ion_moments.external_source_amplitude[iz,ir,index] =
+        #        amplitude * ion_source_settings.controller_source_profile[iz,ir]
+        #end
+        #if moments.evolve_density
+        #    @loop_r_z ir iz begin
+        #            amplitude * ion_source_settings.controller_source_profile[iz,ir]
+        #    end
+        #        ion_moments.external_source_density_amplitude[iz,ir,index] =
+        #end
+        #if moments.evolve_ppar
+        #    @loop_r_z ir iz begin
+        #        ion_moments.external_source_pressure_amplitude[iz,ir,index] =
+        #            (0.5 * ion_source_settings.source_T + upar[iz,ir,is]^2) *
+        #            amplitude * ion_source_settings.controller_source_profile[iz,ir]
+        #    end
+        #end
+        if moments.evolve_upar
+            @loop_r_z ir iz begin
+                ion_moments.external_source_momentum_amplitude[iz,ir,index] =
+                    - density[iz,ir] * upar[iz,ir] * amplitude *
+                    ion_source_settings.controller_source_profile[iz,ir]
+            end
+        end
+        if moments.evolve_ppar
+            @loop_r_z ir iz begin
+                ion_moments.external_source_pressure_amplitude[iz,ir,index] =
+                    (0.5 * ion_source_settings.source_T + upar[iz,ir]^2 - ppar[iz,ir]) *
+                    amplitude * ion_source_settings.controller_source_profile[iz,ir]
+            end
+        end
     elseif ion_source_settings.source_type == "density_profile_control"
         begin_r_z_region()
 
diff --git a/moment_kinetics/src/input_structs.jl b/moment_kinetics/src/input_structs.jl
index 6cf862254..13d6de0cf 100644
--- a/moment_kinetics/src/input_structs.jl
+++ b/moment_kinetics/src/input_structs.jl
@@ -362,6 +362,9 @@ Base.@kwdef struct ion_source_data
     PI_density_target_z_profile::String
     PI_density_target_z_width::mk_float
     PI_density_target_z_relative_minimum::mk_float
+    PI_temperature_controller_P::mk_float
+    PI_temperature_controller_I::mk_float
+    PI_temperature_target_amplitude::mk_float
     recycling_controller_fraction::mk_float
     # r_amplitude through the r coordinate (in 1D this can just be set to 1.0)
     r_amplitude::Vector{mk_float}
@@ -369,12 +372,16 @@ Base.@kwdef struct ion_source_data
     # constant profile, parabolic, etc..
     z_amplitude::Vector{mk_float}
     PI_density_target::Union{mk_float, Nothing, MPISharedArray{mk_float,2}}
+    PI_temperature_target::Union{mk_float, Nothing, MPISharedArray{mk_float,2}}
     PI_controller_amplitude::Union{Nothing, MPISharedArray{mk_float,1}}
     controller_source_profile::Union{Nothing, MPISharedArray{mk_float,2}, Array{mk_float, 2}}
     PI_density_target_ir::Union{mk_int, Nothing}
     PI_density_target_iz::Union{mk_int, Nothing}
     PI_density_target_rank::Union{mk_int, Nothing} #possibly this should have Int64 as well, 
     # in the event that the code is running with mk_int = Int32 but the rank is set to 0::Int64
+    PI_temperature_target_ir::Union{mk_int, Nothing}
+    PI_temperature_target_iz::Union{mk_int, Nothing}
+    PI_temperature_target_rank::Union{mk_int, Nothing}
 end
 
 Base.@kwdef struct electron_source_data

From 499b9ec706305b3d35a490eab3d6d1890ce3f80a Mon Sep 17 00:00:00 2001
From: Lucas McConnell <sjoh5887@ox.ac.uk>
Date: Tue, 15 Oct 2024 10:15:28 +0100
Subject: [PATCH 31/31] Change gamma_i boundary condition for coll_krook heat
 flux closure depending on whether we're 1V or 2V

---
 moment_kinetics/src/external_sources.jl       | 38 ++++++++-----------
 .../src/moment_kinetics_structs.jl            |  2 +-
 moment_kinetics/src/velocity_moments.jl       | 22 ++++++++---
 3 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/moment_kinetics/src/external_sources.jl b/moment_kinetics/src/external_sources.jl
index 6dd44c497..de20e7a44 100644
--- a/moment_kinetics/src/external_sources.jl
+++ b/moment_kinetics/src/external_sources.jl
@@ -1351,7 +1351,7 @@ source amplitude.
         end
     elseif ion_source_settings.source_type == "temperature_midpoint_control"
         begin_serial_region()
-        temperature = 2 * ppar ./ density
+        ion_moments.temp .= 2 .* ppar ./ density
         # controller_amplitude error is a shared memory Vector of length 1
         controller_amplitude = ion_source_settings.PI_controller_amplitude
         @serial_region begin
@@ -1359,7 +1359,7 @@ source amplitude.
                     ion_source_settings.PI_temperature_target_iz !== nothing
                 # This process has the target point
 
-                T_mid = temperature[ion_source_settings.PI_temperature_target_iz,
+                T_mid = ion_moments.temp[ion_source_settings.PI_temperature_target_iz,
                                 ion_source_settings.PI_temperature_target_ir, is]
                 T_error = ion_source_settings.PI_temperature_target - T_mid
 
@@ -1382,24 +1382,11 @@ source amplitude.
         begin_r_z_region()
 
         amplitude = controller_amplitude[1]
-        #println("amplitude is $amplitude")
-        #@loop_r_z ir iz begin
-        #    ion_moments.external_source_amplitude[iz,ir,index] =
-        #        amplitude * ion_source_settings.controller_source_profile[iz,ir]
-        #end
-        #if moments.evolve_density
-        #    @loop_r_z ir iz begin
-        #            amplitude * ion_source_settings.controller_source_profile[iz,ir]
-        #    end
-        #        ion_moments.external_source_density_amplitude[iz,ir,index] =
-        #end
-        #if moments.evolve_ppar
-        #    @loop_r_z ir iz begin
-        #        ion_moments.external_source_pressure_amplitude[iz,ir,index] =
-        #            (0.5 * ion_source_settings.source_T + upar[iz,ir,is]^2) *
-        #            amplitude * ion_source_settings.controller_source_profile[iz,ir]
-        #    end
-        #end
+        @loop_r_z ir iz begin
+            ion_moments.external_source_amplitude[iz,ir,index] =
+                amplitude * ion_source_settings.controller_source_profile[iz,ir]
+        end
+
         if moments.evolve_upar
             @loop_r_z ir iz begin
                 ion_moments.external_source_momentum_amplitude[iz,ir,index] =
@@ -1410,10 +1397,17 @@ source amplitude.
         if moments.evolve_ppar
             @loop_r_z ir iz begin
                 ion_moments.external_source_pressure_amplitude[iz,ir,index] =
-                    (0.5 * ion_source_settings.source_T + upar[iz,ir]^2 - ppar[iz,ir]) *
-                    amplitude * ion_source_settings.controller_source_profile[iz,ir]
+                    ((0.5 * ion_source_settings.source_T + 2 * upar[iz,ir]^2) *
+                    amplitude) * ion_source_settings.controller_source_profile[iz,ir]
             end
         end
+        #if moments.evolve_ppar
+        #    @loop_r_z ir iz begin
+        #        ion_moments.external_source_pressure_amplitude[iz,ir,index] =
+        #            (0.5 * ion_source_settings.source_T + upar[iz,ir]^2 - ppar[iz,ir]) *
+        #            amplitude * ion_source_settings.controller_source_profile[iz,ir]
+        #    end
+        #end
     elseif ion_source_settings.source_type == "density_profile_control"
         begin_r_z_region()
 
diff --git a/moment_kinetics/src/moment_kinetics_structs.jl b/moment_kinetics/src/moment_kinetics_structs.jl
index 65ebf341e..4df2043db 100644
--- a/moment_kinetics/src/moment_kinetics_structs.jl
+++ b/moment_kinetics/src/moment_kinetics_structs.jl
@@ -95,7 +95,7 @@ struct moments_ion_substruct
     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}
-    # this is the temperature, calculated from 2ppar/dens (the comment below may be out of date?)
+    # this is the temperature, calculated from 2ppar/dens (the comment above may be out of date?)
     temp::MPISharedArray{mk_float,3}
     # generalised Chodura integrals for the lower and upper plates
     chodura_integral_lower::MPISharedArray{mk_float,2}
diff --git a/moment_kinetics/src/velocity_moments.jl b/moment_kinetics/src/velocity_moments.jl
index 6d19e409c..c58be5cbf 100644
--- a/moment_kinetics/src/velocity_moments.jl
+++ b/moment_kinetics/src/velocity_moments.jl
@@ -768,7 +768,7 @@ function update_ion_qpar_species!(qpar, density, upar, vth, dT_dz, ff, vpa, vper
         calculate_ion_qpar_from_pdf!(qpar, density, upar, vth, ff, vpa, vperp, z, r, evolve_density,
                                      evolve_upar, evolve_ppar)
     elseif ion_physics == coll_krook_ions
-        calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, 
+        calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, vperp, collisions, evolve_density, 
                                             evolve_upar, evolve_ppar)
     else
         throw(ArgumentError("ion model $ion_physics not implemented for qpar calculation"))
@@ -820,7 +820,7 @@ end
 """
 calculate parallel heat flux if ion composition flag is coll_krook fluid ions
 """
-function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, collisions, evolve_density, evolve_upar, evolve_ppar)
+function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z, r, vperp, collisions, evolve_density, evolve_upar, evolve_ppar)
     # Note that this is a braginskii heat flux for ions using the krook operator. The full Fokker-Planck operator
     # Braginskii heat flux is different! This also assumes one ion species, and so no friction between ions.
     @boundscheck r.n == size(qpar, 2) || throw(BoundsError(qpar))
@@ -858,7 +858,20 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
     else
         return nothing
     end
-
+    # Stangeby (25.2) suggests that, when including kinetic effects, a value
+    # for gamma_i of around 2.5 is sensible.
+    # However, maybe for the purposes of this coll_krook scan, at very high
+    # collisionality we expect the distribution function of the ions at the
+    # sheath entrance to be close to a drifting maxwellian, in which case
+    # the original Stangeby (2.92) would be more appropriate. However, this
+    # also depends on whether we're 1V or 2V - as in 1V gamma_i = 2.5, 
+    # in 2V gamma_i = 3.5.
+
+    if vperp.n == 1
+        gamma_i = 2.5
+    else
+        gamma_i = 3.5
+    end
     @loop_r ir begin
         for iz ∈ z_indices
             this_ppar = vth[iz,ir]^2 * density[iz,ir]/2.0
@@ -867,9 +880,6 @@ function calculate_ion_qpar_from_coll_krook!(qpar, density, upar, vth, dT_dz, z,
             particle_flux = this_dens * this_upar
             T_i = vth[iz,ir]^2
 
-            # Stangeby (25.2)
-            gamma_i = 2.5
-
             # Stangeby (2.92)
             total_heat_flux = gamma_i * T_i * particle_flux