Schedulers no longer break when particle failures occur

src/LearningRateSchedulers.jl

@@ -137,15 +137,17 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Calculates next timestep by pushing to ekp.Δt,
+Calculates next timestep by pushing to ekp.Δt, 
 `!isnothing(return_value)` implies termination condition has been met
 """
 function calculate_timestep!(
     ekp::EnsembleKalmanProcess,
     g::M,
     Δt_new::NFT,
 ) where {M <: AbstractMatrix, NFT <: Union{Nothing, AbstractFloat}}
-    terminate = calculate_timestep!(ekp, g, Δt_new, get_scheduler(ekp))
+    # when using g to calculate Δt, pass only successful particles through
+    successful_ens, _ = split_indices_by_success(g)
+    terminate = calculate_timestep!(ekp, g[:, successful_ens], Δt_new, get_scheduler(ekp))
     return terminate
 end
 
@@ -247,7 +249,6 @@ function calculate_timestep!(
         @info "Cannot override DataMisfitController-type timestep selection, ignoring Δt_new = $(Δt_new)"
     end
 
-
     M, J = size(g)
     T = scheduler.terminate_at
 
@@ -285,7 +286,6 @@ function calculate_timestep!(
         end
     end
 
-    M, J = size(g)
     y_mean = ekp.obs_mean
 
     Φ = [0.5 * norm(inv_sqrt_Γ * (g[:, j] - reshape(y_mean, :, 1)))^2 for j in 1:J]

test/EnsembleKalmanProcess/runtests.jl

@@ -284,6 +284,11 @@ end
 
         # Get inverse problem
         y_obs, G, Γy, A = inv_problem
+        if i_prob == 1
+            scheduler = DataMisfitController(on_terminate = "continue")
+        else
+            scheduler = DefaultScheduler()
+        end
 
         ekiobj = EKP.EnsembleKalmanProcess(
             initial_ensemble,
@@ -293,6 +298,7 @@ end
             rng = rng,
             failure_handler_method = SampleSuccGauss(),
             localization_method = loc_method,
+            scheduler = scheduler,
         )
         ekiobj_unsafe = EKP.EnsembleKalmanProcess(
             initial_ensemble,
@@ -302,6 +308,7 @@ end
             rng = rng,
             failure_handler_method = IgnoreFailures(),
             localization_method = loc_method,
+            scheduler = scheduler,
         )
 
         g_ens = G(get_ϕ_final(prior, ekiobj))