Add helper function to load EKP file (#190)

* Add func to load latest EKP file

* Set JuliaFormatter version to match other repos

* Add compat for docs

Author: nefrathenrici

.github/workflows/julia_formatter.yml

@@ -12,5 +12,5 @@ jobs:
     steps:
       - uses: julia-actions/julia-format@v3
         with:
-          version: '2' # Set `version` to '1.0.54' if you need to use JuliaFormatter.jl v1.0.54 (default: '1')
+          version: '1' # Set `version` to '1.0.54' if you need to use JuliaFormatter.jl v1.0.54 (default: '1')
           suggestion-label: 'format-suggest' # leave this unset or empty to show suggestions for all PRs

docs/Project.toml

@@ -14,3 +14,4 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 
 [compat]
 ClimaAtmos = "=0.28.3"
+EnsembleKalmanProcesses = "2"

docs/src/api.md

@@ -65,4 +65,5 @@ ClimaCalibrate.path_to_model_log
 ClimaCalibrate.parameter_path
 ClimaCalibrate.minibatcher_over_samples
 ClimaCalibrate.observation_series_from_samples
+ClimaCalibrate.load_latest_ekp
 ```

src/ekp_interface.jl

@@ -8,7 +8,7 @@ import EnsembleKalmanProcesses.TOMLInterface as TI
 
 export get_prior, initialize, update_ensemble, save_G_ensemble
 export path_to_ensemble_member,
-    path_to_model_log, path_to_iteration, parameter_path
+    path_to_model_log, path_to_iteration, parameter_path, load_latest_ekp
 
 """
     load_ekp_struct(output_dir, iteration)
@@ -18,6 +18,22 @@ Return the EnsembleKalmanProcess struct for a completed iteration.
 load_ekp_struct(output_dir, iteration) =
     JLD2.load_object(ekp_path(output_dir, iteration))
 
+"""
+    load_latest_ekp(output_dir)
+
+Return the most recent EnsembleKalmanProcess struct from the given output directory.
+
+Returns nothing if no EKP structs are found.
+"""
+function load_latest_ekp(output_dir)
+    iter = -1
+    while isfile(ekp_path(output_dir, iter + 1))
+        iter += 1
+    end
+    iter == -1 && return nothing
+    return load_ekp_struct(output_dir, iter)
+end
+
 """
     path_to_ensemble_member(output_dir, iteration, member)
 

src/workers.jl

@@ -133,8 +133,7 @@ function Distributed.launch(
     output_path = get(params, :o, get(params, :output, default_output))
 
     ntasks = sm.ntasks
-    srun_cmd =
-        `srun -J $jobname -n $ntasks -D $exehome $worker_args -o $output_path -- $exename $exeflags $(worker_cookie_arg())`
+    srun_cmd = `srun -J $jobname -n $ntasks -D $exehome $worker_args -o $output_path -- $exename $exeflags $(worker_cookie_arg())`
     @info "Starting SLURM job $jobname: $srun_cmd"
     pid = open(addenv(srun_cmd, env))
 
@@ -336,8 +335,7 @@ function Distributed.launch(
         -j oe: Send the output and error streams to the same file
         -J 1-ntasks: Job array
         -o: output file =#
-    qsub_cmd =
-        `qsub -V -N $jobname -j oe $job_array_option $worker_args -o $output_path -- $exename $exeflags $(worker_cookie_arg())`
+    qsub_cmd = `qsub -V -N $jobname -j oe $job_array_option $worker_args -o $output_path -- $exename $exeflags $(worker_cookie_arg())`
     @info "Starting PBS job $jobname: $qsub_cmd"
     pid = open(addenv(qsub_cmd, env))
 

test/ekp_interface.jl

@@ -67,6 +67,26 @@ end
     @test eki.u[1].stored_data == user_constructed_eki.u[1].stored_data
 end
 
+@testset "Test loading latest EKP struct" begin
+    # Test loading from directory with no completed iterations
+    empty_dir = joinpath(output_dir, "empty")
+    mkpath(empty_dir)
+    @test isnothing(CAL.load_latest_ekp(empty_dir))
+
+    # Test loading from directory with completed iterations
+    # We already have an EKP struct saved from earlier tests
+    latest_ekp = CAL.load_latest_ekp(output_dir)
+    @test !isnothing(latest_ekp)
+    @test latest_ekp isa EKP.EnsembleKalmanProcess
+    # Compare with the known EKP struct
+    for prop in propertynames(latest_ekp)
+        prop in [:u, :accelerator, :localizer] && continue
+        @test getproperty(latest_ekp, prop) ==
+              getproperty(eki_with_kwargs, prop)
+    end
+    @test latest_ekp.u[1].stored_data == eki_with_kwargs.u[1].stored_data
+end
+
 override_file =
     joinpath(output_dir, "iteration_000", "member_001", "parameters.toml")
 td = CP.create_toml_dict(FT; override_file)