Develop calibration pipeline for GM calibration at half degree

Project.toml

@@ -1,23 +1,34 @@
-authors = ["Gregory Wagner <[email protected]>", "Xin Kai Lee <[email protected]>"]
 name = "ClimaOceanCalibration"
 uuid = "b5b0db0a-afc0-4eaa-813c-9b0f9c9ed209"
+authors = ["Gregory Wagner <[email protected]>", "Xin Kai Lee <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 CFTime = "179af706-886a-5703-950a-314cd64e0468"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ClimaCalibrate = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
 ClimaOcean = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
+ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+NaNStatistics = "b946abbf-3ea7-4610-9019-9858bfdeaf2d"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PkgDev = "149e707d-584d-56d3-88ec-740c18e106ff"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5"
 
 [compat]
@@ -35,4 +46,4 @@ julia = "1.10"
 
 [extras]
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
-MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"

examples/GM_calibration/average_ECCO_data.jl

@@ -0,0 +1,62 @@
+using ClimaOcean
+using Oceananigans
+using Oceananigans.Units
+using Oceananigans.Architectures: on_architecture
+using SeawaterPolynomials.TEOS10
+using ClimaOcean.DataWrangling
+using Printf
+using Dates
+using CUDA
+using ClimaOceanCalibration.DataWrangling
+using JLD2
+using XESMF
+
+arch = GPU()
+
+grid = jldopen(joinpath(pwd(), "examples", "GM_calibration", "grids_and_regridder.jld2"), "r") do file
+    return on_architecture(arch, file["target_grid"])
+end
+
+dataset = ECCO4Monthly()
+
+dir = joinpath(homedir(), "ECCO_data")
+mkpath(dir)
+start_dates = [DateTime(year) for year in 1992:2017] 
+sampling_length = 1 # year
+buoyancy_model = SeawaterBuoyancy(equation_of_state=TEOS10EquationOfState())
+
+for start_date in start_dates
+    @info "Processing data starting from $(start_date)..."
+    end_date = start_date + Year(sampling_length) - Month(1)
+
+    T = Metadata(:temperature; dataset, dir, start_date, end_date)
+    S = Metadata(:salinity; dataset, dir, start_date, end_date)
+
+    T_data = FieldTimeSeries(T, grid, time_indices_in_memory=12)
+    S_data = FieldTimeSeries(S, grid, time_indices_in_memory=12)
+
+    T_averaging = TimeAverageOperator(T_data)
+    T_averaged_fts = AveragedFieldTimeSeries(T_averaging(T_data), T_averaging, nothing)
+
+    S_averaging = TimeAverageOperator(S_data)
+    S_averaged_fts = AveragedFieldTimeSeries(S_averaging(S_data), S_averaging, nothing)
+
+    b_averaging = TimeAverageBuoyancyOperator(T_data)
+    b_averaged_fts = AveragedFieldTimeSeries(b_averaging(T_data, S_data, buoyancy_model), b_averaging, nothing)
+
+    prefix = "$(sampling_length)yearaverage_2degree"
+    date_str = replace(string(start_date), ":" => "-")
+
+    dirname = prefix * date_str
+
+    SAVE_PATH = joinpath(pwd(), "calibration_data", "ECCO4Monthly", dirname)
+    mkpath(SAVE_PATH)
+
+    T_filepath = joinpath(SAVE_PATH, "T.jld2")
+    S_filepath = joinpath(SAVE_PATH, "S.jld2")
+    b_filepath = joinpath(SAVE_PATH, "b.jld2")
+
+    save_averaged_fieldtimeseries(T_averaged_fts, T, filename=T_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(S_averaged_fts, S, filename=S_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(b_averaged_fts, nothing, filename=b_filepath, overwrite_existing=true)
+end

examples/GM_calibration/average_EN4_data.jl

@@ -0,0 +1,111 @@
+using ClimaOcean
+using Oceananigans
+using Oceananigans.Units
+using Oceananigans.Architectures: on_architecture
+using SeawaterPolynomials.TEOS10
+using ClimaOcean.DataWrangling
+using Printf
+using Dates
+using CUDA
+using ClimaOceanCalibration.DataWrangling
+using JLD2
+using XESMF
+
+arch = GPU()
+
+grid = jldopen(joinpath(pwd(), "examples", "GM_calibration", "grids_and_regridder.jld2"), "r") do file
+    return on_architecture(arch, file["target_grid"])
+end
+
+dataset = EN4Monthly()
+
+dir = joinpath(homedir(), "EN4_data")
+mkpath(dir)
+start_dates = [DateTime(1902), DateTime(1912), DateTime(1922), DateTime(1942),
+               DateTime(1952), DateTime(1972),
+               DateTime(1992), DateTime(2002), DateTime(2012)]
+
+# seems that T fields for 1939 1971 1985 is problematic
+
+buoyancy_model = SeawaterBuoyancy(equation_of_state=TEOS10EquationOfState())
+
+for start_date in start_dates
+    @info "Processing data starting from $(start_date)..."
+    end_date = start_date + Year(10) - Month(1)
+
+    T = Metadata(:temperature; dataset, dir, start_date, end_date)
+    S = Metadata(:salinity; dataset, dir, start_date, end_date)
+
+    T_data = FieldTimeSeries(T, grid, time_indices_in_memory=20)
+    S_data = FieldTimeSeries(S, grid, time_indices_in_memory=20)
+
+    T_averaging = TimeAverageOperator(T_data)
+    T_averaged_fts = AveragedFieldTimeSeries(T_averaging(T_data), T_averaging, nothing)
+
+    S_averaging = TimeAverageOperator(S_data)
+    S_averaged_fts = AveragedFieldTimeSeries(S_averaging(S_data), S_averaging, nothing)
+
+    b_averaging = TimeAverageBuoyancyOperator(T_data)
+    b_averaged_fts = AveragedFieldTimeSeries(b_averaging(T_data, S_data, buoyancy_model), b_averaging, nothing)
+
+    prefix = "10yearaverage_2degree"
+    date_str = replace(string(start_date), ":" => "-")
+
+    dirname = prefix * date_str
+
+    SAVE_PATH = joinpath(pwd(), "calibration_data", "EN4Monthly", dirname)
+    mkpath(SAVE_PATH)
+
+    T_filepath = joinpath(SAVE_PATH, "T.jld2")
+    S_filepath = joinpath(SAVE_PATH, "S.jld2")
+    b_filepath = joinpath(SAVE_PATH, "b.jld2")
+
+    save_averaged_fieldtimeseries(T_averaged_fts, T, filename=T_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(S_averaged_fts, S, filename=S_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(b_averaged_fts, nothing, filename=b_filepath, overwrite_existing=true)
+end
+
+# seems that T fields for 1939 1971 1985 is problematic
+start_dates = [DateTime(1902), DateTime(1907), DateTime(1912), DateTime(1917), DateTime(1922),
+               DateTime(1927), DateTime(1932), DateTime(1942),
+               DateTime(1947), DateTime(1952), DateTime(1957), DateTime(1962),
+               DateTime(1972), DateTime(1977),
+               DateTime(1987), DateTime(1992), DateTime(1997), DateTime(2002),
+               DateTime(2007), DateTime(2012)]
+
+for start_date in start_dates
+    @info "Processing data starting from $(start_date)..."
+    end_date = start_date + Year(5) - Month(1)
+
+    T = Metadata(:temperature; dataset, dir, start_date, end_date)
+    S = Metadata(:salinity; dataset, dir, start_date, end_date)
+
+    T_data = FieldTimeSeries(T, grid, time_indices_in_memory=20)
+    S_data = FieldTimeSeries(S, grid, time_indices_in_memory=20)
+
+    T_averaging = TimeAverageOperator(T_data)
+    T_averaged_fts = AveragedFieldTimeSeries(T_averaging(T_data), T_averaging, nothing)
+
+    S_averaging = TimeAverageOperator(S_data)
+    S_averaged_fts = AveragedFieldTimeSeries(S_averaging(S_data), S_averaging, nothing)
+
+    b_averaging = TimeAverageBuoyancyOperator(T_data)
+    b_averaged_fts = AveragedFieldTimeSeries(b_averaging(T_data, S_data, buoyancy_model), b_averaging, nothing)
+
+    prefix = "5yearaverage_2degree"
+    date_str = replace(string(start_date), ":" => "-")
+
+    dirname = prefix * date_str
+
+    SAVE_PATH = joinpath(pwd(), "calibration_data", "EN4Monthly", dirname)
+    mkpath(SAVE_PATH)
+
+    T_filepath = joinpath(SAVE_PATH, "T.jld2")
+    S_filepath = joinpath(SAVE_PATH, "S.jld2")
+    b_filepath = joinpath(SAVE_PATH, "b.jld2")
+
+    save_averaged_fieldtimeseries(T_averaged_fts, T, filename=T_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(S_averaged_fts, S, filename=S_filepath, overwrite_existing=true)
+    save_averaged_fieldtimeseries(b_averaged_fts, nothing, filename=b_filepath, overwrite_existing=true)
+end
+

examples/GM_calibration/calibrate_gm_distributed.jl

@@ -0,0 +1,134 @@
+const ensemble_size = 5
+using Distributed
+using ArgParse
+
+function parse_commandline()
+    s = ArgParseSettings()
+
+    @add_arg_table! s begin
+        "--simulation_length"
+            help = "Length of calibration simulation in years"
+            arg_type = Int
+            default = 6
+        "--sampling_length"
+            help = "Length of sampling period in years"
+            arg_type = Int
+            default = 1
+        "--zonal_average"
+            help = "Whether to perform zonal averaging in loss function"
+            arg_type = Bool
+            default = false
+        "--observation_covariance"
+            help = "Type of covariance to use (observations vs predetermined)"
+            arg_type = Bool
+            default = true
+        "--pickup"
+            help = "Pickup files for simulation spinup"
+            arg_type = Bool
+            default = false
+    end
+    return parse_args(s)
+end
+
+args = parse_commandline()
+
+# Add workers with pre-set environment variables
+nprocs = ensemble_size
+addprocs(nprocs)
+@everywhere @info "Worker $(myid())"
+@everywhere ENV["CUDA_VISIBLE_DEVICES"] = myid() - 1
+
+# Now load CUDA on all workers
+@everywhere using CUDA
+# Verify each worker sees exactly one GPU
+@everywhere println("Worker $(myid()) sees GPU: $(CUDA.NVML.index(CUDA.NVML.Device(CUDA.uuid(CUDA.device()))))")
+
+@everywhere begin
+    using ClimaCalibrate
+    using ClimaOcean
+    using ClimaOceanCalibration.DataWrangling
+    using Oceananigans
+    using EnsembleKalmanProcesses
+    using EnsembleKalmanProcesses.ParameterDistributions
+    using LinearAlgebra
+    using JLD2
+    using Glob
+    using Statistics
+    import ClimaCalibrate: generate_sbatch_script
+    include(joinpath(pwd(), "examples", "GM_calibration", "data_processing.jl"))
+    include(joinpath(pwd(), "examples", "GM_calibration", "model_interface.jl"))
+
+    args = $args
+
+    const simulation_length = args["simulation_length"]
+    const sampling_length = args["sampling_length"]
+    const zonal_average = args["zonal_average"]
+    const observation_covariance = args["observation_covariance"]
+    const pickup = args["pickup"] ? Dict("ocean" => joinpath(pwd(), "pickups", "ocean_pickup.jld2"), "sea_ice" => joinpath(pwd(), "pickups", "seaice_pickup.jld2")) : nothing
+
+    obl_closure = ClimaOcean.OceanSimulations.default_ocean_closure()
+
+    if obl_closure isa RiBasedVerticalDiffusivity
+        obl_str = "RiBased"
+    else
+        obl_str = "CATKE"
+    end
+
+    if observation_covariance
+        cov_str = "obscov"
+    else
+        cov_str = "diagcov"
+    end
+
+    const output_dir = joinpath(pwd(), "calibration_runs", "gm_$(simulation_length)yr_$(sampling_length)yravg_ecco_$(obl_str)_$(cov_str)$(zonal_average ? "_zonalavg" : "")$(pickup !== nothing ? "_pickup" : "")")
+    ClimaCalibrate.forward_model(iteration, member) = gm_forward_model(iteration, member; simulation_length, sampling_length, obl_closure, pickup)
+    ClimaCalibrate.observation_map(iteration) = gm_construct_g_ensemble(iteration, zonal_average)
+end
+
+n_iterations = 10
+
+κ_skew_prior = constrained_gaussian("κ_skew", 5e2, 3e2, 0, Inf)
+κ_symmetric_prior = constrained_gaussian("κ_symmetric", 5e2, 3e2, 0, Inf)
+
+priors = combine_distributions([κ_skew_prior, κ_symmetric_prior])
+
+obs_paths = abspath.(glob("$(sampling_length)yearaverage_2degree*", joinpath("calibration_data", "ECCO4Monthly")))
+
+calibration_target_obs_path = abspath(joinpath("calibration_data", "ECCO4Monthly", "$(sampling_length)yearaverage_2degree2007-01-01T00-00-00"))
+
+Y = hcat(process_observation.(obs_paths, no_tapering, zonal_average)...)
+
+const output_dim = size(Y, 1)
+
+n_trials = size(Y, 2)
+
+if observation_covariance
+    # the noise estimated from the samples (will have rank n_trials-1)
+    internal_cov = tsvd_cov_from_samples(Y) # SVD object
+
+    # the "5%" model error (diagonal)
+    model_error_frac = 0.05
+    data_mean = vec(mean(Y,dims=2))
+    model_error_cov = Diagonal((model_error_frac*data_mean).^2)
+
+    # regularize the model error diagonal (in case of zero entries)
+    model_error_cov += 1e-6*I
+
+    # Combine...
+    covariance = SVDplusD(internal_cov, model_error_cov)
+else
+    T_variance = 0.7^2
+    S_variance = 0.1^2
+    N_data = output_dim ÷ 2
+    covariance = Diagonal(vcat(fill(T_variance, N_data), fill(S_variance, N_data)))
+end
+
+Y_obs = Observation(Dict("samples" => process_observation(calibration_target_obs_path, taper_interior_ocean, zonal_average),
+                         "covariances" => covariance,
+                         "names" => basename(calibration_target_obs_path)))
+
+utki = EnsembleKalmanProcess(Y_obs, TransformUnscented(priors))
+
+backend = ClimaCalibrate.WorkerBackend
+
+ClimaCalibrate.calibrate(ClimaCalibrate.WorkerBackend, utki, n_iterations, priors, output_dir)