Improve ClimaArtifacts Usage

CliMA · Jun 12, 2024 · bf8cdc3 · bf8cdc3
1 parent 43d6005
commit bf8cdc3
Show file tree

Hide file tree

Showing 9 changed files with 29 additions and 41 deletions.
diff --git a/.gitignore b/.gitignore
@@ -46,8 +46,5 @@ test/Manifest.toml
 *.lock
 *.loc
 
-# Calibration experiment restart files
-!calibration/experiments/*/*.hdf5
-
 # misc
 .DS_Store
diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Manifest.toml b/calibration/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.3"
+julia_version = "1.10.4"
 manifest_format = "2.0"
-project_hash = "662aeda187ef55e20754ee64b3bea95f814d7055"
+project_hash = "7277cdc4b48ddb0fe371d652de3b39eeaf24bdde"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "9b3993053b87311e699ac9e1bcabec5ae7444f65"
@@ -343,9 +343,9 @@ version = "0.24.0"
 
 [[deps.ClimaCalibrate]]
 deps = ["Distributions", "EnsembleKalmanProcesses", "JLD2", "Random", "TOML", "YAML"]
-git-tree-sha1 = "acbbe76ab3a4c018d7edd55585fc767fe29f4324"
+git-tree-sha1 = "7f304ff3ca24d417020dcdec5654f1fc91558b86"
 repo-rev = "main"
-repo-url = "https://github.com/CliMA/ClimaCalibrate.jl"
+repo-url = "https://github.com/CliMA/ClimaCalibrate.jl.git"
 uuid = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
 version = "0.1.0"
 
@@ -1654,12 +1654,6 @@ git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4"
 uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 version = "1.0.2"
 
-[[deps.NetCDF]]
-deps = ["DiskArrays", "NetCDF_jll"]
-git-tree-sha1 = "c8ec10c6b72e35f7020b83330f8054908ba454eb"
-uuid = "30363a11-5582-574a-97bb-aa9a979735b9"
-version = "0.11.8"
-
 [[deps.NetCDF_jll]]
 deps = ["Artifacts", "Blosc_jll", "Bzip2_jll", "HDF5_jll", "JLLWrappers", "LazyArtifacts", "LibCURL_jll", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "TOML", "XML2_jll", "Zlib_jll", "Zstd_jll", "libzip_jll"]
 git-tree-sha1 = "ef30054d4e6eab74228cc7beae0c0873129bc078"

diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Project.toml b/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/Project.toml
@@ -7,8 +7,6 @@ ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
-TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
 
 [compat]

diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/experiment_config.yml b/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/experiment_config.yml
diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/observation_map.jl b/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/observation_map.jl
@@ -2,29 +2,26 @@ import EnsembleKalmanProcesses: TOMLInterface
 import ClimaCalibrate: observation_map, ExperimentConfig
 using ClimaAnalysis
 
-const config = ExperimentConfig(@__DIR__)
 function observation_map(iteration)
-    (; ensemble_size, output_dir) = config
-    dims = 1
-    G_ensemble = Array{Float64}(undef, dims..., ensemble_size)
+    ensemble_size = 10
+    output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
+    single_member_dims = 1
+    G_ensemble = Array{Float64}(undef, single_member_dims..., ensemble_size)
     for m in 1:ensemble_size
         member_path =
             TOMLInterface.path_to_ensemble_member(output_dir, iteration, m)
         simdir = SimDir(joinpath(member_path, "output_active"))
-        try
-            G_ensemble[:, m] .= process_member_data(simdir)
-        catch err
-            @info "Error during observation map for ensemble member $m" err
-            G_ensemble[:, m] .= NaN
-        end
+        G_ensemble[:, m] .= process_member_data(simdir)
     end
     return G_ensemble
 end
 
 const meters = 1.0
 const days = 86400.0
+# Cut off first 120 day to get equilibrium, take second level slice
 function process_member_data(simdir)
-    # Cut off first 120 days to get equilibrium, take second level slice
+    isempty(simdir.vars) && return  NaN
+
     ta = get(simdir; short_name = "ta", reduction = "average", period = "60d")
     zonal_avg_temp_observations =
         slice(average_lat(average_lon(ta)), z = 242meters)

diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/pipeline.sbatch b/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/pipeline.sbatch
@@ -4,13 +4,24 @@
 
 julia --project=calibration/experiments/sphere_held_suarez_rhoe_equilmoist -e '
 import Pkg; Pkg.instantiate(;verbose=true)
+import JLD2
+using ClimaUtilities.ClimaArtifacts
 import ClimaCalibrate: calibrate, ExperimentConfig, CaltechHPC, get_prior, kwargs
 
 experiment_dir = dirname(Base.active_project())
 include(joinpath(experiment_dir, "observation_map.jl"))
+artifact_path = @clima_artifact("atmos_held_suarez_obs")
 
+n_iterations = 2
+ensemble_size = 10
+observations = JLD2.load_object(joinpath(artifact_path, "obs_mean.jld2"))
+noise = JLD2.load_object(joinpath(artifact_path, "obs_noise_cov.jld2"))
+prior = get_prior(joinpath(experiment_dir, "prior.toml"))
+output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
+experiment_config = ExperimentConfig(; n_iterations, ensemble_size, 
+                                    observations, noise, output_dir, prior)
 slurm_kwargs = kwargs(time = 50, ntasks = 16)
-eki = calibrate(CaltechHPC, experiment_dir; slurm_kwargs, verbose=true)
+eki = calibrate(CaltechHPC, experiment_config; slurm_kwargs, verbose=true)
 
 include(joinpath(experiment_dir, "postprocessing.jl"))
 prior = get_prior(joinpath(experiment_dir, "prior.toml"))

diff --git a/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/postprocessing.jl b/calibration/experiments/sphere_held_suarez_rhoe_equilmoist/postprocessing.jl
@@ -2,10 +2,7 @@ import EnsembleKalmanProcesses as EKP
 using EnsembleKalmanProcesses.ParameterDistributions
 using EnsembleKalmanProcesses.TOMLInterface
 using Distributions
-import JLD2
 import Statistics: mean
-import YAML
-import TOML
 import CairoMakie: Makie
 import ClimaCalibrate
 
@@ -20,13 +17,12 @@ function convergence_plot(
     error_vec = Float64[]
     spread_vec = Float64[]
     for ensemble in u_vec
-        ensemble_error = 0
-        ensemble_spread = 0
+        ensemble_error = 0.0
+        ensemble_spread = 0.0
         ensemble_mean = mean(ensemble)
         for i in ensemble
             ensemble_error += abs(i - theta_star)^2
             ensemble_spread += abs(i - ensemble_mean)^2
-
         end
         ensemble_error /= length(ensemble)
         ensemble_spread /= length(ensemble)
@@ -101,6 +97,7 @@ function scatter_plot(
 end
 
 # Uncomment for easy plotting
+# import JLD2
 # iteration = 2
 # output_dir = joinpath("output", "sphere_held_suarez_rhoe_equilmoist")
 # eki_filepath = joinpath(ClimaCalibrate.path_to_iteration(output_dir, iteration), "eki_file.jld2")

diff --git a/calibration/model_interface.jl b/calibration/model_interface.jl
@@ -1,7 +1,6 @@
 import ClimaAtmos as CA
 import YAML
 using ClimaUtilities.ClimaArtifacts
-
 import ClimaCalibrate:
     set_up_forward_model, run_forward_model, path_to_ensemble_member
 

diff --git a/test/calibration_interface.jl b/test/calibration_interface.jl
@@ -1,6 +1,7 @@
 # Tests for ensuring that the calibration/model interface sets AtmosConfig correctly.
 # To run: julia --project=calibration/experiments/sphere_held_suarez_rhoe_equilmoist test/calibration_interface.jl
 import ClimaCalibrate
+using ClimaUtilities.ClimaArtifacts
 using Test
 
 experiment_dir = dirname(Base.active_project())
@@ -24,7 +25,7 @@ atmos_config = ClimaCalibrate.set_up_forward_model(1, 1, experiment_dir)
     @test parsed_args["toml"] == [file_path]
     @test parsed_args["output_dir"] == member_path
     @test parsed_args["restart_file"] ==
-          "/groups/esm/ClimaArtifacts/artifacts/atmos_held_suarez_obs/day200.0.hdf5"
+          joinpath(@clima_artifact("atmos_held_suarez_obs"), "day200.0.hdf5")
 end
 
 rm(file_path)