Add stochastic parameter update to EKI

src/EnsembleKalmanInversion.jl

@@ -198,3 +198,75 @@ function update_ensemble!(
         )
     end
 end
+
+"""
+    stochastic_update!(
+        ekp::EnsembleKalmanProcess{FT, IT, Inversion};
+        stochastic_scaling_factor,
+    ) where {FT, IT}
+Applies stochastic prediction step to ekp.
+
+Inputs:
+ - ekp :: The EnsembleKalmanProcess to update.
+ - stochastic_scaling_factor :: Scaling factor for the stochastic perturbation.
+"""
+function stochastic_update!(
+    ekp::EnsembleKalmanProcess{FT, IT, Inversion};
+    stochastic_scaling_factor::FT = 1.0,
+) where {FT, IT}
+    u = get_u_final(ekp)
+    u_mean = get_u_mean_final(ekp)
+    prefactor = sqrt(1 / (1 - stochastic_scaling_factor * ekp.Δt[end]))
+
+    stochastic_Δt = stochastic_scaling_factor * ekp.Δt[end]
+    if stochastic_Δt >= 1.0
+        @warn string(
+            "Parameter noise time step: ",
+            stochastic_Δt,
+            "is >= 1.0",
+            "\nChange stochastic_scaling_factor or EK time step.",
+        )
+    end
+
+    u_updated = u_mean .+ prefactor * (u .- u_mean)
+    ekp.u[end] = DataContainer(u_updated, data_are_columns = true)
+
+end
+
+"""
+    update_ensemble!(
+        ekp::EnsembleKalmanProcess{FT, IT, Inversion},
+        g::AbstractMatrix{FT},
+        stoch_flag::Bool;
+        cov_threshold::Real = 0.01,
+        Δt_new::Union{Nothing, FT} = nothing,
+        deterministic_forward_map::Bool = true,
+        failed_ens = nothing,
+    ) where {FT, IT}
+
+Updates the ensemble according to an Inversion process.
+
+Inputs:
+ - ekp :: The EnsembleKalmanProcess to update.
+ - g :: Model outputs, they need to be stored as a `N_obs × N_ens` array (i.e data are columms).
+ - stoch_flag :: Flag to indicate whether to use stochastic update or not.
+ - cov_threshold :: Threshold below which the reduction in covariance determinant results in a warning.
+ - Δt_new :: Time step to be used in the current update.
+ - deterministic_forward_map :: Whether output `g` comes from a deterministic model.
+ - failed_ens :: Indices of failed particles. If nothing, failures are computed as columns of `g`
+    with NaN entries.
+"""
+function update_ensemble!(
+    ekp::EnsembleKalmanProcess{FT, IT, Inversion},
+    g::AbstractMatrix{FT},
+    stoch_flag::Bool;
+    cov_threshold::Real = 0.01,
+    Δt_new::Union{Nothing, FT} = nothing,
+    deterministic_forward_map::Bool = true,
+    failed_ens = nothing,
+    stochastic_scaling_factor::FT = 1.0,
+) where {FT, IT}
+
+    update_ensemble!(ekp, g; cov_threshold, Δt_new, deterministic_forward_map, failed_ens)
+    stoch_flag ? stochastic_update!(ekp; stochastic_scaling_factor = stochastic_scaling_factor) : nothing
+end