other examples working without error with EKP v2

CliMA · Sep 25, 2024 · 81131f2 · 81131f2
1 parent 0f9e5c6
commit 81131f2
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/examples/Learn_hyperparameters/1d_to_1d_regression_direct_withcov.jl b/examples/Learn_hyperparameters/1d_to_1d_regression_direct_withcov.jl
@@ -218,7 +218,7 @@ for (idx, type) in enumerate(feature_types)
     # Create EKI
     N_ens = 50
     N_iter = 20
-    initial_params = construct_initial_ensemble(priors, N_ens; rng_seed = ekp_seed)
+    initial_params = construct_initial_ensemble(rng, priors, N_ens)
     data = vcat(y[(n_train + 1):end], 0.0)
     ekiobj = [EKP.EnsembleKalmanProcess(initial_params, data, Γ, Inversion())]
 
@@ -247,7 +247,10 @@ for (idx, type) in enumerate(feature_types)
         Γ_tmp[(n_test + 1):end, (n_test + 1):end] += I
 
         #        ekiobj[1] = EKP.EnsembleKalmanProcess(initial_params, data, Γ, Inversion())
-        EKP.update_ensemble!(ekiobj[1], g_ens)
+        terminated = EKP.update_ensemble!(ekiobj[1], g_ens)
+        if !isnothing(terminated)
+            break
+        end
         err[i] = get_error(ekiobj[1])[end] #mean((params_true - mean(params_i,dims=2)).^2)
         println(
             "Iteration: " *

diff --git a/examples/Learn_hyperparameters/nd_to_1d_regression_direct_withcov.jl b/examples/Learn_hyperparameters/nd_to_1d_regression_direct_withcov.jl
@@ -323,7 +323,11 @@ for i in 1:N_iter
 
     end
 
-    EKP.update_ensemble!(ekiobj[1], g_ens)
+    terminated = EKP.update_ensemble!(ekiobj[1], g_ens)
+    if !isnothing(terminated)
+        break
+    end
+
     err[i] = get_error(ekiobj[1])[end] #mean((params_true - mean(params_i,dims=2)).^2)
     constrained_u = transform_unconstrained_to_constrained(priors, get_u_final(ekiobj[1]))
     println(

diff --git a/examples/Learn_hyperparameters/nd_to_md_regression_direct_withcov.jl b/examples/Learn_hyperparameters/nd_to_md_regression_direct_withcov.jl
@@ -510,7 +510,11 @@ end
         end
 
 
-        EKP.update_ensemble!(ekiobj[1], g_ens, Δt_new = Δt[1])
+        terminated = EKP.update_ensemble!(ekiobj[1], g_ens, Δt_new = Δt[1])
+        if !isnothing(terminated)
+            break
+        end
+
         err[i] = get_error(ekiobj[1])[end] #mean((params_true - mean(params_i,dims=2)).^2)
         constrained_u = transform_unconstrained_to_constrained(priors, get_u_final(ekiobj[1]))
         println(