Unexpected behavior getting ϕ final mean distribution

[Julians42]

ArgumentError when return_array = true set for any of the functions similar to EKP.get_ϕ. For this argument configuration, the goal is to get a dictionary with the parameter names and their corresponding parameter entries (which may be an array). This is more useful when parameters have multiple entries per named entry.

One possible solution (for one of the functions) is as follows:

function EKP.get_ϕ_mean_final(prior, eki, return_array = false)
    final_means = EKP.get_ϕ_mean_final(prior, eki_final) # returns a vector of final iteration parameter means
    batch_idx = EKP.batch(prior) # groups the parameters by name (useful for priors whith multiple parameters per named entry)
    param_names = prior.name # get names
    # create new dictionary
    Dict(zip(param_names, [final_means[idx] for idx in batch_idx]))
end

[Julians42]

Another nice feature to add for models where we'd rather get the nearest neighbor to the ensemble mean (as opposed to the ensemble mean itself) is the following. Costa has used this method when the ensemble doesn't completely collapse and the parameter dimension is sufficiently high.

function get_ϕ_mean_final_nn(eki, prior, toml_final = "calibrate/toml/final_1m.toml")
    """
    Get the ensemble member of the final iteration closest to the ensemble mean 
    as measured in unconstrained space wrt 2-norm
    """
    u_final = EKP.get_u_final(eki)
    u_mean_final = EKP.get_u_mean_final(eki)
    best_member = argmin(sum((u_final .- u_mean_final) .^2, dims = 1))[2]

    # pick the best one in physical space
    ϕ_mean_final = EKP.get_ϕ_final(prior, eki)[:, best_member]

    # index and save to toml
    batch_idx = EKP.batch(prior)
    param_names = prior.name

    toml_dict = Dict(zip(param_names, [Dict("value" => length(ϕ_mean_final[idx]) == 1 ? ϕ_mean_final[idx][1] : ϕ_mean_final[idx]) for idx in batch_idx]))

    toml_string = TOML.print(toml_dict)
    open(toml_final, "w") do io 
        TOML.print(io, toml_dict)
    end
end

I'm writing here and saving to a toml file but to be consistent with what's used in EKP.jl we'd probably just want to return the list and the dictionary in two separate function methods.