Allow loading variables from groups other than "dynamic_data"

moment_kinetics/src/load_data.jl

@@ -3657,7 +3657,7 @@ end
 """
     postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                            ir=nothing, iz=nothing, ivperp=nothing, ivpa=nothing,
-                           ivzeta=nothing, ivr=nothing, ivz=nothing)
+                           ivzeta=nothing, ivr=nothing, ivz=nothing, group=nothing)
 
 Load a variable
 
@@ -3671,10 +3671,14 @@ can be set to an integer or a range (e.g. `3:8` or `3:2:8`) to select subsets of
 Only the data for the subset requested will be loaded from the output file (mostly - when
 loading fields or moments from runs which used `parallel_io = false`, the full array will
 be loaded and then sliced).
+
+If a variable is found in a group other than "dynamic_data", the group name should be
+passed to the `group` argument.
 """
 function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                                 ir=nothing, iz=nothing, ivperp=nothing, ivpa=nothing,
-                                ivzeta=nothing, ivr=nothing, ivz=nothing)
+                                ivzeta=nothing, ivr=nothing, ivz=nothing, group=nothing)
+
     nt = run_info.nt
 
     if it === nothing
@@ -3786,8 +3790,12 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
     end
 
     if run_info.parallel_io
+        if group === nothing
+            group = "dynamic_data"
+        end
+
         # Get HDF5/NetCDF variables directly and load slices
-        variable = Tuple(get_group(f, "dynamic_data")[variable_name]
+        variable = Tuple(get_group(f, group)[variable_name]
                          for f ∈ run_info.files)
         nd = ndims(variable[1])
 
@@ -3984,15 +3992,19 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
         end
 
         if diagnostic_variable
+            if group === nothing
+                group = "dynamic_data"
+            end
             fid = open_readonly_output_file(run_info.files[1], run_info.ext, iblock=0)
-            group = get_group(fid, "dynamic_data")
-            result = load_variable(group, variable_name)
+            this_group = get_group(fid, group)
+            result = load_variable(this_group, variable_name)
             result = selectdim(result, ndims(result), it)
         elseif nd == 3
             result = allocate_float(run_info.z.n, run_info.r.n, run_info.nt)
             read_distributed_zr_data!(result, variable_name, run_info.files,
                                       run_info.ext, run_info.nblocks, run_info.z_local.n,
-                                      run_info.r_local.n, run_info.itime_skip)
+                                      run_info.r_local.n, run_info.itime_skip,
+                                      group=group)
             result = result[iz,ir,it]
         elseif nd == 4
             # If we ever have neutrals included but n_neutral_species != n_ion_species,
@@ -4003,16 +4015,25 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                                     run_info.nt)
             read_distributed_zr_data!(result, variable_name, run_info.files,
                                       run_info.ext, run_info.nblocks, run_info.z_local.n,
-                                      run_info.r_local.n, run_info.itime_skip)
+                                      run_info.r_local.n, run_info.itime_skip,
+                                      group=group)
             result = result[iz,ir,is,it]
         elseif nd === 5
+            if group !== nothing
+                error("group argument not supported when reading electron pdf with "
+                      * "parallel_io=false. Got group=$group")
+            end
             result = load_distributed_electron_pdf_slice(run_info.files, run_info.nblocks,
                                                          it, run_info.n_ion_species,
                                                          run_info.r_local,
                                                          run_info.z_local, run_info.vperp,
                                                          run_info.vpa; ir=ir, iz=iz,
                                                          ivperp=ivperp, ivpa=ivpa)
         elseif nd === 6
+            if group !== nothing
+                error("group argument not supported when reading ion pdf with "
+                      * "parallel_io=false. Got group=$group")
+            end
             result = load_distributed_ion_pdf_slice(run_info.files, run_info.nblocks, it,
                                                     run_info.n_ion_species,
                                                     run_info.r_local, run_info.z_local,
@@ -4021,6 +4042,10 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing,
                                                     ir=ir, iz=iz, ivperp=ivperp,
                                                     ivpa=ivpa)
         elseif nd === 7
+            if group !== nothing
+                error("group argument not supported when reading neutral pdf with "
+                      * "parallel_io=false. Got group=$group")
+            end
             result = load_distributed_neutral_pdf_slice(run_info.files, run_info.nblocks,
                                                         it, run_info.n_ion_species,
                                                         run_info.r_local,
@@ -5062,18 +5087,22 @@ restarts (which are sequential in time), so concatenate the data from each entry
 """
 function read_distributed_zr_data!(var::Array{mk_float,N}, var_name::String,
    run_names::Tuple, file_key::String, nblocks::Tuple,
-   nz_local::mk_int,nr_local::mk_int,iskip::mk_int) where N
+   nz_local::mk_int, nr_local::mk_int, iskip::mk_int; group=nothing) where N
     # dimension of var is [z,r,species,t]
 
+    if group === nothing
+        group = "dynamic_data"
+    end
+
     local_tind_start = 1
     local_tind_end = -1
     global_tind_start = 1
     global_tind_end = -1
     for (run_name, nb) in zip(run_names, nblocks)
         for iblock in 0:nb-1
             fid = open_readonly_output_file(run_name,file_key,iblock=iblock,printout=false)
-            group = get_group(fid, "dynamic_data")
-            var_local = load_variable(group, var_name)
+            this_group = get_group(fid, group)
+            var_local = load_variable(this_group, var_name)
 
             ntime_local = size(var_local, N)