|
| 1 | +######## CONFIGS ######## |
| 2 | +# Most of the time users will just need to modify the settings below |
| 3 | +# Our goal is to make it flexible and easy to run different calibration setups |
| 4 | +# And to potentially iterate fast (e.g., low resolution, specific regions, few params, short simulation time...) |
| 5 | +######################### |
| 6 | + |
| 7 | +ENV["JULIA_WORKER_TIMEOUT"] = "1000.0" # This is a ClimaCalibrate setting. Wait time for workers. |
| 8 | + |
| 9 | +using Dates |
| 10 | +using Distributed |
| 11 | +import EnsembleKalmanProcesses as EKP |
| 12 | +import Random |
| 13 | +using ClimaLand |
| 14 | +import ClimaCalibrate: forward_model, parameter_path, path_to_ensemble_member |
| 15 | +import ClimaCalibrate as CAL |
| 16 | +rng_seed = 2 |
| 17 | +rng = Random.MersenneTwister(rng_seed) |
| 18 | +FT = Float64 |
| 19 | + |
| 20 | +# Calibrate the land model with: |
| 21 | +const variable_list = ["swu"] # variables you want to capture by adjusting your priors |
| 22 | +const n_iterations = 10 # 1 iterations takes ~ 1.5 hour with current settings (resolution, 2 year simulation) |
| 23 | +const spinup_period = Year(1) |
| 24 | +# potentially we could add time_for_calibration (currently 1 year) |
| 25 | + |
| 26 | +# Using the following priors: |
| 27 | +include(joinpath(@__DIR__, "priors.jl")) |
| 28 | +# potentially we could add loss-parameters pairing https://clima.github.io/EnsembleKalmanProcesses.jl/dev/update_groups/ |
| 29 | + |
| 30 | +# With the forward model on GPUS. Note: the forward model needs to be adjusted to read priors! |
| 31 | +ekp_process = EKP.Unscented(prior) |
| 32 | +ensemble_size = ekp_process.N_ens |
| 33 | +# Config for workers |
| 34 | +addprocs( |
| 35 | + CAL.PBSManager(ensemble_size), # simulations in parallel. User may change this to use less GPUs. |
| 36 | + q = "main", |
| 37 | + A = "UCIT0011", |
| 38 | + l_select = "1:ngpus=1:ncpus=4", |
| 39 | + l_walltime = "11:30:00", |
| 40 | + l_job_priority = "premium", |
| 41 | +) |
| 42 | +# ClimaLand (forward model) - needs to read priors! |
| 43 | +@everywhere begin |
| 44 | + using Dates # needs to be called again for workers |
| 45 | + const nelements = (50, 10) # resolution - (horizontal elements (lon,lat), vertical elements (soil discretization)) |
| 46 | + const start_date = DateTime(2008, 12, 01) # this is the start of the forward model spinup |
| 47 | + const caldir = "calibration_output_utki" |
| 48 | + using ClimaLand |
| 49 | + dir = pkgdir(ClimaLand) |
| 50 | + include(joinpath(dir, "experiments/calibration/forward_model.jl")) |
| 51 | +end |
| 52 | +@assert month(start_date + spinup_period) == 12 "The start of your calibration period should be December." |
| 53 | + |
| 54 | +# And using those locations (currently all coordinates on land): |
| 55 | +include(joinpath(@__DIR__, "make_training_locations.jl")) |
| 56 | +training_locations = make_training_locations(nelements) |
| 57 | +# potentially we can add regional runs or specific lon lat bands or filter (e.g., regions with snow) |
| 58 | + |
| 59 | +# NOTE1: Don't forget to modify your forward model (land or bucket) to read and use your priors correctly. |
| 60 | +# NOTE2: The noise is set in observationseries_era5.jl - adjust if needed. |
| 61 | +# ^ current noise options: era5 inter-annual variance, era5 seasonal mean * factor, flat noise, weigh by lats. |
| 62 | +# NOTE3: Currently everything is set to use seasonal averages. We could add option to use e.g., monthly or other. |
| 63 | +# NOTE4: We could add option to calibrate single sites (change forward model, observationseries, observation_map). |
| 64 | +# ^ maybe Julia and Thanhthanh SURF? |
| 65 | + |
| 66 | + |
| 67 | + |
| 68 | +########################## |
| 69 | +# Most of the time you won't need to change the code below. |
| 70 | +# Unless you change EKP configurations, or other specific settings. |
| 71 | +########################## |
| 72 | + |
| 73 | + |
| 74 | +# observationseries - era5 data and noise object to compare to model output in EKP (to minimize the loss) |
| 75 | +include(joinpath(@__DIR__, "observationseries_era5.jl")) |
| 76 | + |
| 77 | +# l_obs is the length of Observation objects (observationseries for era5, observationmap for ClimaLand) |
| 78 | +n_locations = length(training_locations) |
| 79 | +n_variables = length(variable_list) |
| 80 | +n_time_points = 4 # 4 seasons (and not, for example, 12 months) |
| 81 | +l_obs = n_time_points * n_variables * n_locations |
| 82 | + |
| 83 | +# build observation from ClimaLand outputs - for one member |
| 84 | +include(joinpath(@__DIR__, "observation_map.jl")) |
| 85 | + |
| 86 | +# build observation from ClimaLand outputs - for all members |
| 87 | +function CAL.observation_map(iteration) |
| 88 | + single_member_dims = (l_obs,) |
| 89 | + G_ensemble = Array{Float64}(undef, single_member_dims..., ensemble_size) |
| 90 | + |
| 91 | + for m in 1:ensemble_size |
| 92 | + member_path = path_to_ensemble_member(caldir, iteration, m) |
| 93 | + simdir_path = |
| 94 | + joinpath(member_path, "global_diagnostics", "output_active") |
| 95 | + simdir = SimDir(simdir_path) |
| 96 | + G_ensemble[:, m] .= |
| 97 | + process_member_data(simdir, training_locations, variable_list) |
| 98 | + end |
| 99 | + |
| 100 | + return G_ensemble |
| 101 | +end |
| 102 | + |
| 103 | +# Build the UTKI object - this is where you set EKP configurations |
| 104 | +utki = EKP.EnsembleKalmanProcess( |
| 105 | + observationseries, |
| 106 | + EKP.TransformUnscented(prior, impose_prior = true); |
| 107 | + verbose = true, |
| 108 | + rng, |
| 109 | + scheduler = EKP.DataMisfitController(terminate_at = 100), |
| 110 | +) |
| 111 | + |
| 112 | +# Run the calibration via ClimaCalibrate using the arguments built above |
| 113 | +CAL.calibrate(CAL.WorkerBackend, utki, n_iterations, prior, caldir) |
0 commit comments