From 25a508ddaaf0c159bae8313e91b91013230ad4fb Mon Sep 17 00:00:00 2001 From: nefrathenrici Date: Tue, 25 Feb 2025 14:03:15 -0800 Subject: [PATCH] Running with single year batch, downsampling outputvars --- experiments/calibration/Project.toml | 4 +- .../calibration/coarse_amip/generate_obs.jl | 16 ++ .../calibration/coarse_amip/model_config.yml | 84 ++++++ .../coarse_amip/model_interface.jl | 35 +++ .../coarse_amip/observation_map.jl | 79 ++++++ .../coarse_amip/observation_utils.jl | 245 ++++++++++++++++++ .../coarse_amip/run_calibration.jl | 93 +++++++ .../coarse_amip/run_calibration.sh | 13 + 8 files changed, 568 insertions(+), 1 deletion(-) create mode 100644 experiments/calibration/coarse_amip/generate_obs.jl create mode 100644 experiments/calibration/coarse_amip/model_config.yml create mode 100644 experiments/calibration/coarse_amip/model_interface.jl create mode 100644 experiments/calibration/coarse_amip/observation_map.jl create mode 100644 experiments/calibration/coarse_amip/observation_utils.jl create mode 100644 experiments/calibration/coarse_amip/run_calibration.jl create mode 100644 experiments/calibration/coarse_amip/run_calibration.sh diff --git a/experiments/calibration/Project.toml b/experiments/calibration/Project.toml index d9821cde5d..2c6608a067 100644 --- a/experiments/calibration/Project.toml +++ b/experiments/calibration/Project.toml @@ -30,5 +30,7 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" [compat] -ClimaCalibrate = "0.0.10" +ClimaCalibrate = "0.0.11" ClimaTimeSteppers = "0.8.2" +ClimaCore = "0.14.26" +EnsembleKalmanProcesses = "2.0" diff --git a/experiments/calibration/coarse_amip/generate_obs.jl b/experiments/calibration/coarse_amip/generate_obs.jl new file mode 100644 index 0000000000..a86c881d9a --- /dev/null +++ b/experiments/calibration/coarse_amip/generate_obs.jl @@ -0,0 +1,16 @@ +# Generate and save experiment observations to disk +using ClimaAnalysis, JLD2 +include("experiments/calibration/coarse_amip/observation_utils.jl") + +const obs_dir = "/home/ext_nefrathe_caltech_edu/calibration_obs" +const diagnostic_dir = "experiments/calibration/output/iteration_000/member_001/model_config/output_0000/clima_atmos/" + +diagnostic_var2d = OutputVar(joinpath(diagnostic_dir, "rsdt_1M_average.nc")); +pressure = OutputVar(joinpath(diagnostic_dir, "pfull_1M_average.nc")); +diagnostic_var3d = OutputVar(joinpath(diagnostic_dir, "ta_1M_average.nc")); +diagnostic_var3d = ClimaAnalysis.Atmos.to_pressure_coordinates(diagnostic_var3d, pressure) + +nt = get_all_output_vars(obs_dir, diagnostic_var2d, diagnostic_var3d) +nyears = 18 +observation_vec = create_observation_vector(nt, nyears) +JLD2.save_object("experiments/calibration/coarse_amip/observations.jld2", observation_vec) diff --git a/experiments/calibration/coarse_amip/model_config.yml b/experiments/calibration/coarse_amip/model_config.yml new file mode 100644 index 0000000000..703d3b1085 --- /dev/null +++ b/experiments/calibration/coarse_amip/model_config.yml @@ -0,0 +1,84 @@ +FLOAT_TYPE: "Float32" +albedo_model: "CouplerAlbedo" +atmos_config_file: "config/longrun_configs/amip_target_edonly.yml" +checkpoint_dt: "99999days" +coupler_toml: ["toml/amip.toml"] +print_config_dict: true +coupler_output_dir: "experiments/calibration/output" +dt: "800secs" +dt_cpl: "800secs" +dt_save_state_to_disk: "99999days" +dt_save_to_sol: "99999days" +dz_bottom: 100.0 +energy_check: false +h_elem: 3 +land_albedo_type: "map_temporal" +mode_name: "amip" +mono_surface: false +netcdf_output_at_levels: true +output_default_diagnostics: false +use_coupler_diagnostics: false +radiation_reset_rng_seed: true +rayleigh_sponge: true +start_date: "20000901" +surface_setup: "PrescribedSurface" +t_end: "457days" +topography: "Earth" +topo_smoothing: true +turb_flux_partition: "CombinedStateFluxesMOST" +viscous_sponge: true +z_elem: 39 +z_max: 60000.0 +extra_atmos_diagnostics: + - reduction_time: average + short_name: rsut + period: 1months + writer: nc + - reduction_time: average + short_name: rlut + period: 1months + writer: nc + - reduction_time: average + short_name: rsdt + period: 1months + writer: nc + - reduction_time: average + short_name: rsutcs + period: 1months + writer: nc + - reduction_time: average + short_name: rlutcs + period: 1months + writer: nc + - reduction_time: average + short_name: pr + period: 1months + writer: nc + - reduction_time: average + short_name: ts + period: 1months + writer: nc + - reduction_time: average + short_name: ta + period: 1months + writer: nc + - reduction_time: average + short_name: hur + period: 1months + writer: nc + - reduction_time: average + short_name: hus + period: 1months + writer: nc + - reduction_time: average + short_name: clw + period: 1months + writer: nc + - reduction_time: average + short_name: cli + period: 1months + writer: nc + - reduction_time: average + short_name: pfull + period: 1months + writer: nc diff --git a/experiments/calibration/coarse_amip/model_interface.jl b/experiments/calibration/coarse_amip/model_interface.jl new file mode 100644 index 0000000000..1780058610 --- /dev/null +++ b/experiments/calibration/coarse_amip/model_interface.jl @@ -0,0 +1,35 @@ +import ClimaCoupler +import ClimaCalibrate +import CUDA +import EnsembleKalmanProcesses as EKP +ENV["CLIMACOMMS_DEVICE"] = "CUDA" +ENV["CLIMACOMMS_CONTEXT"] = "SINGLETON" +import JLD2 +include(joinpath(pkgdir(ClimaCoupler), "experiments", "ClimaEarth", "setup_run.jl")) + +function ClimaCalibrate.forward_model(iter, member) + + config_file = joinpath(pkgdir(ClimaCoupler), "experiments", "calibration", "coarse_amip", "model_config.yml") + config_dict = get_coupler_config_dict(config_file) + + output_dir_root = config_dict["coupler_output_dir"] + eki = JLD2.load_object(ClimaCalibrate.ekp_path(output_dir_root, iter)) + minibatch = EKP.get_current_minibatch(eki) + config_dict["start_date"] = minibatch_to_start_date(minibatch) + + # Set member parameter file + sampled_parameter_file = ClimaCalibrate.parameter_path(output_dir_root, iter, member) + config_dict["calibration_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 + + return setup_and_run(config_dict) +end + + +function minibatch_to_start_date(batch) + start_year = minimum(batch) + 1999 + @assert start_year >= 2000 + return "$(start_year)0901" +end diff --git a/experiments/calibration/coarse_amip/observation_map.jl b/experiments/calibration/coarse_amip/observation_map.jl new file mode 100644 index 0000000000..1487c3564f --- /dev/null +++ b/experiments/calibration/coarse_amip/observation_map.jl @@ -0,0 +1,79 @@ +using ClimaAnalysis, Dates +import ClimaCalibrate +import ClimaCoupler +import JLD2 +import EnsembleKalmanProcesses as EKP +obs = JLD2.load_object("experiments/calibration/coarse_amip/observations.jld2") +const single_member_dims = length(EKP.get_obs(first(obs))) +include(joinpath(pkgdir(ClimaCoupler), "experiments/calibration/coarse_amip/observation_utils.jl")) + +function ClimaCalibrate.observation_map(iteration) + G_ensemble = Array{Float64}(undef, single_member_dims, ensemble_size) + for m in 1:ensemble_size + @show m + member_path = ClimaCalibrate.path_to_ensemble_member(output_dir, iteration, m) + simdir_path = joinpath(member_path, "model_config/output_active") + if isdir(simdir_path) + simdir = SimDir(simdir_path) + G_ensemble[:, m] .= process_member_data(simdir) + else + @info "No data found for member $m." + G_ensemble[:, m] .= NaN + end + end + return G_ensemble +end + +function process_outputvar(simdir, name) + days = 86_400 + + monthly_avgs = get_monthly_averages(simdir, name) + # Preprocess to match observations + if has_altitude(monthly_avgs) + pressure = get_monthly_averages(simdir, "pfull") + monthly_avgs = ClimaAnalysis.Atmos.to_pressure_coordinates(monthly_avgs, pressure) + monthly_avgs = limit_pressure_dim_to_era5_range(monthly_avgs) + end + monthly_avgs = ClimaAnalysis.replace(monthly_avgs, missing => 0.0, NaN => 0.0) + monthly_avgs = ClimaAnalysis.shift_to_start_of_previous_month(monthly_avgs) + + # Cut off first 3 months + single_year = window(monthly_avgs, "time"; left = 92days) + seasons = split_by_season_across_time(single_year) + # Ensure we are splitting evenly across seasons + @assert all(map(x -> length(times(x)) == 3, seasons)) + seasonal_avgs = average_time.(seasons) + + downsampled_seasonal_avg_arrays = downsample.(seasonal_avgs, 3) + return vcat(vec.(downsampled_seasonal_avg_arrays)...) + # return vectorize_nyears_of_seasonal_outputvars(seasonal_avgs, 1) +end + +function process_member_data(simdir::SimDir) + isempty(simdir) && return fill!(zeros(single_member_length), NaN) + + pressure = get_monthly_averages(simdir, "pfull") + + rsdt_full = get_monthly_averages(simdir, "rsdt") + rsut_full = get_monthly_averages(simdir, "rsut") + rlut_full = get_monthly_averages(simdir, "rlut") + year_net_radiation = (rlut_full + rsut_full - rsdt_full) |> average_lat |> average_lon |> average_time + + rsut = process_outputvar(simdir, "rsut") + rlut = process_outputvar(simdir, "rlut") + rsutcs = process_outputvar(simdir, "rsutcs") + rlutcs = process_outputvar(simdir, "rlutcs") + cre = rsut + rlut - rsutcs - rlutcs + + pr = process_outputvar(simdir, "pr") + # shf = process_outputvar(simdir, "shf") + ts = process_outputvar(simdir, "ts") + + ta = process_outputvar(simdir, "ta") + hur = process_outputvar(simdir, "hur") + hus = process_outputvar(simdir, "hus") + # clw = get_seasonal_averages(simdir, "clw") + # cli = get_seasonal_averages(simdir, "cli") + + return vcat(year_net_radiation.data, rsut, rlut, cre, pr, ts)#, ta, hur, hus) +end diff --git a/experiments/calibration/coarse_amip/observation_utils.jl b/experiments/calibration/coarse_amip/observation_utils.jl new file mode 100644 index 0000000000..3e03c6dca0 --- /dev/null +++ b/experiments/calibration/coarse_amip/observation_utils.jl @@ -0,0 +1,245 @@ +# Utilities for processing observations/OutputVars +# Used to generate observations and compute the observation map +using ClimaAnalysis, ClimaCoupler +using Dates, LinearAlgebra, Statistics +import EnsembleKalmanProcesses as EKP + +# Workaround to read ql, qi from file nicely +push!(ClimaAnalysis.Var.TIME_NAMES, "valid_time") + +# Constants +const days_in_seconds = 86_400 +const months = 31days_in_seconds +const years = 365days_in_seconds +const start_date = DateTime(2000, 3, 1) +const first_year_start_date = DateTime(2000, 12, 1) + +include(joinpath(pkgdir(ClimaCoupler), "experiments/ClimaEarth/leaderboard/data_sources.jl")) + +# The ERA5 pressure range is not as large as the ClimaAtmos default pressure levels, +# so we need to limit outputvars to the ERA5 dims +function limit_pressure_dim_to_era5_range(diagnostic_var3d) + @assert has_pressure(diagnostic_var3d) + era5_pressure_min = 100.0 # Pa + era5_pressure_max = 100_000.0 # Pa + pfull_dims = diagnostic_var3d.dims[pressure_name(diagnostic_var3d)] + left = minimum(filter(x -> x >= era5_pressure_min, pfull_dims)) + right = maximum(filter(x -> x <= era5_pressure_max, pfull_dims)) + # Window the diagnostic var to use the era5 pressure bounds + return window(diagnostic_var3d, "pfull"; left, right) +end + +""" + get_all_output_vars(obs_dir, diagnostic_var2d, diagnostic_var3d) + +Return a NamedTuple of OutputVars containing all initial coarse AMIP observations. +Start date is set to `DateTime(2000, 3, 1)`. All OutputVars are resampled to the model diagnostic grid. +""" +function get_all_output_vars(obs_dir, diagnostic_var2d, diagnostic_var3d) + diagnostic_var3d = limit_pressure_dim_to_era5_range(diagnostic_var3d) + + resample_2d(output_var) = resampled_as(output_var, diagnostic_var2d, dim_names = ["longitude", "latitude"]) + resample_3d(output_var) = + resampled_as(output_var, diagnostic_var3d, dim_names = ["longitude", "latitude", "pressure_level"]) + resample(ov) = has_pressure(ov) ? resample_3d(ov) : resample_2d(ov) + + era5_outputvar(path) = OutputVar(path; new_start_date = start_date, shift_by = Dates.firstdayofmonth) + + rad_and_pr_obs_dict = get_obs_var_dict() + + # 2D Fields + # TOA incoming shortwave radiation + rsdt = resample(rad_and_pr_obs_dict["rsdt"](start_date)) + # TOA outgoing long, shortwave radiation + rlut = resample(rad_and_pr_obs_dict["rlut"](start_date)) + rsut = resample(rad_and_pr_obs_dict["rsut"](start_date)) + # TOA clearsky outgoing long, shortwave radiation + rsutcs = resample(rad_and_pr_obs_dict["rsutcs"](start_date)) + rlutcs = resample(rad_and_pr_obs_dict["rlutcs"](start_date)) + + # TOA net radiative flux + net_rad = rlut + rsut - rsdt + # For some reason we need to add the start date back in + net_rad.attributes["start_date"] = string(start_date) + + # cloud radiative effect + cre = rsut + rlut - rsutcs - rlutcs + cre.attributes["start_date"] = string(start_date) + + # Precipitation + pr = resample(rad_and_pr_obs_dict["pr"](start_date)) + + # Latent heat flux + lhf = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_averages_surface_single_level_mslhf.nc"))) + # Sensible heat flux + shf = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_averages_surface_single_level_msshf.nc"))) + shf = lhf + shf + shf.attributes["start_date"] = string(start_date) + # Surface temperature + ts = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_avg_ts.nc"))) + # 3D Fields + # Air temperature + ta = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_avg_pressure_level_t.nc"))) + + # relative humidity + hur = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_avg_pressure_level_r.nc"))) + # specific humidity + hus = resample(era5_outputvar(joinpath(obs_dir, "era5_monthly_avg_pressure_level_q.nc"))) + + # # Cloud specific liquid water content + # ql = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_liquid_water_content.nc")) + # # Cloud specific ice water content + # qi = era5_outputvar(joinpath(obs_dir, "era5_specific_cloud_ice_water_content.nc")) + # foreach((ql, qi)) do var + # # Convert from hPa to Pa in-place so we don't create more huge OutputVars + # @assert var.dim_attributes[pressure_name(var)]["units"] == "hPa" + # var.dims[pressure_name(var)] .*= 100.0 + # set_dim_units!(var, pressure_name(var), "Pa") + # end + + # ql = resample(reverse_dim(reverse_dim(ql, latitude_name(ql)), pressure_name(ql))) + # qi = resample(reverse_dim(reverse_dim(qi, latitude_name(qi)), pressure_name(qi))) + + return (; rlut, rsut, rsutcs, rlutcs, pr, net_rad, cre, shf, ts, ta, hur, hus)#, ql, qi) +end + +##### +# Processing to create EKP.ObservationSeries +##### + +to_datetime(start_date, time) = DateTime(start_date) + Second(time) + +get_monthly_averages(simdir, var_name) = get(simdir; short_name = var_name, reduction = "average", period = "1M") + +get_seasonal_averages(var) = average_time.(split_by_season_across_time(var)) + +function get_seasonal_averages(simdir, var_name) + var = get(simdir; short_name = var_name, reduction = "average", period = "1M") + get_seasonal_averages(var) +end + +function get_yearly_averages(var) + seasonal_avgs = get_seasonal_averages(var) + nyears = fld(length(seasonal_avgs), 4) + matrices = getproperty.(seasonal_avgs, :data) + year_averaged_matrices = map(1:nyears) do i + start_idx = (i - 1) * 4 + 1 + end_idx = i * 4 + group = matrices[start_idx:end_idx] + + # Compute the average matrix for this group + averaged_matrix = sum(group) / 4 + averaged_matrix + end + return year_averaged_matrices +end + +# Given an outputvar, compute the standard deviation at each point for each season. +function get_seasonal_stdev(output_var) + all_seasonal_averages = get_seasonal_averages(output_var) + all_seasonal_averages = downsample.(all_seasonal_averages, 3) + seasonal_average_matrix = cat(all_seasonal_averages...; dims = 3) + interannual_stdev = std(seasonal_average_matrix, dims = 3) + # TODO: Add intraseasonal stdev? + return dropdims(interannual_stdev; dims = 3) +end + +# Given an outputvar, compute the covariance for each season. +function get_seasonal_covariance(output_var) + stdev = get_seasonal_stdev(output_var) + return Diagonal(vec(stdev) .^ 2) +end + +# Given a year, return the indices of that year within a seasonal array +# Assume each year has 4 seasons and starts at index % 4 == 1 +function get_year_indices(year) + start_index = (year * 4) - 3 + end_index = year * 4 + return start_index:end_index +end + +# Take in a vector of seasonal average OutputVars and a range or single number representing the years, +# return a vector of all data within the year range +function vectorize_nyears_of_seasonal_outputvars(vec_of_vars, year_range) + # Generate indices for all specified years + all_year_indices = vcat(get_year_indices.(year_range)...) + result = vcat(vec.(getproperty.(vec_of_vars[all_year_indices], :data))...) + return result +end + +# Make an EKP.Observation of a single year of seasonal averages from an OutputVar +function make_single_year_of_seasonal_observations(output_var, yr) + seasonal_avgs = get_seasonal_averages(output_var) + downsampled_seasonal_avg_arrays = downsample.(seasonal_avgs, 3) + all_year_indices = vcat(get_year_indices.(yr)...) + obs_vec = vcat(vec.(downsampled_seasonal_avg_arrays[all_year_indices])...) + + name = get(output_var.attributes, "CF_name", get(output_var.attributes, "long_name", "")) + cov = get_seasonal_covariance(output_var) + return EKP.Observation(obs_vec, Diagonal(repeat(cov.diag, 4)), "$(yr)_$name") +end + +""" + create_observation_vector(nt, yrs = 19) + +""" +function create_observation_vector(nt, yrs = 19) + # Starting year is 2000-12 to 2001-11 + t_start = Second(first_year_start_date - start_date).value + rsut = window(nt.rsut, "time"; left = t_start) + rlut = window(nt.rlut, "time"; left = t_start) + rsutcs = window(nt.rsutcs, "time"; left = t_start) + rlutcs = window(nt.rlutcs, "time"; left = t_start) + cre = window(nt.cre, "time"; left = t_start) + + # Net radiation uses yearly averages, so we treat it differently + net_rad = window(nt.net_rad, "time"; left = t_start) |> average_lat |> average_lon + net_rad = get_yearly_averages(net_rad) + net_rad_stdev = std(cat(net_rad..., dims = 3), dims = 3) + net_rad_covariance = Diagonal(vec(net_rad_stdev) .^ 2) + + ts = window(nt.ts, "time"; left = t_start) + pr = window(nt.pr, "time"; left = t_start) + shf = window(nt.shf, "time"; left = t_start) + + ta = window(nt.ta, "time"; left = t_start) + + hur = window(nt.hur, "time"; left = t_start) + hus = window(nt.hus, "time"; left = t_start) + + all_observations = map(1:yrs) do yr + net_rad_obs = EKP.Observation(vec(net_rad[yr]), net_rad_covariance, "$(yr)_net_rad") + + rsut_obs = make_single_year_of_seasonal_observations(rsut, yr) + rlut_obs = make_single_year_of_seasonal_observations(rlut, yr) + rsutcs_obs = make_single_year_of_seasonal_observations(rsutcs, yr) + rlutcs_obs = make_single_year_of_seasonal_observations(rlutcs, yr) + cre_obs = make_single_year_of_seasonal_observations(cre, yr) + pr_obs = make_single_year_of_seasonal_observations(pr, yr) + # shf_obs = make_single_year_of_seasonal_observations(shf, yr) + ts_obs = make_single_year_of_seasonal_observations(ts, yr) + + # ta_obs = make_single_year_of_seasonal_observations(ta, yr) + # hur_obs = make_single_year_of_seasonal_observations(hur, yr) + # hus_obs = make_single_year_of_seasonal_observations(hus, yr) + EKP.combine_observations([net_rad_obs, rsut_obs, rlut_obs, cre_obs, pr_obs, ts_obs])#, ta_obs, hur_obs, hus_obs]) + end + + return all_observations # NOT an EKP.ObservationSeries +end + +downsample(var::ClimaAnalysis.OutputVar, n) = downsample(var.data, n) + +function downsample(arr::AbstractArray, n) + if n < 1 + error("Downsampling factor n must be at least 1.") + end + if ndims(arr) == 2 + return arr[1:n:end, 1:n:end] + elseif ndims(arr) == 3 + return arr[1:n:end, 1:n:end, :] + else + error("Only 2D and 3D arrays are supported.") + end +end + diff --git a/experiments/calibration/coarse_amip/run_calibration.jl b/experiments/calibration/coarse_amip/run_calibration.jl new file mode 100644 index 0000000000..8f5889df77 --- /dev/null +++ b/experiments/calibration/coarse_amip/run_calibration.jl @@ -0,0 +1,93 @@ +using Distributed +import ClimaCalibrate as CAL +using ClimaCalibrate +using ClimaAnalysis +import ClimaAnalysis: SimDir, get, slice, average_xy +import ClimaComms +import ClimaCoupler +using EnsembleKalmanProcesses.ParameterDistributions +import EnsembleKalmanProcesses as EKP + +include(joinpath(pkgdir(ClimaCoupler), "experiments/calibration/coarse_amip/observation_map.jl")) + +ENV["JULIA_WORKER_TIMEOUT"] = "300.0" +# addprocs(CAL.SlurmManager()) +# Make variables and the forward model available on the worker sessions +@everywhere begin + import ClimaCoupler + experiment_dir = joinpath(pkgdir(ClimaCoupler), "experiments/calibration/coarse_amip/") + include(joinpath(experiment_dir, "model_interface.jl")) +end + +# Experiment Configuration +output_dir = "experiments/calibration/output" +ensemble_size = 20 +n_iterations = 18 # Cycle through all data +priors = [constrained_gaussian("liquid_cloud_effective_radius", 14e-6, 6e-6, 2.5e-6, 21.5e-6)] +prior = combine_distributions(priors) +observation_vec = JLD2.load_object(joinpath(experiment_dir, "observations.jld2")) + +batch_size = 1 +num_batches = cld(length(observation_vec), batch_size) +batches = map(1:num_batches) do i + start_idx = (i - 1) * batch_size + 1 + end_idx = min(i * batch_size, length(observation_vec)) + start_idx:end_idx +end +minibatcher = EKP.FixedMinibatcher(batches) + +series_names = string.(1:length(observation_vec)) +observation_series = EKP.ObservationSeries(observation_vec, minibatcher, series_names) + +eki = EKP.EnsembleKalmanProcess( + EKP.construct_initial_ensemble(prior, ensemble_size), + observation_series, + EKP.TransformInversion(), +) + +eki = CAL.calibrate(CAL.WorkerBackend, eki, ensemble_size, n_iterations, prior, output_dir) + +# Postprocessing +import Statistics: var, mean +using Test +import CairoMakie + +function scatter_plot(eki::EKP.EnsembleKalmanProcess) + f = CairoMakie.Figure(resolution = (800, 600)) + ax = CairoMakie.Axis(f[1, 1], ylabel = "Parameter Value", xlabel = "Top of atmosphere radiative SW flux") + + g = vec.(EKP.get_g(eki; return_array = true)) + params = vec.((EKP.get_ϕ(prior, eki))) + + for (gg, uu) in zip(g, params) + CairoMakie.scatter!(ax, gg, uu) + end + + CairoMakie.vlines!(ax, observations, linestyle = :dash) + + output = joinpath(output_dir, "scatter.png") + CairoMakie.save(output, f) + return output +end + +function param_versus_iter_plot(eki::EKP.EnsembleKalmanProcess) + f = CairoMakie.Figure(resolution = (800, 600)) + ax = CairoMakie.Axis(f[1, 1], ylabel = "Parameter Value", xlabel = "Iteration") + params = EKP.get_ϕ(prior, eki) + for (i, param) in enumerate(params) + CairoMakie.scatter!(ax, fill(i, length(param)), vec(param)) + end + + output = joinpath(output_dir, "param_vs_iter.png") + CairoMakie.save(output, f) + return output +end + +scatter_plot(eki) +param_versus_iter_plot(eki) + +params = EKP.get_ϕ(prior, eki) +spread = map(var, params) + +# Spread should be heavily decreased as particles have converged +@test last(spread) / first(spread) < 0.1 diff --git a/experiments/calibration/coarse_amip/run_calibration.sh b/experiments/calibration/coarse_amip/run_calibration.sh new file mode 100644 index 0000000000..d9727dde00 --- /dev/null +++ b/experiments/calibration/coarse_amip/run_calibration.sh @@ -0,0 +1,13 @@ +#!/bin/bash + +#SBATCH --partition=a3 +#SBATCH --output="run_calibration.txt" +#SBATCH --time=24:00:00 +#SBATCH --ntasks=20 +#SBATCH --gpus-per-task=1 +#SBATCH --cpus-per-task=4 + +julia --project=experiments/calibration -e 'using Pkg; Pkg.develop(;path="."); Pkg.instantiate(;verbose=true)' + +julia --project=experiments/calibration experiments/calibration/coarse_amip/run_calibration.jl +