Basic pipeline working

Author: nefrathenrici

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -427,12 +427,9 @@ function get_atmos_config_dict(coupler_dict::Dict, job_id::String, atmos_output_
     atmos_toml = joinpath.(pkgdir(CA), atmos_config["toml"])
     coupler_toml = joinpath.(pkgdir(ClimaCoupler), coupler_dict["coupler_toml"])
     toml = isempty(coupler_toml) ? atmos_toml : coupler_toml
-    if haskey(atmos_config, "calibration_toml")
-        push!(toml, atmos_config["calibration_toml"])
-    end
     if !isempty(toml)
         @info "Overwriting Atmos parameters from input TOML file(s): $toml"
-        atmos_config = merge(atmos_config, Dict("toml" => toml))
+        atmos_config["toml"] = toml
     end
 
     # Specify atmos output directory to be inside the coupler output directory

experiments/calibration/cld_eff_rad/generate_observations.jl

@@ -5,7 +5,8 @@ import EnsembleKalmanProcesses as EKP
 
 const start_date = DateTime(2000, 3, 1)
 const slice_time = 0.0
-const diagnostic_var = OutputVar(joinpath(pkgdir(ClimaCoupler), "experiments/calibration/cld_eff_rad/rsut_1M_average_8_helem.nc"))
+const diagnostic_var =
+    OutputVar(joinpath(pkgdir(ClimaCoupler), "experiments/calibration/cld_eff_rad/rsut_1M_average_8_helem.nc"))
 
 include(joinpath(pkgdir(ClimaCoupler), "experiments/ClimaEarth/leaderboard/data_sources.jl"))
 resample(ov) = resampled_as(ov, diagnostic_var, dim_names = ["longitude", "latitude"])
@@ -27,17 +28,9 @@ cre = rsutcs - rsut
 
 # Create an EKP.Observation
 rsut_slice = slice(rsut; time = slice_time)
-rsut_obs = EKP.Observation(
-    vec(rsut_slice.data),
-    Diagonal(vec(auto_covariance(rsut))),
-    "rsut"
-)
+rsut_obs = EKP.Observation(vec(rsut_slice.data), Diagonal(vec(auto_covariance(rsut))), "rsut")
 cre_slice = slice(cre; time = slice_time)
-cre_obs = EKP.Observation(
-    vec(cre_slice.data),
-    Diagonal(vec(auto_covariance(cre))),
-    "cre"
-)
+cre_obs = EKP.Observation(vec(cre_slice.data), Diagonal(vec(auto_covariance(cre))), "cre")
 
 observations = EKP.combine_observations([rsut_obs, cre_obs])
 JLD2.save_object("experiments/calibration/cld_eff_rad/observations.jld2", observations)

experiments/calibration/cld_eff_rad/model_config.yml

@@ -1,5 +1,5 @@
 FLOAT_TYPE: "Float32"
-# use_itime: true
+use_itime: false
 albedo_model: "CouplerAlbedo"
 atmos_config_file: "config/longrun_configs/amip_target_edonly.yml"
 checkpoint_dt: "99999days"

experiments/calibration/cld_eff_rad/model_interface.jl

@@ -7,15 +7,15 @@ ENV["CLIMACOMMS_CONTEXT"] = "SINGLETON"
 include(joinpath(pkgdir(ClimaCoupler), "experiments", "ClimaEarth", "setup_run.jl"))
 
 function ClimaCalibrate.forward_model(iter, member)
-
+    redirect_stderr(stdout) # This should live in ClimaCalibrate.set_worker_logger
     config_file = joinpath(pkgdir(ClimaCoupler), "experiments", "calibration", "cld_eff_rad", "model_config.yml")
     config_dict = get_coupler_config_dict(config_file)
 
     output_dir_root = config_dict["coupler_output_dir"]
 
     # Set member parameter file
     sampled_parameter_file = ClimaCalibrate.parameter_path(output_dir_root, iter, member)
-    config_dict["calibration_toml"] = sampled_parameter_file
+    config_dict["coupler_toml"] = [sampled_parameter_file]
     # Set member output directory
     member_output_dir = ClimaCalibrate.path_to_ensemble_member(output_dir_root, iter, member)
     config_dict["coupler_output_dir"] = member_output_dir

experiments/calibration/cld_eff_rad/observation_map.jl

@@ -6,7 +6,7 @@ import EnsembleKalmanProcesses as EKP
 
 function ClimaCalibrate.observation_map(iteration)
     observation_vec = JLD2.load_object(observation_path)
-    single_member_dims = length(EKP.get_obs(first(observation_vec))) * batch_size
+    single_member_dims = length(EKP.get_obs(observations))
     G_ensemble = Array{Float64}(undef, single_member_dims, ensemble_size)
     for m in 1:ensemble_size
         member_path = ClimaCalibrate.path_to_ensemble_member(output_dir, iteration, m)
@@ -39,11 +39,11 @@ days = 86_400
 spinup_time = 93days
 function preprocess_monthly_averages(simdir, name)
     monthly_avgs = get_monthly_averages(simdir, name)
-    # TODO: Replace NaNs with global mean
-    monthly_avgs = ClimaAnalysis.replace(monthly_avgs, NaN => 0.0)
     monthly_avgs = ClimaAnalysis.shift_to_start_of_previous_month(monthly_avgs)
     # Remove spinup time
     monthly_avgs = window(monthly_avgs, "time"; left = spinup_time)
+    global_mean = monthly_avgs |> average_lat |> average_lon |> average_time # average_time(average_lon(average_lat(monthly_avgs)))
+    monthly_avgs = ClimaAnalysis.replace(monthly_avgs, NaN => first(global_mean.data))
     return monthly_avgs
 end
 

experiments/calibration/cld_eff_rad/run_calibration.jl

@@ -13,7 +13,9 @@ rng = Random.MersenneTwister(rng_seed)
 
 include(joinpath(pkgdir(ClimaCoupler), "experiments/calibration/cld_eff_rad/observation_map.jl"))
 
+ensemble_size = 10
 addprocs(CAL.SlurmManager())
+# addprocs(CAL.SlurmManager(ensemble_size); cpus_per_task = 4, gpus_per_task = 1, partition = "a3", time = "08:00:00")
 
 # Make variables and the forward model available on the worker sessions
 @everywhere begin
@@ -24,20 +26,17 @@ end
 
 # Experiment Configuration
 output_dir = "experiments/calibration/output"
-ensemble_size = 10
 n_iterations = 10
-priors = [
-    constrained_gaussian("prescribed_cloud_droplet_number_concentration", 1.5e8, 9e7, 1e7, 1e9),
-]
+priors = [constrained_gaussian("prescribed_cloud_droplet_number_concentration", 1e8, 7e7, 1e7, 1e9)]
 prior = combine_distributions(priors)
-obs_path = joinpath(experiment_dir, "observations.jld2")
-observations = JLD2.load_object(obs_path)
+observation_path = joinpath(experiment_dir, "observations.jld2")
+observations = JLD2.load_object(observation_path)
 
 eki = EKP.EnsembleKalmanProcess(
     EKP.construct_initial_ensemble(rng, prior, ensemble_size),
     observations,
     EKP.TransformInversion(),
-    verbose=true
-)
-
+    verbose = true,
+);
+# using ClimaAtmos#main
 eki = CAL.calibrate(CAL.WorkerBackend, eki, ensemble_size, n_iterations, prior, output_dir)