working well!

CliMA · Oct 31, 2023 · 563bb7b · 563bb7b
1 parent 77a2439
commit 563bb7b
Show file tree

Hide file tree

Showing 12 changed files with 22 additions and 354 deletions.
diff --git a/examples/Manifest.toml b/examples/Manifest.toml
diff --git a/examples/driver-job-array.sbatch b/examples/driver-job-array.sbatch
diff --git a/examples/driver.sbatch b/examples/driver.sbatch
diff --git a/examples/job_array.sh b/examples/job_array.sh
@@ -1,23 +1,23 @@
 #!/bin/bash
+
+# Set number of tasks per model run
 tasks_per_ensemble=8
 
-# Load required modules
+# Configure for cluster
 module load julia/1.9.3 cuda/12.2 ucx/1.14.1_cuda-12.2 openmpi/4.1.5_cuda-12.2 hdf5/1.12.2-ompi415
 export OMPI_MCA_opal_warn_on_missing_libcuda=0
+export JULIA_NUM_PRECOMPILE_TASKS=8
+export JULIA_CPU_TARGET='broadwell;skylake'
 
 # Check if a job ID is provided as a command-line argument
 job_id=${1?Error: no job ID given}
-# exec > output/$job_id/log.out 2>&1
 
 echo "Running job: $job_id"
 
 # Instantiate the Julia project
-export JULIA_CPU_TARGET='broadwell;skylake'
-
 julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'
 
 echo 'Initializing ensemble for calibration.'
-
 julia --project=examples -e "
 import CalibrateAtmos
 using Distributions
@@ -43,6 +43,3 @@ do
     dep="afterany:$update_id"
     echo "Model run $(($i+1)) will wait for $update_id"
 done
-
-# Restore output to stdout
-# exec >&-
diff --git a/examples/model_run.sbatch b/examples/model_run.sbatch
@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --time=4:00:00
+#SBATCH --time=2:00:00
 #SBATCH --cpus-per-task=8
 #SBATCH --mem-per-cpu=8G
 #SBATCH --mail-user=nat.henrici@gmail.com

diff --git a/examples/sphere_held_suarez_rhoe_equilmoist/ekp_config.yml b/examples/sphere_held_suarez_rhoe_equilmoist/ekp_config.yml
@@ -2,6 +2,6 @@ output_dir: output/sphere_held_suarez_rhoe_equilmoist
 prior_path: examples/sphere_held_suarez_rhoe_equilmoist/prior.toml
 parameter_names: ["equator_pole_temperature_gradient_wet"]
 ensemble_size: 10
-n_iterations: 5
+n_iterations: 3
 obs_mean: examples/sphere_held_suarez_rhoe_equilmoist/obs_mean.jld2
 obs_noise_cov: examples/sphere_held_suarez_rhoe_equilmoist/obs_noise_cov.jld2
diff --git a/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl b/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl
@@ -63,4 +63,5 @@ function observation_map(iteration, job_id)
     )
     iter_path = joinpath(member_path[1:(length(member_path) - 1)])
     JLD2.save_object(joinpath(iter_path, "observation_map.jld2"), G_ens)
+    return G_ens
 end
diff --git a/examples/sphere_held_suarez_rhoe_equilmoist/prior.toml b/examples/sphere_held_suarez_rhoe_equilmoist/prior.toml
@@ -1,6 +1,6 @@
 ["equator_pole_temperature_gradient_wet"]
-prior = "Parameterized(Normal(4.1291,0.3007608))"
-# approximates a positive distribution with mean 65, std of 20
+prior = "Parameterized(Normal(4.779568,0.31223328))"
 constraint = "[bounded_below(0)]"
+# EKP.ParameterDistributions.constrained_gaussian("name", 125, 40, 0, Inf)
 type = "float"
 alias = "ΔT_y_wet"
diff --git a/examples/update.sbatch b/examples/update.sbatch
@@ -10,6 +10,8 @@ echo "Iteration $i"
 
 julia --project=examples -e '
     import CalibrateAtmos
+    using Distributions
+    using EnsembleKalmanProcesses.ParameterDistributions
     include("examples/'$job_id'/observation_map.jl")
     observation_map('$i', "'$job_id'")
     CalibrateAtmos.update_ensemble("'${job_id}'", '$i')

diff --git a/plot/Project.toml b/plot/Project.toml
@@ -2,3 +2,8 @@
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
diff --git a/src/atmos_interface.jl b/src/atmos_interface.jl
@@ -16,7 +16,7 @@ function get_atmos_config(member, iteration, config_dict::AbstractDict)
     config_dict["output_dir"] = member_path
     parameter_path = joinpath(member_path, "parameters.toml")
     config_dict["toml"] = [parameter_path]
-    
+
     # Turn off default diagnostics
     config_dict["output_default_diagnostics"] = false
 

diff --git a/src/ekp_interface.jl b/src/ekp_interface.jl
@@ -43,13 +43,13 @@ function initialize(
 
     output_dir = config["output_dir"]
     prior_path = config["prior_path"]
-    names = config["parameter_names"]
+    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 names]
+    prior_vec = [get_parameter_distribution(param_dict, n) for n in param_names]
     prior = combine_distributions(prior_vec)
 
     initial_ensemble =
@@ -115,7 +115,6 @@ function update_ensemble(
         output_dir,
         "parameters.toml",
         iteration,
-        apply_constraints = false,
     )
 
     # Save EKI object for next iteration