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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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",