API

forward_model(iteration, member)

Execute the forward model simulation with the given configuration.

This function must be overridden by a component's model interface and should set things like the parameter path and other member-specific settings.

observation_map(iteration)

Runs the observation map for the specified iteration. This function must be implemented for each calibration experiment.

WorkerBackend

Used to run calibrations on Distributed.jl's workers. For use on a Slurm cluster, see SlurmManager.

SlurmManager(ntasks=get(ENV, "SLURM_NTASKS", 1))

The ClusterManager for Slurm clusters, taking in the number of tasks to request with srun.

To execute the srun command, run addprocs(SlurmManager(ntasks))

Keyword arguments can be passed to srun: addprocs(SlurmManager(ntasks), gpus_per_task=1)

By default the workers will inherit the running Julia environment.

To run a calibration, call calibrate(WorkerBackend, ...)

To run functions on a worker, call remotecall(func, worker_id, args...)

PBSManager(ntasks)

The ClusterManager for PBS/Torque clusters, taking in the number of tasks to request with qsub.

To execute the qsub command, run addprocs(PBSManager(ntasks)). Unlike the SlurmManager, this will not nest scheduled jobs, but will acquire new resources.

Keyword arguments can be passed to qsub: addprocs(PBSManager(ntasks), nodes=2)

By default, the workers will inherit the running Julia environment.

To run a calibration, call calibrate(WorkerBackend, ...)

To run functions on a worker, call remotecall(func, worker_id, args...)

set_worker_loggers(workers = workers())

Set the global logger to a simple file logger for the given workers.

map_remotecall_fetch(f::Function, args...; workers = workers())

Call function f from each worker and wait for the results to return.

foreach_remotecall_wait(f::Function, args...; workers = workers())

Call function f from each worker.

calibrate(backend, ensemble_size, n_iterations, observations, noise, prior, output_dir; ekp_kwargs...)
calibrate(backend, ekp::EnsembleKalmanProcess, ensemble_size, n_iterations, prior, output_dir)
calibrate(backend, config::ExperimentConfig; ekp_kwargs...)

Run a full calibration on the given backend.

If the EKP struct is not given, it will be constructed upon initialization. The experiment configuration (ensemble size, prior, observations, etc) can be wrapped in an ExperimentConfig or passed in as arguments to the function.

Available Backends: WorkerBackend, CaltechHPCBackend, ClimaGPUBackend, DerechoBackend, JuliaBackend

Derecho, ClimaGPU, and CaltechHPC backends are designed to run on a specific high-performance computing cluster. WorkerBackend uses Distributed.jl to run the forward model on workers.

Keyword Arguments for HPC backends

• `model_interface: Path to the model interface file.
• hpc_kwargs: Dictionary of resource arguments for HPC clusters, passed to the job scheduler.
• verbose::Bool: Enable verbose logging.
• Any keyword arguments for the EnsembleKalmanProcess constructor, such as scheduler

JuliaBackend

The simplest backend, used to run a calibration in Julia without any parallelization.

DerechoBackend

Used for NSF NCAR's Derecho supercomputing system.

CaltechHPCBackend

Used for Caltech's high-performance computing cluster.

ClimaGPUBackend

Used for CliMA's private GPU server.

get_backend()

Get ideal backend for deploying forward model runs. Each backend is found via gethostname(). Defaults to JuliaBackend if none is found.

model_run(backend, iter, member, output_dir, experiment_dir; model_interface, verbose, hpc_kwargs)

Construct and execute a command to run a single forward model on a given job scheduler.

Dispatches on backend to run slurm_model_run or pbs_model_run.

Arguments:

• iter: Iteration number
• member: Member number
• output_dir: Calibration experiment output directory
• project_dir: Directory containing the experiment's Project.toml
• model_interface: Model interface file
• moduleloadstr: Commands which load the necessary modules
• hpc_kwargs: Dictionary containing the resources for the job. Easily generated using kwargs.

module_load_string(backend)

Return a string that loads the correct modules for a given backend when executed via bash.

wait_for_jobs(jobids, output_dir, iter, experiment_dir, model_interface, module_load_str, model_run_func; verbose, hpc_kwargs, reruns=1)

Wait for a set of jobs to complete. If a job fails, it will be rerun up to reruns times.

This function monitors the status of multiple jobs and handles failures by rerunning the failed jobs up to the specified number of reruns. It logs errors and job completion status, ensuring all jobs are completed before proceeding.

Arguments:

• jobids: Vector of job IDs.
• output_dir: Directory for output.
• iter: Iteration number.
• experiment_dir: Directory for the experiment.
• model_interface: Interface to the model.
• module_load_str: Commands to load necessary modules.
• model_run_func: Function to run the model.
• verbose: Print detailed logs if true.
• hpc_kwargs: HPC job parameters.
• reruns: Number of times to rerun failed jobs.

log_member_error(output_dir, iteration, member, verbose=false)

Log a warning message when an error occurs. If verbose, includes the ensemble member's output.

kill_job(jobid::SlurmJobID)
kill_job(jobid::PBSJobID)

End a running job, catching errors in case the job can not be ended.

job_status(job_id)

Parse the slurm job_id's state and return one of three status symbols: :PENDING, :RUNNING, or :COMPLETED.

kwargs(; kwargs...)

Create a dictionary from keyword arguments.

slurm_model_run(iter, member, output_dir, experiment_dir, model_interface, module_load_str; hpc_kwargs)

Construct and execute a command to run a single forward model on Slurm. Helper function for model_run.

generate_sbatch_script(iter, member, output_dir, experiment_dir, model_interface; module_load_str, hpc_kwargs)

Generate a string containing an sbatch script to run the forward model.

hpc_kwargs is turned into a series of sbatch directives using generate_sbatch_directives. module_load_str is used to load the necessary modules and can be obtained via module_load_string.

generate_sbatch_directives(hpc_kwargs)

Generate Slurm sbatch directives from HPC kwargs.

submit_slurm_job(sbatch_filepath; env=deepcopy(ENV))

Submit a job to the Slurm scheduler using sbatch, removing unwanted environment variables.

Unset variables: "SLURMMEMPERCPU", "SLURMMEMPERGPU", "SLURMMEMPER_NODE"

pbs_model_run(iter, member, output_dir, experiment_dir, model_interface, module_load_str; hpc_kwargs)

Construct and execute a command to run a single forward model on PBS Pro. Helper function for model_run.

generatepbsscript( iter, member, outputdir, experimentdir, modelinterface; moduleloadstr, hpckwargs, )

Generate a string containing a PBS script to run the forward model.

Returns:

• qsub_contents::Function: A function generating the content of the PBS script based on the provided arguments. This will run the contents of the julia_script, which have to be run from a file due to Derecho's set_gpu_rank.
• julia_script::String: The Julia script string to be executed by the PBS job.

Helper function for pbs_model_run.

submit_pbs_job(sbatch_filepath; env=deepcopy(ENV))

Submit a job to the PBS Pro scheduler using qsub, removing unwanted environment variables.

Unset variables: "PBSMEMPERCPU", "PBSMEMPERGPU", "PBSMEMPER_NODE"

initialize(ensemble_size, observations, noise, prior, output_dir)
initialize(eki::EKP.EnsembleKalmanProcess, prior, output_dir)
initialize(config)

Initialize a calibration, saving the initial parameter ensemble to a folder within output_dir.

If no EKP struct is given, construct an EKP struct and return it.

save_G_ensemble(config::ExperimentConfig, iteration, G_ensemble)
save_G_ensemble(output_dir::AbstractString, iteration, G_ensemble)

Saves the ensemble's observation map output to the correct directory based on the provided configuration. Takes an output directory, either extracted from an ExperimentConfig or passed directly.

update_ensemble(output_dir::AbstractString, iteration, prior)
update_ensemble(config::ExperimentConfig, iteration)
update_ensemble(config_file::AbstractString, iteration)

Updates the EnsembleKalmanProcess object and saves the parameters for the next iteration.

update_ensemble!(ekp, G_ens, output_dir, iteration, prior)

Updates an EKP object with data G_ens, saving the object and final parameters to disk.

observation_map_and_update!(ekp, output_dir, iteration, prior)

Compute the observation map and update the given EKP object.

ExperimentConfig(
    n_iterations::Integer,
    ensemble_size::Integer,
    observations,
    noise,
    prior::ParameterDistribution,
    output_dir,
)
ExperimentConfig(filepath::AbstractString; kwargs...)

Construct an ExperimentConfig from a given YAML file or directory containing 'experiment_config.yml'.

ExperimentConfig holds the configuration for a calibration experiment. This can be constructed from a YAML configuration file or directly using individual parameters.

get_prior(param_dict::AbstractDict; names = nothing)
get_prior(prior_path::AbstractString; names = nothing)

Constructs the combined prior distribution from a param_dict or a TOML configuration file specified by prior_path. If names is provided, only those parameters are used.

get_param_dict(distribution; names)

Generates a dictionary for parameters based on the specified distribution, assumed to be of floating-point type. If names is not provided, the distribution's names will be used.

path_to_iteration(output_dir, iteration)

Return the path to the directory for a given iteration within the specified output directory.

path_to_ensemble_member(output_dir, iteration, member)

Return the path to an ensemble member's directory for a given iteration and member number.

path_to_model_log(output_dir, iteration, member)

Return the path to an ensemble member's forward model log for a given iteration and member number.

parameter_path(output_dir, iteration, member)

Return the path to an ensemble member's parameter file.