Skip to content

Commit

Permalink
Explicitly defined struct for I/O input
Browse files Browse the repository at this point in the history 
Also define a `setup_io_input()` function to simplify `mk_input()`.
  • Loading branch information
@johnomotani
johnomotani committed Sep 16, 2024
1 parent 813c319 commit 0405d08
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 65 deletions.
82 changes: 73 additions & 9 deletions moment_kinetics/src/file_io.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ module file_io
export input_option_error
export get_group
export open_output_file, open_ascii_output_file
export setup_io_input
export setup_file_io, finish_file_io
export write_data_to_ascii

Expand All @@ -19,6 +20,7 @@ using ..type_definitions: mk_float, mk_int
using LibGit2
using MPI
using Pkg
using UUIDs

@debug_shared_array using ..communication: DebugMPISharedArray

Expand All @@ -32,6 +34,23 @@ function __init__()
end
end

"""
Container for I/O settings
"""
Base.@kwdef struct io_input_struct
run_name::String
base_directory::String
ascii_output::Bool
binary_format::binary_format_type
parallel_io::Bool
run_id::String
output_dir::String
write_error_diagnostics::Bool
write_steady_state_diagnostics::Bool
write_electron_error_diagnostics::Bool
write_electron_steady_state_diagnostics::Bool
end

"""
structure containing the various input/output streams
"""
Expand All @@ -56,8 +75,7 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
Tchodura_upper, Texti1, Texti2, Texti3, Texti4,
Texti5, Textn1, Textn2, Textn3, Textn4, Textn5, Texte1, Texte2,
Texte3, Texte4, Tconstri, Tconstrn, Tconstre, Tint, Tfailcause,
Telectrontime, Telectronint, Telectronfailcause, Tnldiagnostics,
Tinput}
Telectrontime, Telectronint, Telectronfailcause, Tnldiagnostics}
# file identifier for the binary file to which data is written
fid::Tfile
# handle for the time variable
Expand Down Expand Up @@ -211,14 +229,14 @@ struct io_moments_info{Tfile, Ttime, Tphi, Tmomi, Tmome, Tmomn, Tchodura_lower,
nl_solver_diagnostics::Tnldiagnostics

# Settings for I/O
io_input::Tinput
io_input::io_input_struct
end

"""
structure containing the data/metadata needed for binary file i/o
distribution function data only
"""
struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments, Tinput}
struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments}
# file identifier for the binary file to which data is written
fid::Tfile
# handle for the ion species distribution function variable
Expand All @@ -241,7 +259,7 @@ struct io_dfns_info{Tfile, Tfi, Tfe, Tfn, Tmoments, Tinput}
f_neutral_start_last_timestep::Union{Tfn,Nothing}

# Settings for I/O
io_input::Tinput
io_input::io_input_struct

# Handles for moment variables
io_moments::Tmoments
Expand All @@ -252,8 +270,7 @@ structure containing the data/metadata needed for binary file i/o
for electron initialization
"""
struct io_initial_electron_info{Tfile, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4,
Tconstr, Telectrontime, Telectronint, Telectronfailcause,
Tinput}
Tconstr, Telectrontime, Telectronint, Telectronfailcause}
# file identifier for the binary file to which data is written
fid::Tfile
time::Telectrontime
Expand Down Expand Up @@ -313,7 +330,54 @@ struct io_initial_electron_info{Tfile, Tfe, Tmom, Texte1, Texte2, Texte3, Texte4
electron_dt_before_last_fail::Telectrontime

# Settings for I/O
io_input::Tinput
io_input::io_input_struct
end

"""
Read the settings for I/O
"""
function setup_io_input(input_dict, timestepping_section; ignore_MPI=false)
io_settings = set_defaults_and_check_section!(
input_dict, "output";
run_name="",
base_directory="runs",
ascii_output=false,
binary_format=hdf5,
parallel_io=nothing,
)
if io_settings["run_name"] == ""
error("When passing a Dict directly for input, it is required to set `run_name` "
* "in the `[output]` section")
end
if io_settings["parallel_io"] === nothing
io_settings["parallel_io"] = io_has_parallel(Val(io_settings["binary_format"]))
end
# Make copy of the section to avoid modifying the passed-in Dict
io_settings = copy(io_settings)
run_id = string(uuid4())
if !ignore_MPI
# Communicate run_id to all blocks
# Need to convert run_id to a Vector{Char} for MPI
run_id_chars = [run_id...]
MPI.Bcast!(run_id_chars, 0, comm_world)
run_id = string(run_id_chars...)
end
io_settings["run_id"] = run_id
io_settings["output_dir"] = joinpath(io_settings["base_directory"], io_settings["run_name"])
io_settings["write_error_diagnostics"] = timestepping_section["write_error_diagnostics"]
io_settings["write_steady_state_diagnostics"] = timestepping_section["write_steady_state_diagnostics"]
io_settings["write_electron_error_diagnostics"] = timestepping_section["electron_t_input"]["write_error_diagnostics"]
io_settings["write_electron_steady_state_diagnostics"] = timestepping_section["electron_t_input"]["write_steady_state_diagnostics"]

# Create output_dir if it does not exist.
if !ignore_MPI
if global_rank[] == 0
mkpath(io_settings["output_dir"])
end
_block_synchronize()
end

return io_input_struct(; (Symbol(k) => v for (k,v) ∈ io_settings)...)
end

"""
Expand Down Expand Up @@ -1931,7 +1995,7 @@ function define_dynamic_dfn_variables!(fid, r, z, vperp, vpa, vzeta, vr, vz, com
io_f_electron, io_f_electron_loworder,
io_f_electron_start_last_timestep, io_f_neutral,
io_f_neutral_loworder, io_f_neutral_start_last_timestep,
parallel_io, io_moments)
io_input, io_moments)
end

# For processes other than the root process of each shared-memory group...
Expand Down
69 changes: 13 additions & 56 deletions moment_kinetics/src/moment_kinetics_input.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ using ..array_allocation: allocate_float
using ..communication
using ..coordinates: define_coordinate, get_coordinate_input
using ..external_sources
using ..file_io: io_has_parallel, input_option_error, open_ascii_output_file
using ..file_io: io_has_parallel, input_option_error, open_ascii_output_file,
setup_io_input
using ..krook_collisions: setup_krook_collisions_input
using ..maxwell_diffusion: setup_mxwl_diff_collisions_input
using ..fokker_planck: setup_fkpl_collisions_input
Expand All @@ -25,7 +26,6 @@ using ..geo: init_magnetic_geometry, setup_geometry_input
using ..species_input: get_species_input
using MPI
using TOML
using UUIDs

"""
Read input from a TOML file
Expand Down Expand Up @@ -94,24 +94,6 @@ function mk_input(input_dict=OptionsDict(); save_inputs_to_txt=false, ignore_MPI
n_ion_species = composition.n_ion_species
n_neutral_species = composition.n_neutral_species

# Start processing inputs for file I/O. This is done early because we need to work
# out what `output_dir` should be. The setup is completed later, after some other
# sections have been read.
io_settings = set_defaults_and_check_section!(
input_dict, "output";
run_name="",
base_directory="runs",
ascii_output=false,
binary_format=hdf5,
parallel_io=nothing,
)
if io_settings["run_name"] == ""
error("When passing a Dict directly for input, it is required to set `run_name` "
* "in the `[output]` section")
end
# this is the directory where the simulation data will be stored
output_dir = joinpath(io_settings["base_directory"], io_settings["run_name"])

# if evolve_moments.density = true, evolve density via continuity eqn
# and g = f/n via modified drift kinetic equation
evolve_moments_settings = set_defaults_and_check_section!(
Expand Down Expand Up @@ -277,7 +259,6 @@ function mk_input(input_dict=OptionsDict(); save_inputs_to_txt=false, ignore_MPI
# `electron_timestepping_section` into the Dict that is used to make
# `timestepping_input`, so that it becomes part of `timestepping_input`.
timestepping_section = copy(timestepping_section)
timestepping_section["stopfile_name"] = joinpath(output_dir, "stop")
timestepping_section["electron_t_input"] = electron_timestepping_section
if !(0.0 < timestepping_section["step_update_prefactor"] < 1.0)
error("step_update_prefactor=$(timestepping_section["step_update_prefactor"]) must "
Expand Down Expand Up @@ -312,6 +293,11 @@ function mk_input(input_dict=OptionsDict(); save_inputs_to_txt=false, ignore_MPI
error("maximum_dt=$(timestepping_section["maximum_dt"]) must be positive")
end

io_immutable = setup_io_input(input_dict, timestepping_section; ignore_MPI=ignore_MPI)

# this is the directory where the simulation data will be stored
timestepping_section["stopfile_name"] = joinpath(io_immutable.output_dir, "stop")

manufactured_solns_input = setup_manufactured_solutions(input_dict)

#########################################################################
Expand All @@ -335,45 +321,15 @@ function mk_input(input_dict=OptionsDict(); save_inputs_to_txt=false, ignore_MPI
r_coord_input.nelement_local)
end

# Create output_dir if it does not exist.
if !ignore_MPI
if global_rank[] == 0
mkpath(output_dir)
end
_block_synchronize()
end

em_input = set_defaults_and_check_section!(
input_dict, em_fields_input, "em_fields"
)

# Complete setup of io_settings
if io_settings["parallel_io"] === nothing
io_settings["parallel_io"] = io_has_parallel(Val(io_settings["binary_format"]))
end
# Make copy of the section to avoid modifying the passed-in Dict
io_settings = copy(io_settings)
run_id = string(uuid4())
if !ignore_MPI
# Communicate run_id to all blocks
# Need to convert run_id to a Vector{Char} for MPI
run_id_chars = [run_id...]
MPI.Bcast!(run_id_chars, 0, comm_world)
run_id = string(run_id_chars...)
end
io_settings["run_id"] = run_id
io_settings["output_dir"] = output_dir
io_settings["write_error_diagnostics"] = timestepping_section["write_error_diagnostics"]
io_settings["write_steady_state_diagnostics"] = timestepping_section["write_steady_state_diagnostics"]
io_settings["write_electron_error_diagnostics"] = timestepping_section["electron_t_input"]["write_error_diagnostics"]
io_settings["write_electron_steady_state_diagnostics"] = timestepping_section["electron_t_input"]["write_steady_state_diagnostics"]
io_immutable = Dict_to_NamedTuple(io_settings)

# initialize z grid and write grid point locations to file
if ignore_MPI
run_directory = nothing
else
run_directory = output_dir
run_directory = io_immutable.output_dir
end
z, z_spectral = define_coordinate(z_coord_input; parallel_io=io_immutable.parallel_io,
run_directory=run_directory, ignore_MPI=ignore_MPI,
Expand Down Expand Up @@ -438,15 +394,16 @@ function mk_input(input_dict=OptionsDict(); save_inputs_to_txt=false, ignore_MPI

if global_rank[] == 0 && save_inputs_to_txt
# Make file to log some information about inputs into.
io = open_ascii_output_file(string(output_dir,"/",io_settings["run_name"]), "input")
io = open_ascii_output_file(joinpath(io_immutable.output_dir, io_immutable.run_name), "input")
else
io = devnull
end

# check input (and initialized coordinate structs) to catch errors/unsupported options
check_input(io, output_dir, timestepping_section["nstep"], timestepping_section["dt"], r, z,
vpa, vperp, composition, species_immutable, evolve_moments,
num_diss_params, save_inputs_to_txt, collisions)
check_input(io, io_immutable.output_dir, timestepping_section["nstep"],
timestepping_section["dt"], r, z, vpa, vperp, composition,
species_immutable, evolve_moments, num_diss_params, save_inputs_to_txt,
collisions)

# return immutable structs for z, vpa, species and composition
all_inputs = (io_immutable, evolve_moments, timestepping_section, z, z_spectral, r,
Expand Down

0 comments on commit 0405d08

Please sign in to comment.