Merge #101

101: Handle parameter names corresponding to vectors r=costachris a=costachris

In preparation for function learning we'll need to seamlessly deal with parameter names that correspond to vectors. This is handled here by treating each vector component as a unique parameter in EKP.

Co-authored-by: costachris <christopouloscosta@gmail.com>

src/DistributionUtils.jl

@@ -39,16 +39,50 @@ function construct_priors(
     outdir_path::String = pwd(),
     to_file::Bool = true,
 )
-    u_names = collect(keys(params))
-    constraints = collect(values(params))
-    n_param = length(u_names)
 
+    # if parameter vectors found => flatten
+    if any(1 .< [length(val) for val in collect(values(params))])
+        u_names, constraints = flatten_param_dict(params)
+    else
+        u_names = collect(keys(params))
+        constraints = collect(values(params))
+    end
+
+    n_param = length(u_names)
     # All vars are approximately uniform in unconstrained space
     distributions = repeat([Parameterized(Normal(0.0, unconstrained_σ))], n_param)
     to_file ? jldsave(joinpath(outdir_path, "prior.jld2"); distributions, constraints, u_names) : nothing
     return ParameterDistribution(distributions, constraints, u_names)
 end
 
+"""
+    flatten_param_dict(param_dict::Dict{String, Vector{Constraint}})
+
+For parameter names that correspond to vectors, assign unique name to each vector component.
+Inputs:
+    param_dict :: Dictionary of parameter names to constraints.
+Outputs:
+    u_names :: vector of parameter names
+    constraints :: vector of constraints
+"""
+function flatten_param_dict(param_dict::Dict{String, Vector{Constraint}})
+
+    u_names = Vector{String}()
+    constraints = Vector{Vector{Constraint}}()
+    for (param, value) in param_dict
+        if length(value) > 1
+            for j in 1:length(value)
+                push!(u_names, "$(param)_{$j}")
+                push!(constraints, [value[j]])
+            end
+        else
+            push!(u_names, param)
+            push!(constraints, [value])
+        end
+    end
+    return (u_names, constraints)
+end
+
 """
     deserialize_prior(prior_dict::Dict{String, Any})
 

src/Pipeline.jl

@@ -96,7 +96,7 @@ function init_calibration(config::Dict{Any, Any}; mode::String = "hpc", job_id::
     )
 
     # Dimensionality
-    n_param = length(collect(keys(params)))
+    n_param = sum(map(length, collect(values(params))))
     d = length(ref_stats.y)
     if algo_name == "Unscented"
         N_ens = 2 * n_param + 1

src/TurbulenceConvectionUtils.jl

@@ -175,6 +175,8 @@ function run_SCM_handler(
     inputdir = scm_dir(m)
     namelist = JSON.parsefile(namelist_directory(inputdir, m))
 
+    u_names, u = create_parameter_vectors(u_names, u)
+
     # update parameter values
     for (pName, pVal) in zip(u_names, u)
         namelist["turbulence"]["EDMF_PrognosticTKE"][pName] = pVal
@@ -197,6 +199,47 @@ function run_SCM_handler(
     return data_directory(tmpdir, m.case_name, uuid), model_error
 end
 
+"""
+    create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}
+
+Given vector of parameter names and corresponding values, combine any vector components
+into single parameter vectors for input into SCM.
+
+Inputs:
+    u_names :: SCM names for parameters `u`, which may contain vector components.
+    u :: Values of parameters to be used in simulations, which may contain vector components.
+Outputs:
+    u_names_out :: SCM names for parameters `u`.
+    u_out :: Values of parameters to be used in simulations.
+"""
+function create_parameter_vectors(u_names::Vector{String}, u::Vector{FT}) where {FT <: AbstractFloat}
+
+    u_names_out = []
+    u_out = []
+    u_vec_dict = Dict()
+
+    for i in 1:length(u_names)
+        param_name_i = u_names[i]
+        if occursin(r"{?}", param_name_i)
+            param_name_split = split(param_name_i, "_")
+            param_vec_name = join(param_name_split[1:(length(param_name_split) - 1)], "_")
+            if param_vec_name in keys(u_vec_dict)
+                push!(u_vec_dict[param_vec_name], u[i])
+            else
+                u_vec_dict[param_vec_name] = [u[i]]
+            end
+        else
+            push!(u_names_out, u_names[i])
+            push!(u_out, u[i])
+
+        end
+    end
+    append!(u_names_out, collect(keys(u_vec_dict)))
+    append!(u_out, collect(values(u_vec_dict)))
+
+    return (u_names_out, u_out)
+end
+
 """
     run_SCM_handler(
         m::ReferenceModel,