Skip to content

Commit c275449

Browse files
committed
Add run_initial_ensemble
1 parent b3ed289 commit c275449

File tree

1 file changed

+84
-1
lines changed

1 file changed

+84
-1
lines changed

src/ekp_interface.jl

Lines changed: 84 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,11 @@ import EnsembleKalmanProcesses.TOMLInterface as TI
88

99
export get_prior, initialize, update_ensemble, save_G_ensemble
1010
export path_to_ensemble_member,
11-
path_to_model_log, path_to_iteration, parameter_path, load_latest_ekp
11+
path_to_model_log,
12+
path_to_iteration,
13+
parameter_path,
14+
load_latest_ekp,
15+
run_initial_ensemble
1216

1317
"""
1418
load_ekp_struct(output_dir, iteration)
@@ -308,6 +312,85 @@ function save_eki_and_parameters(eki, output_dir, iteration, prior)
308312
JLD2.save_object(ekp_path(output_dir, iteration), eki)
309313
end
310314

315+
"""
316+
run_initial_ensemble(
317+
backend::AbstractBackend,
318+
process::EKP.Process,
319+
prior::PD.ParameterDistribution,
320+
output_dir;
321+
ensemble_size = 0,
322+
rng = MersenneTwister(1234),
323+
run_forward_models_kwargs = (),
324+
)
325+
326+
Run the initial ensemble without needing the `EKP.EnsembleKalmanProcess`
327+
object and return the unconstrained parameters of the ensemble.
328+
329+
This is useful if the observations depend on the grid resolution of the
330+
simulation data as the first ensemble can be ran, the observations can be
331+
generated, and the `EKP.EnsembleKalmanProcess` object can be created afterward.
332+
333+
!!! warning "Initial ensemble"
334+
TODO: Add a warning about when to pass the unconstrained parameters to
335+
the EKP object
336+
"""
337+
function run_initial_ensemble( # TODO: Maybe move this to backends?
338+
backend,
339+
process::EKP.Process,
340+
prior::PD.ParameterDistribution,
341+
output_dir;
342+
ensemble_size = 0,
343+
rng = Random.MersenneTwister(1234),
344+
run_forward_models_kwargs = (),
345+
)
346+
# Cannot use last_completed_iteration because the JLD2 files storing the
347+
# G ensemble object and the ekp objects will never exists when calling this
348+
# function is used as intended
349+
iter0_path = joinpath(output_dir, "iteration_000")
350+
if isdir(iter0_path)
351+
@info "Detected first iteration. Not running the initial ensemble. If this is not the case, then delete the directory $iter0_path."
352+
end
353+
354+
initial_ensemble = _construct_intial_unconstrained_ensemble(
355+
process,
356+
prior,
357+
ensemble_size = ensemble_size,
358+
rng = rng,
359+
)
360+
361+
# The size of initial_ensemble is (number of parameters, ensemble members)
362+
# For some processes, the ensemble_size is ignored
363+
ensemble_size = last(size(initial_ensemble))
364+
365+
param_dict = get_param_dict(prior)
366+
TI.save_parameter_ensemble(
367+
initial_ensemble,
368+
prior,
369+
param_dict,
370+
output_dir,
371+
DEFAULT_PARAMETER_FILE,
372+
0,
373+
)
374+
375+
# Cannot save eki object because the goal of this function is to be able
376+
# to run the first ensemble without the eki object
377+
JLD2.save_object(
378+
joinpath(path_to_iteration(output_dir, 0), "prior.jld2"),
379+
prior,
380+
)
381+
382+
# TODO: Make the text files? Maybe this is done automatically by this
383+
# functions?
384+
run_forward_models(
385+
backend,
386+
0,
387+
ensemble_size;
388+
output_dir = output_dir,
389+
run_forward_models_kwargs...,
390+
)
391+
return initial_ensemble
392+
end
393+
311394
"""
312395
_construct_intial_unconstrained_ensemble(
313396
process::EKP.Process,

0 commit comments

Comments
 (0)