add stochastic_update function

src/EnsembleKalmanProcess.jl

@@ -448,6 +448,38 @@ function get_cov_blocks(cov::AbstractMatrix{FT}, p::IT) where {FT <: Real, IT <:
     return uu_cov, ug_cov, gg_cov
 end
 
+"""
+    stochastic_update!(
+        ekp::EnsembleKalmanProcess;
+        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; stochastic_scaling_factor::FT = 1.0) where {FT <: Real}
+
+    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,
@@ -458,10 +490,19 @@ 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 indicating whether to use stochastic update.
+ - stochastic_scaling_factor :: Scaling factor for time step in stochastic perturbation.
  - ekp_kwargs :: Keyword arguments to pass to standard ekp update_ensemble!.
 """
-function update_ensemble!(ekp::EnsembleKalmanProcess, g::AbstractMatrix{FT}; ekp_kwargs...) where {FT, IT}
+function update_ensemble!(
+    ekp::EnsembleKalmanProcess,
+    g::AbstractMatrix{FT};
+    stoch_flag::Bool = false,
+    stochastic_scaling_factor::FT = 1.0,
+    ekp_kwargs...,
+) where {FT, IT}
     update_ensemble!(ekp, g, get_process(ekp); ekp_kwargs...)
+    stoch_flag ? stochastic_update!(ekp; stochastic_scaling_factor = stochastic_scaling_factor) : nothing
 end
 
 ## include the different types of Processes and their exports: