Merge pull request #55 from CliMA/ne/config

Update Atmos, simplify EKP config, other small fixes

docs/src/experiment_setup_guide.md

@@ -215,15 +215,13 @@ The EKP configuration file must contain the following:
 - `n_iterations`, the number of iterations to run
 - `ensemble_size`, the ensemble size
 - `prior_path`, the path to the TOML file with the prior parameter distributions
-- `parameter_names`, a list of the long names of the parameters being calibrated
 - `truth_data`, the truth data
 - `truth_noise`, the covariance of the truth data
 - `output_dir`, the folder where you want calibration data and logs to be output. This must be the same as the `output_dir` in the model configuration file.
 Example:
 ```
 output_dir: output/sphere_held_suarez_rhoe_equilmoist
 prior_path: experiments/sphere_held_suarez_rhoe_equilmoist/prior.toml
-parameter_names: ["equator_pole_temperature_gradient_wet"]
 ensemble_size: 10
 n_iterations: 3
 truth_data: experiments/sphere_held_suarez_rhoe_equilmoist/obs_mean.jld2

experiments/sphere_held_suarez_rhoe_equilmoist/Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.0"
+julia_version = "1.10.1"
 manifest_format = "2.0"
-project_hash = "77d5fb58df087bb3ab845a50e744683f56660cad"
+project_hash = "830cda58ec7ee94d462f4e2c264d62d1988abf42"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
@@ -115,9 +115,9 @@ version = "7.7.1"
 
 [[deps.ArrayLayouts]]
 deps = ["FillArrays", "LinearAlgebra"]
-git-tree-sha1 = "64d582bcb9c93ac741234789eeb4f16812413efb"
+git-tree-sha1 = "e46675dbc095ddfdf2b5fba247d5a25f34e1f8a2"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
-version = "1.6.0"
+version = "1.6.1"
 weakdeps = ["SparseArrays"]
 
     [deps.ArrayLayouts.extensions]
@@ -237,9 +237,9 @@ version = "0.5.0"
 
 [[deps.CFTime]]
 deps = ["Dates", "Printf"]
-git-tree-sha1 = "ed2e76c1c3c43fd9d0cb9248674620b29d71f2d1"
+git-tree-sha1 = "5afb5c5ba2688ca43a9ad2e5a91cbb93921ccfa1"
 uuid = "179af706-886a-5703-950a-314cd64e0468"
-version = "0.1.2"
+version = "0.1.3"
 
 [[deps.CLIMAParameters]]
 deps = ["DocStringExtensions", "TOML", "Test"]
@@ -296,9 +296,9 @@ version = "0.5.0"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra"]
-git-tree-sha1 = "aef70bb349b20aa81a82a19704c3ef339d4ee494"
+git-tree-sha1 = "575cd02e080939a33b6df6c5853d14924c08e35b"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.22.1"
+version = "1.23.0"
 weakdeps = ["SparseArrays"]
 
     [deps.ChainRulesCore.extensions]
@@ -312,9 +312,9 @@ version = "0.20.1"
 
 [[deps.ClimaComms]]
 deps = ["CUDA", "MPI"]
-git-tree-sha1 = "57c054ddd4280ca8e2b5915ef1cf1395c4edbc78"
+git-tree-sha1 = "f0350e34c91c8f3b5a11b5e39990439303d727b1"
 uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
-version = "0.5.6"
+version = "0.5.7"
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "Unrolled"]
@@ -334,9 +334,9 @@ version = "0.10.2"
 
 [[deps.ClimaTimeSteppers]]
 deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "DiffEqCallbacks", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"]
-git-tree-sha1 = "216385dda9966d820b484f399b90447f073893a6"
+git-tree-sha1 = "9c203f39784c968700c55f555754a7771b3410df"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
-version = "0.7.16"
+version = "0.7.19"
 
 [[deps.CloseOpenIntervals]]
 deps = ["Static", "StaticArrayInterface"]
@@ -404,7 +404,7 @@ weakdeps = ["Dates", "LinearAlgebra"]
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
-version = "1.0.5+1"
+version = "1.1.0+0"
 
 [[deps.CompositionsBase]]
 git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad"
@@ -474,9 +474,9 @@ version = "1.6.1"
 
 [[deps.DataStructures]]
 deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
-git-tree-sha1 = "1fb174f0d48fe7d142e1109a10636bc1d14f5ac2"
+git-tree-sha1 = "0f4b5d62a88d8f59003e43c25a8a90de9eb76317"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.18.17"
+version = "0.18.18"
 
 [[deps.DataValueInterfaces]]
 git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
@@ -512,9 +512,9 @@ version = "0.1.0+0"
 
 [[deps.DiffEqBase]]
 deps = ["ArrayInterface", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces"]
-git-tree-sha1 = "2ad3a2dcd5f28f535aa884d199cc2f0a9d335729"
+git-tree-sha1 = "aee5798bd4f1ed1260bd1741221f9589e5ee8a90"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
-version = "6.147.1"
+version = "6.147.3"
 
     [deps.DiffEqBase.extensions]
     DiffEqBaseChainRulesCoreExt = "ChainRulesCore"
@@ -639,9 +639,9 @@ uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 version = "1.0.4"
 
 [[deps.EnzymeCore]]
-git-tree-sha1 = "59c44d8fbc651c0395d8a6eda64b05ce316f58b4"
+git-tree-sha1 = "496c5455d6a61c2a6f2233ce07c1fcdbe4995ab6"
 uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
-version = "0.6.5"
+version = "0.7.0"
 weakdeps = ["Adapt"]
 
     [deps.EnzymeCore.extensions]
@@ -807,10 +807,10 @@ uuid = "88fa7841-ef32-4516-bb70-c6ec135699d9"
 version = "0.1.0"
 
 [[deps.GnuTLS_jll]]
-deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Nettle_jll", "P11Kit_jll", "Pkg", "Zlib_jll"]
-git-tree-sha1 = "266fe9b2335527cbf569ba4fd0979e3d8c6fd491"
+deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Nettle_jll", "P11Kit_jll", "Zlib_jll"]
+git-tree-sha1 = "f3c0936dd685d57fa0b1eee7dbebf382b969ea63"
 uuid = "0951126a-58fd-58f1-b5b3-b08c7c4a876d"
-version = "3.7.8+1"
+version = "3.8.3+0"
 
 [[deps.HDF5]]
 deps = ["Compat", "HDF5_jll", "Libdl", "MPIPreferences", "Mmap", "Preferences", "Printf", "Random", "Requires", "UUIDs"]
@@ -927,9 +927,9 @@ version = "1.0.0"
 
 [[deps.JLD2]]
 deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"]
-git-tree-sha1 = "7c0008f0b7622c6c0ee5c65cbc667b69f8a65672"
+git-tree-sha1 = "5ea6acdd53a51d897672edb694e3cc2912f3f8a7"
 uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-version = "0.4.45"
+version = "0.4.46"
 
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
@@ -951,9 +951,9 @@ version = "0.2.1+0"
 
 [[deps.KernelAbstractions]]
 deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
-git-tree-sha1 = "c7753cc3febe006708ce6798482004241f7d890b"
+git-tree-sha1 = "ed7167240f40e62d97c1f5f7735dea6de3cc5c49"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
-version = "0.9.17"
+version = "0.9.18"
 weakdeps = ["EnzymeCore"]
 
     [deps.KernelAbstractions.extensions]
@@ -1208,9 +1208,9 @@ uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [[deps.MathOptInterface]]
 deps = ["BenchmarkTools", "CodecBzip2", "CodecZlib", "DataStructures", "ForwardDiff", "JSON", "LinearAlgebra", "MutableArithmetics", "NaNMath", "OrderedCollections", "PrecompileTools", "Printf", "SparseArrays", "SpecialFunctions", "Test", "Unicode"]
-git-tree-sha1 = "e8b98c868029d007102dc5f98986c81f33b0ec37"
+git-tree-sha1 = "679c1aec6934d322783bd15db4d18f898653be4f"
 uuid = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
-version = "1.26.0"
+version = "1.27.0"
 
 [[deps.MbedTLS_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -1295,10 +1295,10 @@ uuid = "c020b1a1-e9b0-503a-9c33-f039bfc54a85"
 version = "1.0.0"
 
 [[deps.NetCDF]]
-deps = ["DiskArrays", "Formatting", "NetCDF_jll"]
-git-tree-sha1 = "328178762645783b20495d408ab317b4c2d25b1a"
+deps = ["DiskArrays", "NetCDF_jll"]
+git-tree-sha1 = "c8ec10c6b72e35f7020b83330f8054908ba454eb"
 uuid = "30363a11-5582-574a-97bb-aa9a979735b9"
-version = "0.11.7"
+version = "0.11.8"
 
 [[deps.NetCDF_jll]]
 deps = ["Artifacts", "Blosc_jll", "Bzip2_jll", "HDF5_jll", "JLLWrappers", "LibCURL_jll", "Libdl", "OpenMPI_jll", "XML2_jll", "Zlib_jll", "Zstd_jll", "libzip_jll"]
@@ -1334,7 +1334,7 @@ version = "0.3.24+0"
 [[deps.OpenBLAS_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
 uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
-version = "0.3.23+2"
+version = "0.3.23+4"
 
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -1449,9 +1449,9 @@ version = "1.2.0"
 
 [[deps.Preferences]]
 deps = ["TOML"]
-git-tree-sha1 = "00805cd429dcb4870060ff49ef443486c262e38e"
+git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6"
 uuid = "21216c6a-2e73-6563-6e65-726566657250"
-version = "1.4.1"
+version = "1.4.3"
 
 [[deps.PrettyTables]]
 deps = ["Crayons", "LaTeXStrings", "Markdown", "PrecompileTools", "Printf", "Reexport", "StringManipulation", "Tables"]
@@ -1481,9 +1481,9 @@ version = "0.1.4"
 
 [[deps.ProgressMeter]]
 deps = ["Distributed", "Printf"]
-git-tree-sha1 = "00099623ffee15972c16111bcf84c58a0051257c"
+git-tree-sha1 = "763a8ceb07833dd51bb9e3bbca372de32c0605ad"
 uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
-version = "1.9.0"
+version = "1.10.0"
 
 [[deps.PyCall]]
 deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
@@ -1931,9 +1931,9 @@ uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 version = "0.5.23"
 
 [[deps.TranscodingStreams]]
-git-tree-sha1 = "54194d92959d8ebaa8e26227dbe3cdefcdcd594f"
+git-tree-sha1 = "3caa21522e7efac1ba21834a03734c57b4611c7e"
 uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
-version = "0.10.3"
+version = "0.10.4"
 weakdeps = ["Random", "Test"]
 
     [deps.TranscodingStreams.extensions]

experiments/sphere_held_suarez_rhoe_equilmoist/ekp_config.yml

@@ -1,6 +1,5 @@
 output_dir: output/sphere_held_suarez_rhoe_equilmoist
 prior_path: experiments/sphere_held_suarez_rhoe_equilmoist/prior.toml
-parameter_names: ["equator_pole_temperature_gradient_wet"]
 ensemble_size: 10
 n_iterations: 3
 truth_data: experiments/sphere_held_suarez_rhoe_equilmoist/obs_mean.jld2

experiments/sphere_held_suarez_rhoe_equilmoist/generate_observations.sbatch

@@ -2,6 +2,7 @@
 #SBATCH --time=24:00:00
 #SBATCH --ntasks=32
 #SBATCH --cpus-per-task=8
+#SBATCH --partition=expansion
 #SBATCH --output="experiments/sphere_held_suarez_rhoe_equilmoist/truth_simulation/model_log.out"
 #SBATCH --partition=expansion
 

experiments/sphere_held_suarez_rhoe_equilmoist/observation_map.jl

@@ -1,4 +1,4 @@
-using NetCDF
+import NetCDF as NC
 using Statistics
 import YAML
 import EnsembleKalmanProcesses: TOMLInterface
@@ -40,8 +40,13 @@ function observation_map(::Val{:sphere_held_suarez_rhoe_equilmoist}, iteration)
     for m in 1:ensemble_size
         member_path =
             TOMLInterface.path_to_ensemble_member(output_dir, iteration, m)
-        ta = ncread(joinpath(member_path, model_output), "ta")
-        G_ensemble[:, m] = process_member_data(ta)
+        try
+            ta = NC.ncread(joinpath(member_path, model_output), "ta")
+            G_ensemble[:, m] = process_member_data(ta)
+        catch e
+            @assert e isa NC.NetCDFError
+            G_ensemble[:, m] .= NaN
+        end
     end
     return G_ensemble
 end

experiments/surface_fluxes_perfect_model/ekp_config.yml

@@ -1,6 +1,5 @@
 output_dir: output/surface_fluxes_perfect_model
 prior_path: experiments/surface_fluxes_perfect_model/prior.toml
-parameter_names: ["coefficient_a_m_businger", "coefficient_a_h_businger"]
 ensemble_size: 10
 n_iterations: 10
 truth_data: experiments/surface_fluxes_perfect_model/data/obs_mean.jld2

plot/convergence_surface_fluxes_perfect_model.jl

@@ -117,10 +117,7 @@ eki_path = joinpath(
 eki = JLD2.load_object(eki_path);
 EKP.get_u(eki)
 prior_path = joinpath(pkg_dir, ekp_config["prior_path"])
-param_names = ekp_config["parameter_names"]
-param_dict = TOML.parsefile(prior_path)
-prior_vec = [get_parameter_distribution(param_dict, n) for n in param_names]
-prior = combine_distributions(prior_vec)
+prior = CalibrateAtmos.get_prior(prior_path)
 
 theta_star_vec =
     (; coefficient_a_m_businger = 4.7, coefficient_a_h_businger = 4.7)
@@ -129,7 +126,7 @@ convergence_plot(
     eki,
     prior,
     theta_star_vec,
-    param_names,
+    ["coefficient_a_m_businger", "coefficient_a_h_businger"],
     joinpath(pkg_dir, "output"),
 )
 

src/ekp_interface.jl

@@ -14,13 +14,18 @@ path_to_iteration(output_dir, iteration) =
     joinpath(output_dir, join(["iteration", lpad(iteration, 3, "0")], "_"))
 
 """
-    get_prior(prior_path; names = nothing)
+    get_prior(prior_pathAbstractString; names = nothing)
+    get_prior(param_dict::AbstractDict; names = nothing)
 
 Constructs the combined prior distribution from the TOML file at the `prior_path`.
 If no parameter names are passed in, all parameters in the TOML are used in the distribution.
 """
-function get_prior(prior_path; names = nothing)
+function get_prior(prior_path::AbstractString; names = nothing)
     param_dict = TOML.parsefile(prior_path)
+    return get_prior(param_dict; names)
+end
+
+function get_prior(param_dict::AbstractDict; names = nothing)
     names = isnothing(names) ? keys(param_dict) : names
     prior_vec = [get_parameter_distribution(param_dict, n) for n in names]
     prior = combine_distributions(prior_vec)
@@ -52,15 +57,12 @@ function initialize(
     rng_ekp = Random.MersenneTwister(rng_seed)
 
     output_dir = config["output_dir"]
-    prior_path = config["prior_path"]
-    param_names = config["parameter_names"]
     ensemble_size = config["ensemble_size"]
     # Save in EKI object in iteration_000 folder
     eki_path = joinpath(output_dir, "iteration_000", "eki_file.jld2")
 
-    param_dict = TOML.parsefile(prior_path)
-    prior_vec = [get_parameter_distribution(param_dict, n) for n in param_names]
-    prior = combine_distributions(prior_vec)
+    param_dict = TOML.parsefile(config["prior_path"])
+    prior = get_prior(param_dict)
 
     initial_ensemble =
         EKP.construct_initial_ensemble(rng_ekp, prior, ensemble_size)
@@ -100,7 +102,6 @@ function update_ensemble(
     config = YAML.load_file("experiments/$experiment_id/ekp_config.yml"),
 )
     output_dir = config["output_dir"]
-    names = config["parameter_names"]
     # Load EKI object from iteration folder
     iter_path = path_to_iteration(output_dir, iteration)
     eki_path = joinpath(iter_path, "eki_file.jld2")
@@ -114,10 +115,9 @@ function update_ensemble(
     iteration += 1
 
     # Update and save parameters for next iteration
-    prior_path = config["prior_path"]
-    param_dict = TOML.parsefile(prior_path)
-    prior_vec = [get_parameter_distribution(param_dict, n) for n in names]
-    prior = combine_distributions(prior_vec)
+    param_dict = TOML.parsefile(config["prior_path"])
+    prior = get_prior(param_dict)
+
     save_parameter_ensemble(
         EKP.get_u_final(eki),  # constraints applied when saving
         prior,