diff --git a/moment_kinetics/src/load_data.jl b/moment_kinetics/src/load_data.jl index 7781e38df0..8279942c7d 100644 --- a/moment_kinetics/src/load_data.jl +++ b/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]) @@ -3985,14 +3993,15 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, if diagnostic_variable 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,9 +4012,14 @@ 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, @@ -4013,6 +4027,10 @@ function postproc_load_variable(run_info, variable_name; it=nothing, is=nothing, 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 +4039,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,9 +5084,13 @@ 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 @@ -5072,8 +5098,8 @@ function read_distributed_zr_data!(var::Array{mk_float,N}, var_name::String, 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)