From 55157de9a125cf7b7417d428906eb1ee760f17eb Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 13 May 2026 16:41:43 -0700 Subject: [PATCH 01/86] some initial files from EKP --- examples/EKIRace/Lorenz63.jl | 138 +++++++++++ examples/EKIRace/Lorenz96.jl | 192 +++++++++++++++ examples/EKIRace/Project.toml | 12 + examples/EKIRace/calibrate_l96.jl | 377 ++++++++++++++++++++++++++++++ 4 files changed, 719 insertions(+) create mode 100644 examples/EKIRace/Lorenz63.jl create mode 100644 examples/EKIRace/Lorenz96.jl create mode 100644 examples/EKIRace/Project.toml create mode 100644 examples/EKIRace/calibrate_l96.jl diff --git a/examples/EKIRace/Lorenz63.jl b/examples/EKIRace/Lorenz63.jl new file mode 100644 index 000000000..1cff36af5 --- /dev/null +++ b/examples/EKIRace/Lorenz63.jl @@ -0,0 +1,138 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using StatsPlots +using Plots +using Random +using JLD2 +using Statistics + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +# This will change for different Lorenz simulators +struct LorenzConfig{FT1 <: Real, FT2 <: Real} + "Length of a fixed integration timestep" + dt::FT1 + "Total duration of integration (T = N*dt)" + T::FT2 +end + +# This will change for each ensemble member +struct EnsembleMemberConfig{VV <: AbstractVector} + "rho, beta (unknowns)" + u::VV +end + +# This will change for different "Observations" of Lorenz +struct ObservationConfig{FT1 <: Real, FT2 <: Real} + "initial time to gather statistics (T_start = N_start*dt)" + T_start::FT1 + "end time to gather statistics (T_end = N_end*dt)" + T_end::FT2 +end + +######################################################################### +############################ Model Functions ############################ +######################################################################### + +# Forward pass of forward model +# Inputs: +# - params: structure with u (unknowns vector) +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_forward( + params::EnsembleMemberConfig, + x0::VorM, + config::LorenzConfig, + observation_config::ObservationConfig, +) where {VorM <: AbstractVecOrMat} + # run the Lorenz simulation + xn = lorenz_solve(params, x0, config) + # Get statistics + gt = stats(xn, config, observation_config) + return gt +end + +#Calculates statistics for forward model output +# Inputs: +# - xn: timeseries of states for length of simulation through Lorenz63 +function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat} + T_start = observation_config.T_start + T_end = observation_config.T_end + dt = config.dt + N_start = Int(ceil(T_start / dt)) + N_end = Int(ceil(T_end / dt)) + xn_stat = xn[:, N_start:N_end] + N_state = size(xn_stat, 1) + gt = zeros(9) # Might want to switch to more general statement? + gt[1:3] = mean(xn_stat, dims = 2) + xn_stat_cov = cov(xn_stat, dims = 2) + gt[4:6] = diag(xn_stat_cov) + gt[7:8] = xn_stat_cov[1, 2:3] + gt[9] = xn_stat_cov[2, 3] + return gt +end + +# Forward pass of the Lorenz 96 model +# Inputs: +# - params: structure with u (unknowns vector) +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + # Initialize + nstep = Int(ceil(config.T / config.dt)) + state_dim = isa(x0, AbstractVector) ? length(x0) : size(x0, 1) + xn = zeros(size(x0, 1), nstep + 1) + xn[:, 1] = x0 + + # March forward in time + for j in 1:nstep + xn[:, j + 1] = RK4(params, xn[:, j], config) + end + # Output + return xn +end + +# Lorenz 96 system +# f = dx/dt +# Inputs: +# - params: structure with u (unknowns vector) +# - x: current state +function f(params::EnsembleMemberConfig, x::VorM) where {VorM <: AbstractVecOrMat} + u = params.u + N = length(x) + f = zeros(N) + + f[1] = 10.0 * (x[2] - x[1]) + f[2] = x[1] * (u[1] - x[3]) - x[2] + f[3] = x[1] * x[2] - u[2] * x[3] + + # Output + return f +end + +# RK4 solve +# Inputs: +# - params: structure with F (state-dependent-forcing vector) +# - xold: current state +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function RK4(params::EnsembleMemberConfig, xold::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + N = length(xold) + dt = config.dt + + # Predictor steps (note no time-dependence is needed here) + k1 = f(params, xold) + k2 = f(params, xold + k1 * dt / 2.0) + k3 = f(params, xold + k2 * dt / 2.0) + k4 = f(params, xold + k3 * dt) + # Step + xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + # Output + return xnew +end diff --git a/examples/EKIRace/Lorenz96.jl b/examples/EKIRace/Lorenz96.jl new file mode 100644 index 000000000..da75d4f93 --- /dev/null +++ b/examples/EKIRace/Lorenz96.jl @@ -0,0 +1,192 @@ +# Import modules +using LinearAlgebra +using Statistics +using Random, Distributions +using Flux + + + +# This will change for different Lorenz simulators +struct LorenzConfig{FT1 <: Real, FT2 <: Real} + "Length of a fixed integration timestep" + dt::FT1 + "Total duration of integration (T = N*dt)" + T::FT2 +end + +# This will change for each ensemble member +abstract type EnsembleMemberConfig end +# struct EnsembleMemberConfig{FT} +# val::FT +# end + +# Sub-type of ensemble config for constant forcing +struct ConstantEMC{FT <: Real} <: EnsembleMemberConfig + val::FT +end +build_forcing(::T, val::FT, args...) where {T <: ConstantEMC, FT <: Real} = ConstantEMC(val) +build_forcing(::T, val::FT, args...) where {T <: ConstantEMC, FT <: AbstractVector} = ConstantEMC(val[1]) + +# Sub-type of ensemble config for spatially-dependent forcing +struct VectorEMC{VV <: AbstractVector} <: EnsembleMemberConfig + val::VV +end +build_forcing(::T, val::VV, args...) where {T <: VectorEMC, VV <: AbstractVector} = VectorEMC(val) + +# Sub-type of ensemble config for spatially-dependent forcing with neural network approximation +struct FluxEMC{FC <: Flux.Chain, VV <: AbstractVector} <: EnsembleMemberConfig + model::FC + sample_range::VV +end +function build_forcing(::T, params, model, sample_range) where {T <: FluxEMC} + _, reconstructor = Flux.destructure(model) + return FluxEMC(reconstructor(params), Float32.(sample_range)) +end + +# Constant-global +forcing(params::ConstantEMC, x, i) = params.val +forcing(params::ConstantEMC, x) = repeat([params.val], length(x)) + +# Constant-vector +forcing(params::VectorEMC, x, i) = params.val[i] +forcing(params::VectorEMC, x) = params.val + +# Flux +forcing(params::FluxEMC, x, i) = Float64(params.model([params.sample_range[i]])[1]) +forcing(params::FluxEMC, x) = Float64.(params.model([sr])[1] for sr in params.sample_range) + + + + +# This will change for different "Observations" of Lorenz +struct ObservationConfig{FT1 <: Real, FT2 <: Real} + "initial time to gather statistics (T_start = N_start*dt)" + T_start::FT1 + "end time to gather statistics (T_end = N_end*dt)" + T_end::FT2 +end + +######################################################################### +############################ Model Functions ############################ +######################################################################### + +# Forward pass of forward model +# Inputs: +# - params: structure with F +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_forward( + params::EnsembleMemberConfig, + x0::VorM, + config::LorenzConfig, + observation_config::ObservationConfig, +) where {VorM <: AbstractVecOrMat} + # run the Lorenz simulation + xn = lorenz_solve(params, x0, config) + # Get statistics + gt = stats(xn, config, observation_config) + return gt +end + +#Calculates statistics for forward model output +# Inputs: +# - xn: timeseries of states for length of simulation through Lorenz96 +function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat} + T_start = observation_config.T_start + T_end = observation_config.T_end + dt = config.dt + N_start = Int(ceil(T_start / dt)) + N_end = Int(ceil(T_end / dt)) + xn_stat = xn[:, N_start:N_end] + N_state = size(xn_stat, 1) + gt = zeros(2 * N_state) + gt[1:N_state] = mean(xn_stat, dims = 2) + gt[(N_state + 1):(2 * N_state)] = std(xn_stat, dims = 2) + return gt +end + +# Forward pass of the Lorenz 96 model +# Inputs: +# - params: structure with F +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + # Initialize + nstep = Int(ceil(config.T / config.dt)) + state_dim = isa(x0, AbstractVector) ? length(x0) : size(x0, 1) + xn = zeros(size(x0, 1), nstep + 1) + xn[:, 1] = x0 + + # March forward in time + forcing_vec = forcing(params, x0) # not state dependent so evaluate once here + for j in 1:nstep + xn[:, j + 1] = RK4(forcing_vec, xn[:, j], config) + end + # Output + return xn +end + +# Lorenz 96 system +# f = dx/dt +# Inputs: +# - params: structure with F +# - x: current state +function f(forcing_vec::VV, x::VorM) where {VV <: AbstractVector, VorM <: AbstractVecOrMat} + N = length(x) + f = zeros(N) + # Loop over N positions + for i in 3:(N - 1) + f[i] = -x[i - 2] * x[i - 1] + x[i - 1] * x[i + 1] - x[i] + forcing_vec[i] + end + # Periodic boundary conditions + f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + forcing_vec[1] + f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + forcing_vec[2] + f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + forcing_vec[N] + + # Output + return f +end + +# RK4 solve +# Inputs: +# - params: structure with F +# - xold: current state +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function RK4(forcing_vec::VV, xold::VorM, config::LorenzConfig) where {VV <: AbstractVector, VorM <: AbstractVecOrMat} + N = length(xold) + dt = config.dt + + # Predictor steps (note no time-dependence is needed here) + k1 = f(forcing_vec, xold) + k2 = f(forcing_vec, xold + k1 * dt / 2.0) + k3 = f(forcing_vec, xold + k2 * dt / 2.0) + k4 = f(forcing_vec, xold + k3 * dt) + # Step + xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + # Output + return xnew +end + + +# Neural network functions +function train_network(model, x_train, y_train) + loss(model, x, y) = Flux.Losses.mse(model(x), y) + + # Reshape x_train and y_train for Flux compatibility + x_train = reshape(x_train, 1, :) + y_train = reshape(y_train, 1, :) + x_train = Float32.(x_train) + y_train = Float32.(y_train) + + opt = Flux.setup(Adam(), model) + data = Flux.DataLoader((x_train, y_train), batchsize = 32, shuffle = true) # train the model + + # Train the model over multiple epochs + epochs = 5000 + for epoch in 1:epochs + Flux.train!(loss, model, data, opt) + end + + params, _ = Flux.destructure(model) + return model, params +end diff --git a/examples/EKIRace/Project.toml b/examples/EKIRace/Project.toml new file mode 100644 index 000000000..889f334d6 --- /dev/null +++ b/examples/EKIRace/Project.toml @@ -0,0 +1,12 @@ +[deps] +BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" +CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl new file mode 100644 index 000000000..4e1fc02ee --- /dev/null +++ b/examples/EKIRace/calibrate_l96.jl @@ -0,0 +1,377 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates + + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +include("Lorenz96.jl") # Contains Lorenz 96 source code + +verbose_flag = false +save_all_ekp = true +######################################################################## +############### Choose problem type and structure ###################### +######################################################################## + +cases = ["const-force", "vec-force", "flux-force"] # problem types +# User specifications +case = cases[1] # choose problem type +N_ens_sizes = [80, 90, 100] # list of number of ensemble members (should be problem dependent) +N_iter = 20 # maximum number of EKI iterations allowed +target_rmse = 1.0 # target RMSE +n_repeats = 4 +rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds +@info "Running $case case" +@info "Maximum number of EKI iterations: $N_iter" +@info "RMSE target: $target_rmse" +# saved and loaded for plotting etc. +configuration = Dict( + "case" => case, + "N_iter" => N_iter, + "N_ens_sizes" => N_ens_sizes, + "target_rmse" => target_rmse, + "rng_seeds" => rng_seeds, +) + +if case == "const-force" + nx = 40 #dimensions of parameter vector + nu = 1 + phi = ConstantEMC(8.0) # forcing + phi_structure = nothing + sample_range = nothing + prior = constrained_gaussian("φ", 10.0, 4.0, 0, Inf) + T = 14.0 + inff = 2 + +elseif case == "vec-force" + nx = 40 # dimensions of parameter vector + nu = nx + sinusoid = 8 .+ 6 * sin.((4 * pi * range(0, stop = nx - 1, step = 1)) / nx) + phi = VectorEMC(sinusoid) + phi_structure = nothing + sample_range = nothing + pl, psig = 2.0, 3.0 + prior_cov = zeros(nx, nx) # prior covariance + for ii in 1:nx + for jj in 1:nx + prior_cov[ii, jj] = psig^2 * exp(-abs(ii - jj) / pl) + end + end + prior_mean = 8.0 * ones(nx) # prior mean + #Creating prior distribution + distribution = Parameterized(MvNormal(prior_mean, prior_cov)) + constraint = repeat([no_constraint()], 40) + name = "ml96_prior" + prior = ParameterDistribution(distribution, constraint, name) + T = 54.0 + inff = 2 + +elseif case == "flux-force" + nx = 100 #dimensions of parameter vector + nu = 61 + true_sinusoid(x) = 8 .+ 6 * sin.((4 * pi * x) / 10) + x_train = collect(-5.0:0.01:5.0) + y_train = true_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) + phi_structure = Chain(Dense(1 => 20, tanh), Dense(20 => 1)) + true_model, _ = train_network(phi_structure, x_train, y_train) + sample_range = Float32.(collect(-5.0:0.1:4.9)) + phi = FluxEMC(true_model, sample_range) + + prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) + prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) + prior_model, prior_mean = train_network(phi_structure, x_train, prior_train) + + y_pred = true_model(reshape(sample_range, 1, :)) + prior_pred = prior_model(reshape(sample_range, 1, :)) + # Plot for visualizing training results + # p = plot(sample_range, true_sinusoid.(sample_range), label="True function", lw=2, color=:blue) + # scatter!(p, x_train, y_train, label="Training data", ms=1, color=:red, alpha=0.5) + # plot!(p, sample_range, y_pred[1,:], label="Predicted function", lw=2, color=:green) + # xlabel!("x") + # ylabel!("y") + # title!("Offline 1D DNN") + # display(p) + + # p = plot(sample_range, prior_sinusoid.(sample_range), label="True prior function", lw=2, color=:blue) + # scatter!(p, x_train, prior_train, label="Training data", ms=1, color=:red, alpha=0.5) + # plot!(p, sample_range, prior_pred[1,:], label="Predicted prior function", lw=2, color=:green) + # xlabel!("x") + # ylabel!("y") + # title!("Offline 1D DNN") + # display(p) + + # p = plot(sample_range, prior_sinusoid.(sample_range), label="Prior function", lw=2, color=:blue) + # plot!(p, sample_range, true_sinusoid.(sample_range), label="True function", lw=2, color=:green) + # xlabel!("x") + # ylabel!("y") + # title!("Offline 1D DNN") + # display(p) + prior_cov = (0.1^2) * I(length(prior_mean)) + distribution = Parameterized(MvNormal(prior_mean, prior_cov)) + constraint = repeat([no_constraint()], 61) + name = "l96_nn_prior" + prior = ParameterDistribution(distribution, constraint, name) + T = 54.0 + inff = 2.5 + +end + + +######################################################################## +############################ Problem setup ############################# +######################################################################## +rng_seed_init = 11 +rng_i = MersenneTwister(rng_seed_init) + +output_dir = joinpath(@__DIR__, "output") +if !isdir(output_dir) + mkdir(output_dir) +end + +ny = nx * 2 #number of data points + +#Creating my sythetic data +prelim_file = joinpath(output_dir, "l96_computed_preliminaries_$(case).jld2") +if isfile(prelim_file) + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + R_inv_var = loaded_data["R_inv_var"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + @info "loaded precomputed preliminary quantities from $(prelim_file)" +else + t = 0.01 #time step + T_long = 1000.0 #total time + lorenz_config_settings = LorenzConfig(t, T) + @info "Initial spin up stage: Running $( (T_long)) time units" + picking_initial_condition = LorenzConfig(t, T_long) + x_initial = rand(rng_i, Normal(0.0, 1.0), nx) # initial condition for spinning up Lorenz system + x_spun_up = lorenz_solve(phi, x_initial, picking_initial_condition) # spinning up Lorenz system + + x0 = x_spun_up[:, end] #last element of the run is the initial condition for creating the data + + # construct how we compute Observations + T_start = 4.0 #2*max + T_end = T + observation_config = ObservationConfig(T_start, T_end) + @info "Compute synthetic data: Running $(T) time units" + y = lorenz_forward(phi, x0, lorenz_config_settings, observation_config) # synthetic data + + #Observation covariance R + window = T_end - T_start + T_R = 10 * window * ny + T_start + R_config = LorenzConfig(t, T_R) + @info "Estimating noise covariance: Running $(T_R) time units" + R_run = lorenz_solve(phi, x_initial, R_config) + + R_sample_size = Int(ceil(10 * ny)) + R_samples = zeros(ny, R_sample_size) + for ii in 1:R_sample_size + local_obs_config = ObservationConfig(T_start + (ii - 1) * window, T_start + ii * window) + R_samples[:, ii] = stats(R_run, R_config, local_obs_config) + end + R = cov(R_samples, dims = 2) * inff + R_sqrt = sqrt(R) + R_inv_var = sqrt(inv(R)) + + # Need a way to perturb the initial condition when doing the EKI updates + # Solving for initial condition perturbation covariance + covT = 2000.0 #time to simulate to calculate a covariance matrix of the system + @info "Estimating initial condition covariance: Running $(covT) time units" + cov_solve = lorenz_solve(phi, x0, LorenzConfig(t, covT)) + ic_cov = 0.1 * cov(cov_solve, dims = 2) + ic_cov_sqrt = sqrt(ic_cov) + + JLD2.save( + prelim_file, + "x0", + x0, + "y", + y, + "lorenz_config_settings", + lorenz_config_settings, + "observation_config", + observation_config, + "ic_cov_sqrt", + ic_cov_sqrt, + "R", + R, + "R_inv_var", + R_inv_var, + ) + @info "saving computed quantities in $prelim_file" +end + + +######################################################################## +########################### Running EKI Race ########################### +######################################################################## + +# Counters +conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how many iterations it takes to converge (per algorithm, per rand seed, per ense size) +final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) +final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) + +method_names = [ + ("Inversion(prior)", "TEKI"), + ("TransformInversion(prior)", "ETKI"), + ("GaussNewtonInversion(prior)", "GNKI"), + ("Unscented(prior; impose_prior=true)", "UKI"), +] + +for (rr, rng_seed) in enumerate(rng_seeds) + @info "Random seed: $(rng_seed)" + rng = MersenneTwister(rng_seed) + + for (ee, N_ens) in enumerate(N_ens_sizes) + # initial parameters: N_params x N_ens + initial_params = construct_initial_ensemble(rng, prior, N_ens) + methods = [ + Inversion(prior), + TransformInversion(prior), + GaussNewtonInversion(prior), + Unscented(prior; impose_prior = true), + ] + + @info "Ensemble size: $(N_ens)" + for (kk, method) in enumerate(methods) + @info "Method: $(nameof(typeof(method)))" + if isa(method, Unscented) + ekpobj = EKP.EnsembleKalmanProcess( + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + accelerator = DefaultAccelerator(), + localization_method = NoLocalization(), + scheduler = DefaultScheduler(), + ) + else + ekpobj = EKP.EnsembleKalmanProcess( + initial_params, + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + accelerator = DefaultAccelerator(), + localization_method = NoLocalization(), + scheduler = DefaultScheduler(), + ) + end + Ne = get_N_ens(ekpobj) + + count = 0 + for i in 1:N_iter + params_i = get_ϕ_final(prior, ekpobj) + + # Calculating RMSE_e + ens_mean = mean(params_i, dims = 2)[:] + # forcing = build_forcing(ens_mean, phi_structure) + forcing = build_forcing(phi, ens_mean, phi_structure, sample_range) + G_ens_mean = lorenz_forward( + forcing, + x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + RMSE_e = norm(R_inv_var * (y - G_ens_mean[:])) / sqrt(size(y, 1)) + @info "RMSE (at G(u_mean)): $(RMSE_e)" + # Convergence criteria + if RMSE_e < target_rmse + conv_alg_iters[kk, ee, rr] = count * Ne + final_parameters[kk, ee, rr, :] = ens_mean + final_model_output[kk, ee, rr, :] = G_ens_mean + break + end + + # If RMSE convergence criteria is not satisfied + G_ens = hcat( + [ + lorenz_forward( + build_forcing(phi, params_i[:, j], phi_structure, sample_range), + (x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, Ne))[:, j], + lorenz_config_settings, + observation_config, + ) for j in 1:Ne + ]..., + ) + # Update + EKP.update_ensemble!(ekpobj, G_ens) + count = count + 1 + + + + end + final_ensemble = get_ϕ_final(prior, ekpobj) + + # save ekp files + per_method_dir = joinpath(output_dir,"$(nameof(typeof(method)))_$(today())") + if rr == 1 && ee == 1 + if !isdir(per_method_dir) + mkpath(per_method_dir) + end + prior_filename = joinpath(per_method_dir,"priors_l96.jld2") + JLD2.save(prior_filename,"prior", prior) + end + if save_all_ekp + # JLD2 + ekp_filename = joinpath(per_method_dir, "l96_ekp_$(case)_$(N_ens)_$(rr).jld2") + JLD2.save( + ekp_filename, + "N_ens", N_ens, + "method", method, + "ekpobj", ekpobj, + ) + results_filename = joinpath(per_method_dir, "l96_calibrate_results_$(case)_$(N_ens)_$(rr).jld2") + u_stored = get_u(ekpobj, return_array = false) + g_stored = get_g(ekpobj, return_array = false) + + JLD2.save( + results_filename, + "y", y, + "R", R, + "inputs", u_stored, + "outputs", g_stored, + "truth_params_structure", (phi, phi_structure) + ) + end + + end + + end +end + + +# Saving summary data: +data_filename = joinpath(output_dir, "l96_calibrate_$(case)_$(today()).jld2") +JLD2.save( + data_filename, + "configuration", + configuration, + "method_names", + method_names, + "conv_alg_iters", + conv_alg_iters, + "final_parameters", + final_parameters, + "final_model_output", + final_model_output, +) From 2e8555bdc653752920480b7ddaae30377443c385 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 18 May 2026 16:11:45 -0700 Subject: [PATCH 02/86] improved more filenaming etc. --- examples/EKIRace/calibrate_l96.jl | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index 4e1fc02ee..efdf9f47d 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -27,8 +27,16 @@ save_all_ekp = true cases = ["const-force", "vec-force", "flux-force"] # problem types # User specifications -case = cases[1] # choose problem type -N_ens_sizes = [80, 90, 100] # list of number of ensemble members (should be problem dependent) +#case = cases[1] # choose problem type +#N_ens_sizes = [5, 15, 30] # e.g. for case 1 + +#case = cases[2] # choose problem type +#N_ens_sizes = [50, 75, 100] + +case = cases[3] # choose problem type +N_ens_sizes = [50, 75, 100] # e.g. for case 1 + + N_iter = 20 # maximum number of EKI iterations allowed target_rmse = 1.0 # target RMSE n_repeats = 4 @@ -252,6 +260,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) @info "Ensemble size: $(N_ens)" for (kk, method) in enumerate(methods) @info "Method: $(nameof(typeof(method)))" + ekp_dt = 0.2 if isa(method, Unscented) ekpobj = EKP.EnsembleKalmanProcess( y, @@ -259,9 +268,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, - accelerator = DefaultAccelerator(), +# accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(), + scheduler = DataMisfitController(terminate_at=100), ) else ekpobj = EKP.EnsembleKalmanProcess( @@ -271,9 +280,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, - accelerator = DefaultAccelerator(), + # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(), + scheduler = DataMisfitController(terminate_at=100), ) end Ne = get_N_ens(ekpobj) @@ -328,7 +337,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) if !isdir(per_method_dir) mkpath(per_method_dir) end - prior_filename = joinpath(per_method_dir,"priors_l96.jld2") + prior_filename = joinpath(per_method_dir,"l96_priors_$(case).jld2") JLD2.save(prior_filename,"prior", prior) end if save_all_ekp From ceaa3894d93ab9142d561185c4bdfa5f13edaf9e Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 18 May 2026 16:12:43 -0700 Subject: [PATCH 03/86] initial emulate_sample and plot_forcing for l96 (currently does one case at a time) --- examples/EKIRace/emulate_sample_l96.jl | 273 +++++++++++++++++++++++ examples/EKIRace/plot_l96_forcing.jl | 292 +++++++++++++++++++++++++ 2 files changed, 565 insertions(+) create mode 100644 examples/EKIRace/emulate_sample_l96.jl create mode 100644 examples/EKIRace/plot_l96_forcing.jl diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl new file mode 100644 index 000000000..bddcd3f07 --- /dev/null +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -0,0 +1,273 @@ +# Import modules +include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) + +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +ENV["GKSwstype"] = "100" +using StatsPlots +using Plots +using Random +using JLD2 +using Dates + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers + +include("Lorenz96.jl") # Contains Lorenz 96 source code + + +function main() + + method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] + force_cases = ["const-force", "vec-force", "flux-force"] # problem types + cases = [ + "GP", + "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs + "RF-nonsep", # diagonalize, train scalar RF, don't asume diag inputs + ] + + #### CHOOSE YOUR CASE: (todo: looped over in some fashion) + # calib_data_dir + method=method_cases[1] + calibrate_date=Date("2026-05-18", "yyyy-mm-dd") + calib_directory="$(method)_$(calibrate_date)" + + # calib_filename_suffix + force_case=force_cases[3] + N_ens=100 + rng_idx=1 + calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + + # emulate_sample cases + case = cases[2] # e.g. 1:2 or [3] + + @info("case: ", case) + @info("force_case: ", force_case) + min_iter = 1 + skip_iter = 1 + max_iter = 10 # number of EKP iterations to use data from is at most this + + #### + + # need to loop over + # method, rng_index, N_ens = "TransformInversion", "1", "100"] + exp_name = "Lorenz_histogram_l96" + rng_seed = 44011 + rng = Random.MersenneTwister(rng_seed) + + # loading relevant data + homedir = pwd() + println(homedir) + figure_save_directory = joinpath(homedir, "output/", calib_directory) + data_save_directory = joinpath(homedir, "output", calib_directory) + loaded_calib_files = [ + joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, "l96_priors_$(force_case).jld2"), + joinpath(data_save_directory, "l96_calibrate_results_$(calib_filename_suffix).jld2"), + ] + if !isfile(loaded_calib_files[1]) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + + ekpobj = load(loaded_calib_files[1])["ekpobj"] + priors = load(loaded_calib_files[2])["prior"] + truth_sample = load(loaded_calib_files[3])["y"] + (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space + + truth_params_constrained = if force_case == "const-force" + [truth_params_obj.val] #1d + elseif force_case == "vec-force" + truth_params_obj.val + elseif force_case == "flux-force" + truth_params_model = truth_params_obj.model + truth_params_constrained, rebuild = destructure(truth_params_model) + truth_params_constrained + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = get_obs_noise_cov(ekpobj) + + n_params = length(truth_params) # "input dim" + output_dim = size(Γy, 1) + ### + ### Emulate: Gaussian Process Regression + ### + + # Emulate-sample settings + # choice of machine-learning tool in the emulation stage + if case == "GP" + gppackage = Emulators.GPJL() + pred_type = Emulators.YType() + mlt = GaussianProcess( + gppackage; + kernel = nothing, # use default squared exponential kernel + prediction_type = pred_type, + noise_learn = false, + ) + elseif case == "RF-scalar" + nugget = 1e-6 + overrides = Dict( + "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 1.0, + "n_iteration" => 10, + "n_features_opt" => 160, + ) + n_features = 200 + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) + mlt = ScalarRandomFeatureInterface( + n_features, + n_params, + rng = rng, + kernel_structure = kernel_structure, + optimizer_options = overrides, + ) + elseif case ∈ ["RF-nonsep"] + nugget = 1e-6 + overrides = Dict( + "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 1.0, + "n_iteration" => 10, + "n_features_opt" => 160, + ) + kernel_structure = NonseparableKernel(LowRankFactor(1, nugget)) + n_features = 200 + + mlt = VectorRandomFeatureInterface( + n_features, + n_params, + output_dim, + rng = rng, + kernel_structure = kernel_structure, + optimizer_options = overrides, + ) + + else + throw(ArgumentError("case $(case) not recognised, please choose from the list $(cases)")) + end + + + # Get training points from the EKP iteration number in the second input term + N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP + min_iter = min(max_iter, max(1, min_iter)) + if N_iter < 2 + @warn("Convergence in only 1 iteration, removing train-test split based on iterations") + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) + @info "Train iterations: $(min_iter:skip_iter:N_iter)" + + else + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" + @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" + end + + # Save data + @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs + + # data processing configuration + retain_var = 0.95 + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # Check how well the Gaussian Process regression predicts on the + # true parameters + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + if !(N_iter<2) + y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) + end + println("ML prediction on true parameters: ") + println(vec(y_mean)) + println("true data: ") + println(truth_sample) # what was used as truth + println(" ML predicted standard deviation") + println(sqrt.(diag(y_var[1], 0))) + println("ML MSE (truth): ") + println(mean((truth_sample - vec(y_mean)) .^ 2)) + if !(N_iter<2) + println("ML MSE (test ensemble): ") + println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) + end + + #end + ### + ### Sample: Markov Chain Monte Carlo + ### + # initial values + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + println("initial parameters: ", u0) + + # First let's run a short chain to determine a good step size + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0) + + # Now begin the actual MCMC + println("Begin MCMC - with step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + + post_mean = mean(posterior) + post_cov = cov(posterior) + println("post_mean") + println(post_mean) + println("post_cov") + println(post_cov) + println("D util") + println(det(inv(post_cov))) + println(" ") + + param_names = get_name(posterior) + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + + # marginal histogram + pp = plot(priors, c = :gray) + plot!(pp, posterior) + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "posterior_hist_" * case) + savefig(figpath * ".png") + savefig(figpath * ".pdf") + + # Save data + save( + joinpath(data_save_directory, "posterior_$(case).jld2"), + "posterior", + posterior, + "priors", + priors, + "truth_params_constrained", + truth_params_constrained, + "final_params_constrained", + final_params_constrained, + "input_output_pairs", + input_output_pairs, + "truth_params", + truth_params, + "encoder_schedule", + get_encoder_schedule(emulator), + ) +end + +main() diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl new file mode 100644 index 000000000..4c5f3e9ae --- /dev/null +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -0,0 +1,292 @@ +# some context values (from calibrate_spatial_dep.jl) + +# some context calues from emulate_sample_spatial_dep.jl +cases = ["GP", "RF-scalar", "RF-nonsep"] +case = cases[2] + +# load packages +# CES + +using Random +using JLD2 +using Statistics +using LinearAlgebra + +using Plots +using Plots.Measures + +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions +const EKP = EnsembleKalmanProcesses + + +method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] +force_cases = ["const-force", "vec-force", "flux-force"] # problem types +cases = [ + "GP", + "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs + "RF-nonsep", # diagonalize, train scalar RF, don't asume diag inputs +] + +#### CHOOSE YOUR CASE: (todo: looped over in some fashion) +# calib_data_dir +method=method_cases[1] +calibrate_date=Date("2026-05-18", "yyyy-mm-dd") +calib_directory="$(method)_$(calibrate_date)" + +# calib_filename_suffix +force_case=force_cases[3] +N_ens=60 +rng_idx=1 +calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + +# emulate_sample cases +case = cases[2] # e.g. 1:2 or [3] + +@info("case: ", case) +@info("force_case: ", force_case) + +# load +homedir = pwd() +println(homedir) +figure_save_directory = joinpath(homedir, "output/", calib_directory) +data_save_directory = joinpath(homedir, "output", calib_directory) +data_file = joinpath(data_save_directory, "posterior_$(force_case).jld2") +truth_params = load(data_file)["truth_params_constrained"] +final_params = load(data_file)["final_params_constrained"] +posterior = load(data_file)["posterior"] + +prior_file = joinpath(data_save_directory, "l96_priors_$(force_case).jld2") +prior = load(prior_file)["prior"] + +ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") +ekpobj = load(ekp_file)["ekpobj"] +N_ens = get_N_ens(ekpobj) +nx = length(truth_params) +y = get_obs(ekpobj) +ny = length(y) +# get samples and quantiles from posterior +param_names = get_name(posterior) +posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns +constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + +quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' # rows are quantiles for row of posterior samples + + +# plot - EKI results +gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) +p1 = plot( + range(0, nx - 1, step = 1), + truth_params, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (input)", + left_margin = 15mm, + bottom_margin = 15mm, +) + +p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, +# xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) +# xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), +) + +l = @layout [a b] +plt = plot(p1, p2, layout = l) + +savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) +savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + +p3 = deepcopy(p1) +p4 = deepcopy(p2) + +plot!(p1, range(0, nx - 1, step = 1), final_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) +plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + +l = @layout [a b] +plt = plot(p1, p2, layout = l) + +savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) +savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + +# +plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 1], + label = "ensemble inputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, +) + +plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 2:end], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, +) +plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 1], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, +) + +plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 2:end], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, +) + +plot!( + p4, + 1:length(y), + get_g_final(ekpobj)[:, 1], + color = :lightgreen, + label = "ensemble outputs", + linewidth = 4, + linealpha = 0.1, +) +plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) +plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) +plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + +plt = plot(p3, p4, layout = l) + +savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) +savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + +# plot - UQ results + +gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) +p1 = plot( + range(0, nx - 1, step = 1), + [truth_params final_params], + label = ["solution" "EKI-opt"], + color = [:black :lightgreen], + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (input)", + left_margin = 15mm, + bottom_margin = 15mm, +) + +plot!( + p1, + range(0, nx - 1, step = 1), + quantiles[:, 2], # median of all vals + color = :blue, + label = "posterior", + linewidth = 4, + ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], + fillalpha = 0.1, +) + +figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)") +savefig(figpath * ".png") +savefig(figpath * ".pdf") + +########################## +#cases = ["GP", "RF-scalar", "RF-nonsep"] +#= +# load +homedir = pwd() +println(homedir) +data_file_GP = joinpath(data_save_directory, "posterior_$(cases[1]).jld2") +data_file_RF = joinpath(data_save_directory, "posterior_$(case).jld2") +truth_params_GP = load(data_file_GP)["truth_params_constrained"] +final_params_GP = load(data_file_GP)["final_params_constrained"] +posterior_GP = load(data_file_GP)["posterior"] +truth_params_RF = load(data_file_RF)["truth_params_constrained"] +final_params_RF = load(data_file_RF)["final_params_constrained"] +posterior_RF = load(data_file_RF)["posterior"] + +ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") +ekpobj = load(ekp_file)["ekpobj"] +N_ens = get_N_ens(ekpobj) +nx = length(truth_params) +y = get_obs(ekpobj) +ny = length(y) +# get samples and quantiles from posterior +param_names_GP = get_name(posterior_GP) +posterior_samples_GP = vcat([get_distribution(posterior_GP)[name] for name in get_name(posterior_GP)]...) #samples are columns +constrained_posterior_samples_GP = + mapslices(x -> transform_unconstrained_to_constrained(posterior_GP, x), posterior_samples_GP, dims = 1) + +quantiles_GP = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_GP)])' # rows are quantiles for row of posterior samples + +param_names_RF = get_name(posterior_RF) +posterior_samples_RF = vcat([get_distribution(posterior_RF)[name] for name in get_name(posterior_RF)]...) #samples are columns +constrained_posterior_samples_RF = + mapslices(x -> transform_unconstrained_to_constrained(posterior_RF, x), posterior_samples_RF, dims = 1) + +quantiles_RF = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_RF)])' # rows are quantiles for row of posterior samples + +# plot - UQ results - both + +gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) +p1 = plot( + range(0, nx - 1, step = 1), + [truth_params final_params], + label = ["solution" "EKI-opt"], + color = [:black :lightgreen], + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (input)", + left_margin = 15mm, + bottom_margin = 15mm, +) + +plot!( + p1, + range(0, nx - 1, step = 1), + quantiles_GP[:, 2], # median of all vals + color = :blue, + label = "GP posterior", + linewidth = 4, + ribbon = [quantiles_GP[:, 2] - quantiles_GP[:, 1] quantiles_GP[:, 3] - quantiles_GP[:, 2]], + linealpha = 0.5, + fillalpha = 0.1, +) + +plot!( + p1, + range(0, nx - 1, step = 1), + quantiles_RF[:, 2], # median of all vals + color = :red, + label = "posterior_RF", + linewidth = 4, + ribbon = [quantiles_RF[:, 2] - quantiles_RF[:, 1] quantiles_RF[:, 3] - quantiles_RF[:, 2]], + linealpha = 0.5, + fillalpha = 0.1, +) + + +figpath = joinpath(figure_save_directory, "posterior_ribbons_$(cases[1])_$(case)") +savefig(figpath * ".png") +savefig(figpath * ".pdf") +=# From 30ec7cc167e3a1fef7d2acebb711e488b42ec367 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 20 May 2026 09:34:20 -0700 Subject: [PATCH 04/86] added the looping over the experiments --- examples/EKIRace/emulate_sample_l96.jl | 420 +++++++++++-------------- examples/EKIRace/plot_l96_forcing.jl | 360 +++++++++++---------- 2 files changed, 369 insertions(+), 411 deletions(-) diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index bddcd3f07..ebae06425 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -27,11 +27,6 @@ function main() method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] force_cases = ["const-force", "vec-force", "flux-force"] # problem types - cases = [ - "GP", - "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs - "RF-nonsep", # diagonalize, train scalar RF, don't asume diag inputs - ] #### CHOOSE YOUR CASE: (todo: looped over in some fashion) # calib_data_dir @@ -39,235 +34,200 @@ function main() calibrate_date=Date("2026-05-18", "yyyy-mm-dd") calib_directory="$(method)_$(calibrate_date)" - # calib_filename_suffix - force_case=force_cases[3] - N_ens=100 - rng_idx=1 - calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + # calib_filename_suffix items to loop over + force_cases=[force_cases[3]] + N_enss = [50, 75, 100] # [5,15,30] #[50,75,100] + rng_idxs=[1,2,3,4] # emulate_sample cases - case = cases[2] # e.g. 1:2 or [3] - - @info("case: ", case) - @info("force_case: ", force_case) - min_iter = 1 - skip_iter = 1 - max_iter = 10 # number of EKP iterations to use data from is at most this - - #### - - # need to loop over - # method, rng_index, N_ens = "TransformInversion", "1", "100"] - exp_name = "Lorenz_histogram_l96" - rng_seed = 44011 - rng = Random.MersenneTwister(rng_seed) - - # loading relevant data - homedir = pwd() - println(homedir) - figure_save_directory = joinpath(homedir, "output/", calib_directory) - data_save_directory = joinpath(homedir, "output", calib_directory) - loaded_calib_files = [ - joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2"), - joinpath(data_save_directory, "l96_priors_$(force_case).jld2"), - joinpath(data_save_directory, "l96_calibrate_results_$(calib_filename_suffix).jld2"), - ] - if !isfile(loaded_calib_files[1]) - throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) - end - - ekpobj = load(loaded_calib_files[1])["ekpobj"] - priors = load(loaded_calib_files[2])["prior"] - truth_sample = load(loaded_calib_files[3])["y"] - (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space - - truth_params_constrained = if force_case == "const-force" - [truth_params_obj.val] #1d - elseif force_case == "vec-force" - truth_params_obj.val - elseif force_case == "flux-force" - truth_params_model = truth_params_obj.model - truth_params_constrained, rebuild = destructure(truth_params_model) - truth_params_constrained - end - truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) - Γy = get_obs_noise_cov(ekpobj) - - n_params = length(truth_params) # "input dim" - output_dim = size(Γy, 1) - ### - ### Emulate: Gaussian Process Regression - ### - - # Emulate-sample settings - # choice of machine-learning tool in the emulation stage - if case == "GP" - gppackage = Emulators.GPJL() - pred_type = Emulators.YType() - mlt = GaussianProcess( - gppackage; - kernel = nothing, # use default squared exponential kernel - prediction_type = pred_type, - noise_learn = false, - ) - elseif case == "RF-scalar" - nugget = 1e-6 - overrides = Dict( - "verbose" => true, - "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 1.0, - "n_iteration" => 10, - "n_features_opt" => 160, - ) - n_features = 200 - kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) - mlt = ScalarRandomFeatureInterface( - n_features, - n_params, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - elseif case ∈ ["RF-nonsep"] - nugget = 1e-6 - overrides = Dict( - "verbose" => true, - "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 1.0, - "n_iteration" => 10, - "n_features_opt" => 160, - ) - kernel_structure = NonseparableKernel(LowRankFactor(1, nugget)) - n_features = 200 - - mlt = VectorRandomFeatureInterface( - n_features, - n_params, - output_dim, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - - else - throw(ArgumentError("case $(case) not recognised, please choose from the list $(cases)")) + for force_case in force_cases + for N_ens in N_enss + for rng_idx in rng_idxs + calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + @info("Perform emulate-sample for L96: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") + min_iter = 1 + skip_iter = 1 + max_iter = 15 # number of EKP iterations to use data from is at most this + + #### + + # need to loop over + rng_seed = 44011 + rng = Random.MersenneTwister(rng_seed) + + # loading relevant data + homedir = pwd() + println(homedir) + figure_save_directory = joinpath(homedir, "output/", calib_directory) + data_save_directory = joinpath(homedir, "output", calib_directory) + loaded_calib_files = [ + joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, "l96_priors_$(force_case).jld2"), + joinpath(data_save_directory, "l96_calibrate_results_$(calib_filename_suffix).jld2"), + ] + if !isfile(loaded_calib_files[1]) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + + ekpobj = load(loaded_calib_files[1])["ekpobj"] + priors = load(loaded_calib_files[2])["prior"] + truth_sample = load(loaded_calib_files[3])["y"] + (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space + + truth_params_constrained = if force_case == "const-force" + [truth_params_obj.val] #1d + elseif force_case == "vec-force" + truth_params_obj.val + elseif force_case == "flux-force" + truth_params_model = truth_params_obj.model + truth_params_constrained, rebuild = destructure(truth_params_model) + truth_params_constrained + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = get_obs_noise_cov(ekpobj) + + n_params = length(truth_params) # "input dim" + output_dim = size(Γy, 1) + ### + ### Emulate: + ### + # choice of machine-learning tool in the emulation stage + nugget = 1e-6 + overrides = Dict( +# "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 20.0, + "n_iteration" => 10, + "n_features_opt" => 160, + ) + n_features = 200 + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) + mlt = ScalarRandomFeatureInterface( + n_features, + n_params, + rng = rng, + kernel_structure = kernel_structure, + optimizer_options = overrides, + ) + + + # Get training points from the EKP iteration number in the second input term + N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP + min_iter = min(max_iter, max(1, min_iter)) + if N_iter < 2 + @warn("Convergence in only 1 iteration, removing train-test split based on iterations") + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) + @info "Train iterations: $(min_iter:skip_iter:N_iter)" + + else + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" + @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" + end + + # Save data + @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs + + # data processing configuration + retain_var = 0.95 + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # Check how well the Gaussian Process regression predicts on the + # true parameters + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + if !(N_iter<2) + y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) + end + println("ML prediction on true parameters: ") + println(vec(y_mean)) + println("true data: ") + println(truth_sample) # what was used as truth + println(" ML predicted standard deviation") + println(sqrt.(diag(y_var[1], 0))) + println("ML MSE (truth): ") + println(mean((truth_sample - vec(y_mean)) .^ 2)) + if !(N_iter<2) + println("ML MSE (test ensemble): ") + println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) + end + + #end + ### + ### Sample: Markov Chain Monte Carlo + ### + # initial values + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + println("initial parameters: ", u0) + + # First let's run a short chain to determine a good step size + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + + # Now begin the actual MCMC + println("Begin MCMC - with step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + + post_mean = mean(posterior) + post_cov = cov(posterior) + println("post_mean") + println(post_mean) + println("post_cov") + println(post_cov) + println("D util") + println(det(inv(post_cov))) + println(" ") + + param_names = get_name(posterior) + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + + # marginal histogram + pp = plot(priors, c = :gray) + plot!(pp, posterior) + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + + # Save data + save( + joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2"), + "posterior", + posterior, + "priors", + priors, + "truth_params_constrained", + truth_params_constrained, + "final_params_constrained", + final_params_constrained, + "input_output_pairs", + input_output_pairs, + "truth_params", + truth_params, + "encoder_schedule", + get_encoder_schedule(emulator), + ) + end end - - - # Get training points from the EKP iteration number in the second input term - N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP - min_iter = min(max_iter, max(1, min_iter)) - if N_iter < 2 - @warn("Convergence in only 1 iteration, removing train-test split based on iterations") - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) - @info "Train iterations: $(min_iter:skip_iter:N_iter)" - - else - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations - @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" - @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" end - - # Save data - @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs - - # data processing configuration - retain_var = 0.95 - encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] - encoder_kwargs = encoder_kwargs_from(ekpobj, priors) - - emulator = Emulator( - mlt, - input_output_pairs; - encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, - ) - optimize_hyperparameters!(emulator) - - # Check how well the Gaussian Process regression predicts on the - # true parameters - y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) - if !(N_iter<2) - y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) - end - println("ML prediction on true parameters: ") - println(vec(y_mean)) - println("true data: ") - println(truth_sample) # what was used as truth - println(" ML predicted standard deviation") - println(sqrt.(diag(y_var[1], 0))) - println("ML MSE (truth): ") - println(mean((truth_sample - vec(y_mean)) .^ 2)) - if !(N_iter<2) - println("ML MSE (test ensemble): ") - println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) - end - - #end - ### - ### Sample: Markov Chain Monte Carlo - ### - # initial values - u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) - println("initial parameters: ", u0) - - # First let's run a short chain to determine a good step size - mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) - new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0) - - # Now begin the actual MCMC - println("Begin MCMC - with step size ", new_step) - chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) - - posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) - - post_mean = mean(posterior) - post_cov = cov(posterior) - println("post_mean") - println(post_mean) - println("post_cov") - println(post_cov) - println("D util") - println(det(inv(post_cov))) - println(" ") - - param_names = get_name(posterior) - - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - - # marginal histogram - pp = plot(priors, c = :gray) - plot!(pp, posterior) - final_params_constrained = get_ϕ_mean_final(priors, ekpobj) - for (i, sp) in enumerate(pp.subplots) - vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) - vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) - end - figpath = joinpath(figure_save_directory, "posterior_hist_" * case) - savefig(figpath * ".png") - savefig(figpath * ".pdf") - - # Save data - save( - joinpath(data_save_directory, "posterior_$(case).jld2"), - "posterior", - posterior, - "priors", - priors, - "truth_params_constrained", - truth_params_constrained, - "final_params_constrained", - final_params_constrained, - "input_output_pairs", - input_output_pairs, - "truth_params", - truth_params, - "encoder_schedule", - get_encoder_schedule(emulator), - ) end main() diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl index 4c5f3e9ae..f8784ea6b 100644 --- a/examples/EKIRace/plot_l96_forcing.jl +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -1,9 +1,3 @@ -# some context values (from calibrate_spatial_dep.jl) - -# some context calues from emulate_sample_spatial_dep.jl -cases = ["GP", "RF-scalar", "RF-nonsep"] -case = cases[2] - # load packages # CES @@ -11,6 +5,7 @@ using Random using JLD2 using Statistics using LinearAlgebra +using Dates using Plots using Plots.Measures @@ -34,182 +29,185 @@ method=method_cases[1] calibrate_date=Date("2026-05-18", "yyyy-mm-dd") calib_directory="$(method)_$(calibrate_date)" -# calib_filename_suffix -force_case=force_cases[3] -N_ens=60 -rng_idx=1 -calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - -# emulate_sample cases -case = cases[2] # e.g. 1:2 or [3] - -@info("case: ", case) -@info("force_case: ", force_case) - -# load -homedir = pwd() -println(homedir) -figure_save_directory = joinpath(homedir, "output/", calib_directory) -data_save_directory = joinpath(homedir, "output", calib_directory) -data_file = joinpath(data_save_directory, "posterior_$(force_case).jld2") -truth_params = load(data_file)["truth_params_constrained"] -final_params = load(data_file)["final_params_constrained"] -posterior = load(data_file)["posterior"] - -prior_file = joinpath(data_save_directory, "l96_priors_$(force_case).jld2") -prior = load(prior_file)["prior"] - -ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") -ekpobj = load(ekp_file)["ekpobj"] -N_ens = get_N_ens(ekpobj) -nx = length(truth_params) -y = get_obs(ekpobj) -ny = length(y) -# get samples and quantiles from posterior -param_names = get_name(posterior) -posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns -constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - -quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' # rows are quantiles for row of posterior samples - - -# plot - EKI results -gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) -p1 = plot( - range(0, nx - 1, step = 1), - truth_params, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing (input)", - left_margin = 15mm, - bottom_margin = 15mm, -) - -p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, -# xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) -# xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), -) - -l = @layout [a b] -plt = plot(p1, p2, layout = l) - -savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) -savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) - -p3 = deepcopy(p1) -p4 = deepcopy(p2) - -plot!(p1, range(0, nx - 1, step = 1), final_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) -plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) - -l = @layout [a b] -plt = plot(p1, p2, layout = l) - -savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) -savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) - -# -plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1], - label = "ensemble inputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, -) - -plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end], - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, -) -plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1], - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, -) - -plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end], - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, -) - -plot!( - p4, - 1:length(y), - get_g_final(ekpobj)[:, 1], - color = :lightgreen, - label = "ensemble outputs", - linewidth = 4, - linealpha = 0.1, -) -plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) -plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) -plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - - -plt = plot(p3, p4, layout = l) - -savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) -savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - -# plot - UQ results - -gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) -p1 = plot( - range(0, nx - 1, step = 1), - [truth_params final_params], - label = ["solution" "EKI-opt"], +# calib_filename_suffix items to loop over +force_cases=[force_cases[3]] +N_enss=[50,75,100] +rng_idxs=[1,2,3,4] + + +for force_case in force_cases + for N_ens in N_enss + for rng_idx in rng_idxs + calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + + @info("Plotting for L96: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") + + # load + homedir = pwd() + println(homedir) + figure_save_directory = joinpath(homedir, "output/", calib_directory) + data_save_directory = joinpath(homedir, "output", calib_directory) + data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + truth_params = load(data_file)["truth_params_constrained"] + final_params = load(data_file)["final_params_constrained"] + posterior = load(data_file)["posterior"] + + prior_file = joinpath(data_save_directory, "l96_priors_$(force_case).jld2") + prior = load(prior_file)["prior"] + + ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") + ekpobj = load(ekp_file)["ekpobj"] + N_ens = get_N_ens(ekpobj) + nx = length(truth_params) + y = get_obs(ekpobj) + ny = length(y) + # get samples and quantiles from posterior + param_names = get_name(posterior) + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + + quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' # rows are quantiles for row of posterior samples + + + # plot - EKI results + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + truth_params, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + # xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) + # xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), + ) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + p3 = deepcopy(p1) + p4 = deepcopy(p2) + + plot!(p1, range(0, nx - 1, step = 1), final_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + + # + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 1], + label = "ensemble inputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 2:end], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 1], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 2:end], + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p4, + 1:length(y), + get_g_final(ekpobj)[:, 1], + color = :lightgreen, + label = "ensemble outputs", + linewidth = 4, + linealpha = 0.1, + ) + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + + plt = plot(p3, p4, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + + # plot - UQ results + + gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + [truth_params final_params], + label = ["solution" "EKI-opt"], color = [:black :lightgreen], - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing (input)", - left_margin = 15mm, - bottom_margin = 15mm, -) - -plot!( - p1, - range(0, nx - 1, step = 1), - quantiles[:, 2], # median of all vals - color = :blue, - label = "posterior", - linewidth = 4, - ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], - fillalpha = 0.1, -) - -figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)") -savefig(figpath * ".png") -savefig(figpath * ".pdf") - + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plot!( + p1, + range(0, nx - 1, step = 1), + quantiles[:, 2], # median of all vals + color = :blue, + label = "posterior", + linewidth = 4, + ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], + fillalpha = 0.1, + ) + + figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + end + end +end ########################## #cases = ["GP", "RF-scalar", "RF-nonsep"] #= From c89b1200c7c50d2defafb78f0a6b9c3ba5678cd1 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 22 May 2026 14:52:10 -0700 Subject: [PATCH 05/86] update error catch --- examples/EKIRace/emulate_sample_l96.jl | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index ebae06425..82884a3da 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -36,9 +36,9 @@ function main() # calib_filename_suffix items to loop over force_cases=[force_cases[3]] - N_enss = [50, 75, 100] # [5,15,30] #[50,75,100] - rng_idxs=[1,2,3,4] - + N_enss = [100] # [5,15,30] #[50,75,100] + rng_idxs=[2,3,4] + # emulate_sample cases for force_case in force_cases for N_ens in N_enss @@ -112,18 +112,21 @@ function main() # Get training points from the EKP iteration number in the second input term - N_iter = min(max_iter, length(get_u(ekpobj)) - 1) # number of paired iterations taken from EKP + N_iter = min(max_iter, length(get_u(ekpobj))-1) # number of paired iterations taken from EKP min_iter = min(max_iter, max(1, min_iter)) - if N_iter < 2 + if N_iter < 1 + @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") + continue + elseif N_iter == 1 @warn("Convergence in only 1 iteration, removing train-test split based on iterations") - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) @info "Train iterations: $(min_iter:skip_iter:N_iter)" - + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter+1) else - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + end # Save data From adc6c5fd5680536d6412000b252c75ecef4f3ccf Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 22 May 2026 14:52:49 -0700 Subject: [PATCH 06/86] plot tool can now plot an exp series --- examples/EKIRace/plot_l96_forcing.jl | 81 ++++++++++++++++++++-------- 1 file changed, 59 insertions(+), 22 deletions(-) diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl index f8784ea6b..b8b5aad38 100644 --- a/examples/EKIRace/plot_l96_forcing.jl +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -17,11 +17,6 @@ const EKP = EnsembleKalmanProcesses method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] force_cases = ["const-force", "vec-force", "flux-force"] # problem types -cases = [ - "GP", - "RF-scalar", # diagonalize, train scalar RF, don't asume diag inputs - "RF-nonsep", # diagonalize, train scalar RF, don't asume diag inputs -] #### CHOOSE YOUR CASE: (todo: looped over in some fashion) # calib_data_dir @@ -34,19 +29,35 @@ force_cases=[force_cases[3]] N_enss=[50,75,100] rng_idxs=[1,2,3,4] +homedir = pwd() +valid_file_items = [] +valid_files = [] for force_case in force_cases for N_ens in N_enss for rng_idx in rng_idxs + data_save_directory = joinpath(homedir, "output", calib_directory) calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + if isfile(data_file) + push!(valid_files, calib_filename_suffix) + push!(valid_file_items, (force_case, N_ens, rng_idx)) + end + end + end +end +@info "Plotting valid files:" +display(valid_files) +println(" ") + +for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items,valid_files) + calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" @info("Plotting for L96: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") # load - homedir = pwd() - println(homedir) - figure_save_directory = joinpath(homedir, "output/", calib_directory) data_save_directory = joinpath(homedir, "output", calib_directory) + figure_save_directory = data_save_directory data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") truth_params = load(data_file)["truth_params_constrained"] final_params = load(data_file)["final_params_constrained"] @@ -79,7 +90,7 @@ for force_case in force_cases color = :black, linewidth = 4, xlabel = "Spatial index", - ylabel = "Forcing (input)", + ylabel = "Forcing difference (input)", left_margin = 15mm, bottom_margin = 15mm, ) @@ -104,11 +115,39 @@ for force_case in force_cases savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + zeros(nx), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) - p3 = deepcopy(p1) - p4 = deepcopy(p2) - - plot!(p1, range(0, nx - 1, step = 1), final_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + # xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) + # xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), + ) + + p3 = deepcopy(p1) + p4 = deepcopy(p2) + + plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) l = @layout [a b] @@ -121,7 +160,7 @@ for force_case in force_cases plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1], + get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, label = "ensemble inputs", color = :lightgreen, linewidth = 4, @@ -131,7 +170,7 @@ for force_case in force_cases plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end], + get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, label = "", color = :lightgreen, linewidth = 4, @@ -140,7 +179,7 @@ for force_case in force_cases plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1], + get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, label = "", color = :lightgreen, linewidth = 4, @@ -150,7 +189,7 @@ for force_case in force_cases plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end], + get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, label = "", color = :lightgreen, linewidth = 4, @@ -181,12 +220,12 @@ for force_case in force_cases gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), - [truth_params final_params], + [truth_params final_params] .- truth_params, label = ["solution" "EKI-opt"], color = [:black :lightgreen], linewidth = 4, xlabel = "Spatial index", - ylabel = "Forcing (input)", + ylabel = "Forcing difference (input)", left_margin = 15mm, bottom_margin = 15mm, ) @@ -194,7 +233,7 @@ for force_case in force_cases plot!( p1, range(0, nx - 1, step = 1), - quantiles[:, 2], # median of all vals + quantiles[:, 2] .- truth_params, # median of all vals color = :blue, label = "posterior", linewidth = 4, @@ -206,8 +245,6 @@ for force_case in force_cases savefig(figpath * ".png") savefig(figpath * ".pdf") end - end -end ########################## #cases = ["GP", "RF-scalar", "RF-nonsep"] #= From c74367aac0ac40e2ed9f482e8f208d934b7cf51b Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 27 May 2026 17:37:18 -0700 Subject: [PATCH 07/86] added conversion from file experiment to nc dataset for the experiment --- .../EKIRace/exp_to_leaderboard_utilities.jl | 217 ++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 examples/EKIRace/exp_to_leaderboard_utilities.jl diff --git a/examples/EKIRace/exp_to_leaderboard_utilities.jl b/examples/EKIRace/exp_to_leaderboard_utilities.jl new file mode 100644 index 000000000..6528a64f0 --- /dev/null +++ b/examples/EKIRace/exp_to_leaderboard_utilities.jl @@ -0,0 +1,217 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +# Step 1, Read and extract the available experiments +method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] +force_cases = ["const-force", "vec-force", "flux-force"] # problem types + +#### CHOOSE YOUR CASE: +# calib_data_dir +method = method_cases[1] +calibrate_date = Date("2026-05-18", "yyyy-mm-dd") +calib_directory = "$(method)_$(calibrate_date)" + +# calib_filename_suffix items to loop over +force_case = force_cases[2] +N_enss = [50, 75, 100] +rng_idxs = [1, 2, 3, 4] + +### For saving the cases in the leaderboard style +forcing_cases_key = Dict( + "const-force" => "ensemble_results", + "vec-force" => "spatial_forcing_ensemble_results", + "flux-force" => "nn_forcing_ensemble_results", +) + +method_cases_key = Dict( + "Inversion" => "ces-eki-dmc", + "TransformInversion" => "ces-etki-dmc", + "Unscented" => "ces-uki-dmc", + "GaussNewtonInversion" => "ces-iekf", +) +nc_save_filename = "$(method_cases_key[method])_l96_$(forcing_cases_key[force_case])_$(calibrate_date).nc" + +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_directory) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" + data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + if isfile(data_file) + push!(valid_files, calib_filename_suffix) + push!(valid_file_items, (force_case, N_ens, rng_idx)) + end + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") +end + +### Load data + +# Load first valid file to determine parameter dimension +first_loaded = JLD2.load(joinpath(data_save_directory, "posterior_$(valid_files[1]).jld2")) +n_params = length(vec(mean(first_loaded["posterior"]))) + +targets = [1.0] + +n_rng = length(rng_idxs) +n_ens = length(N_enss) +n_targets = length(targets) + +# Pre-allocate storage arrays indexed as (random_seed, ensemble_size, rmse_target, algorithm_type, ...) +true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +logpdf_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) + +for (fc, N_ens, rng_idx) in valid_file_items + calib_filename_suffix = "$(fc)_$(N_ens)_$(rng_idx)" + post_filename = "posterior_$(calib_filename_suffix).jld2" + @info "loading case $(post_filename)" + loaded = JLD2.load(joinpath(data_save_directory, post_filename)) + + post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) + truth_params = loaded["truth_params"] # truth in unconstrained space + post_name = get_name(post_dist) + + pm = vec(mean(post_dist)) # posterior mean, shape (n_params,) + pc = cov(post_dist) # posterior covariance, shape (n_params, n_params) + + C_reg = Symmetric(pc + 1e-6 * I(n_params)) + post_normal = MvNormal(pm, C_reg) + post_array = get_distribution(post_dist)[get_name(post_dist)[1]] + pmode_idx = argmax(logpdf(post_normal, post_array)) + pmode = post_array[:, pmode_idx] + # Load ekpobj to compute calibration evaluation count + ekp_loaded = JLD2.load(joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2")) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps + + # Array indices + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + diff = pm - truth_params + + for l in 1:n_targets + true_param_arr[i, j, l, 1, :] = truth_params + post_mean_arr[i, j, l, 1, :] = pm + post_cov_arr[i, j, l, 1, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, l, 1] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + # total forward model evaluations to reach target + n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens + end + +end + + +### Saving the data in nc + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "rmse_target", n_targets) +defDim(ds, "algorithm_type", 1) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) # second param axis for covariance matrix + +# Coordinate variables +alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) +alg_var[1] = method_cases_key[method] + +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) +rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" +rmse_var[:] = targets + +# Data variables + +n_evals_v = defVar( + ds, + "n_evals_to_target", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +n_evals_v.attrib["description"] = "Number of forward model evaluations (i.e., algorithm cost): conv_alg_iters * ensemble_size. Stored as float in case of NaN values." +n_evals_v[:, :, :, :] = n_evals_arr + +true_param_v = defVar( + ds, + "true_param", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + fillvalue = NaN, +) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :, :, :] = true_param_arr + +post_mean_v = defVar( + ds, + "post_mean", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + fillvalue = NaN, +) +post_mean_v.attrib["description"] = "mean(posterior)" +post_mean_v[:, :, :, :, :] = post_mean_arr + +post_cov_v = defVar( + ds, + "post_cov", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim", "param_dim_2"); + fillvalue = NaN, +) +post_cov_v.attrib["description"] = "cov(posterior)" +post_cov_v[:, :, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar( + ds, + "mahalanobis", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" +mahalanobis_v[:, :, :, :] = mahal_arr + +posterior_logpdf_v = defVar( + ds, + "posterior_logpdf", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +posterior_logpdf_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" +posterior_logpdf_v[:, :, :, :] = logpdf_arr + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" From 704c2568977bc3cc7806183a7986ab42c3c87710 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 27 May 2026 17:39:54 -0700 Subject: [PATCH 08/86] better posterior naming --- examples/EKIRace/exp_to_leaderboard_utilities.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/EKIRace/exp_to_leaderboard_utilities.jl b/examples/EKIRace/exp_to_leaderboard_utilities.jl index 6528a64f0..e859b93b3 100644 --- a/examples/EKIRace/exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/exp_to_leaderboard_utilities.jl @@ -79,7 +79,7 @@ true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -logpdf_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) for (fc, N_ens, rng_idx) in valid_file_items @@ -117,7 +117,7 @@ for (fc, N_ens, rng_idx) in valid_file_items # (m - truth)' C^{-1} (m - truth) mahal_arr[i, j, l, 1] = diff' * (C_reg \ diff) # log p(truth | posterior) - log p(mode | posterior) - logpdf_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + logpdf_true_v_map_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) # total forward model evaluations to reach target n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens end @@ -203,15 +203,15 @@ mahalanobis_v = defVar( mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" mahalanobis_v[:, :, :, :] = mahal_arr -posterior_logpdf_v = defVar( +posterior_logpdf_true_v_map_v = defVar( ds, - "posterior_logpdf", + "posterior_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); fillvalue = NaN, ) -posterior_logpdf_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" -posterior_logpdf_v[:, :, :, :] = logpdf_arr +posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" +posterior_logpdf_true_v_map_v[:, :, :, :] = logpdf_true_v_map_arr close(ds) @info "Saved leaderboard data to $(nc_save_filename)" From c23eb1649f34cc0d4458d55cd77c8f4c3e43aa82 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 27 May 2026 18:06:14 -0700 Subject: [PATCH 09/86] better naming --- examples/EKIRace/exp_to_leaderboard_utilities.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/EKIRace/exp_to_leaderboard_utilities.jl b/examples/EKIRace/exp_to_leaderboard_utilities.jl index e859b93b3..be717b479 100644 --- a/examples/EKIRace/exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/exp_to_leaderboard_utilities.jl @@ -19,7 +19,7 @@ calib_directory = "$(method)_$(calibrate_date)" # calib_filename_suffix items to loop over force_case = force_cases[2] -N_enss = [50, 75, 100] +N_enss = [50,75,100] #[5, 15, 30] rng_idxs = [1, 2, 3, 4] ### For saving the cases in the leaderboard style @@ -95,11 +95,13 @@ for (fc, N_ens, rng_idx) in valid_file_items pm = vec(mean(post_dist)) # posterior mean, shape (n_params,) pc = cov(post_dist) # posterior covariance, shape (n_params, n_params) - C_reg = Symmetric(pc + 1e-6 * I(n_params)) + C_reg = Symmetric(pc + 1e-10 * I) post_normal = MvNormal(pm, C_reg) - post_array = get_distribution(post_dist)[get_name(post_dist)[1]] + post_array = get_distribution(post_dist)[get_name(post_dist)[1]] # safe as posterior never has constraints - just samples pmode_idx = argmax(logpdf(post_normal, post_array)) pmode = post_array[:, pmode_idx] + + # Load ekpobj to compute calibration evaluation count ekp_loaded = JLD2.load(joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2")) conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps From 8f394dae2ca2f462a567e14a2467fa52747140ac Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 28 May 2026 15:35:24 -0700 Subject: [PATCH 10/86] adding l63 experiment, tighten up l96 cases --- examples/EKIRace/Project.toml | 1 + examples/EKIRace/calibrate_l63.jl | 268 +++++++++++++++++++++++++ examples/EKIRace/calibrate_l96.jl | 18 +- examples/EKIRace/emulate_sample_l63.jl | 224 +++++++++++++++++++++ examples/EKIRace/plot_l96_forcing.jl | 6 +- 5 files changed, 506 insertions(+), 11 deletions(-) create mode 100644 examples/EKIRace/calibrate_l63.jl create mode 100644 examples/EKIRace/emulate_sample_l63.jl diff --git a/examples/EKIRace/Project.toml b/examples/EKIRace/Project.toml index 889f334d6..39c1c6f84 100644 --- a/examples/EKIRace/Project.toml +++ b/examples/EKIRace/Project.toml @@ -7,6 +7,7 @@ EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl new file mode 100644 index 000000000..61be6b6c8 --- /dev/null +++ b/examples/EKIRace/calibrate_l63.jl @@ -0,0 +1,268 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Dates + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +include("Lorenz63.jl") # Contains Lorenz 96 source code + +verbose_flag = false +save_all_ekp = true +######################################################################## +############### Choose problem type and structure ###################### +######################################################################## + +N_ens_sizes = [20, 25, 30] # list of number of ensemble members (should be problem dependent) +N_iter = 20 # maximum number of EKI iterations allowed +target_rmse = 1.0 # target RMSE +n_repeats = 4 +rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds +@info "Running Lorenz 63 problem" +@info "Maximum number of EKI iterations: $N_iter" +@info "RMSE target: $target_rmse" +configuration = + Dict("N_iter" => N_iter, "N_ens_sizes" => N_ens_sizes, "target_rmse" => target_rmse, "rng_seeds" => rng_seeds) + +nx = 3 # dimensions of parameter vector +nu = 2 +truth_params = EnsembleMemberConfig([28.0, 8.0 / 3.0]) + +prior_mean = [3.3, 1.2] +prior_cov = [ + 0.15^2 0 + 0 0.5^2 +] +#Creating prior distribution +distribution = Parameterized(MvNormal(prior_mean, prior_cov)) +constraint = repeat([no_constraint()], 2) +name = "l63_prior" +prior = ParameterDistribution(distribution, constraint, name) +T = 40.0 + + + +######################################################################## +############################ Problem setup ############################# +######################################################################## +rng_seed_init = 11 +rng_i = MersenneTwister(rng_seed_init) + +output_dir = joinpath(@__DIR__, "output") +if !isdir(output_dir) + mkdir(output_dir) +end + +#Creating my sythetic data + +t = 0.01 #time step +T_long = 1000.0 #total time +picking_initial_condition = LorenzConfig(t, T_long) +x_initial = rand(rng_i, Normal(0.0, 1.0), nx) # initial condition for spinning up Lorenz system +x_spun_up = lorenz_solve(truth_params, x_initial, picking_initial_condition) # spinning up Lorenz system + +x0 = x_spun_up[:, end] #last element of the run is the initial condition for creating the data + +ny = 9 #number of data points +lorenz_config_settings = LorenzConfig(t, T) + +# construct how we compute Observations +T_start = 30.0 +T_end = T +observation_config = ObservationConfig(T_start, T_end) +y = lorenz_forward(truth_params, x0, lorenz_config_settings, observation_config) # synthetic data + +#Observation covariance R +multiple = 36 +window = T_end - T_start +T_R = multiple * window + T_start +R_config = LorenzConfig(t, T_R) +R_run = lorenz_solve(truth_params, x_initial, R_config) +R_sample_size = Int(ceil(multiple)) +R_samples = zeros(ny, R_sample_size) +for ii in 1:R_sample_size + local_obs_config = ObservationConfig(T_start + (ii - 1) * window, T_start + ii * window) + R_samples[:, ii] = stats(R_run, R_config, local_obs_config) +end +R = cov(R_samples, dims = 2) +R_sqrt = sqrt(R) +R_inv_var = sqrt(inv(R)) + + +# Need a way to perturb the initial condition when doing the EKI updates +# Solving for initial condition perturbation covariance +covT = 2000.0 #time to simulate to calculate a covariance matrix of the system +cov_solve = lorenz_solve(truth_params, x0, LorenzConfig(t, covT)) +ic_cov = 0.1 * cov(cov_solve, dims = 2) +ic_cov_sqrt = sqrt(ic_cov) + +######################################################################## +########################### Running EKI Race ########################### +######################################################################## + +# Counters +conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how many iterations it takes to converge (per algorithm, per rand seed, per ense size) +final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) +final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) + +method_names = [ + ("Inversion(prior)", "TEKI"), + ("TransformInversion(prior)", "ETKI"), + ("GaussNewtonInversion(prior)", "GNKI"), + ("Unscented(prior; impose_prior=true)", "UKI"), +] + + +for (rr, rng_seed) in enumerate(rng_seeds) + @info "Random seed: $(rng_seed)" + rng = MersenneTwister(rng_seed) + + for (ee, N_ens) in enumerate(N_ens_sizes) + # initial parameters: N_params x N_ens + initial_params = construct_initial_ensemble(rng, prior, N_ens) + methods = [ + Inversion(prior), + TransformInversion(prior), + GaussNewtonInversion(prior), + Unscented(prior; impose_prior = true), + ] + + @info "Ensemble size: $(N_ens)" + for (kk, method) in enumerate(methods) + if isa(method, Unscented) + ekpobj = EKP.EnsembleKalmanProcess( + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + accelerator = DefaultAccelerator(), + localization_method = NoLocalization(), + scheduler = DefaultScheduler(), + ) + else + ekpobj = EKP.EnsembleKalmanProcess( + initial_params, + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + accelerator = DefaultAccelerator(), + localization_method = NoLocalization(), + scheduler = DefaultScheduler(), + ) + end + Ne = get_N_ens(ekpobj) + + count = 0 + for i in 1:N_iter + params_i = get_ϕ_final(prior, ekpobj) + + # Calculating RMSE_e + ens_mean = mean(params_i, dims = 2)[:] + G_ens_mean = lorenz_forward( + EnsembleMemberConfig(exp.(ens_mean)), + x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + RMSE_e = norm(R_inv_var * (y - G_ens_mean[:])) / sqrt(size(y, 1)) + @info "RMSE (at G(u_mean)): $(RMSE_e)" + # Convergence criteria + if RMSE_e < target_rmse + conv_alg_iters[kk, ee, rr] = count * Ne + final_parameters[kk, ee, rr, :] = ens_mean + final_model_output[kk, ee, rr, :] = G_ens_mean + break + end + + # If RMSE convergence criteria is not satisfied + G_ens = hcat( + [ + lorenz_forward( + EnsembleMemberConfig(exp.(params_i[:, j])), + (x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, Ne))[:, j], + lorenz_config_settings, + observation_config, + ) for j in 1:Ne + ]..., + ) + # Update + EKP.update_ensemble!(ekpobj, G_ens) + count = count + 1 + + # # Calculate RMSE_f to the mean G(u) if desired + # RMSE_f = sqrt(get_error_metrics(ekpobj)["loss"][end]) # equivalently + # @info "RMSE (at mean(G(u)): $(RMSE_f)" + # # Convergence criteria + # if RMSE_f < target_rmse + # conv_alg_iters[kk, ee, rr] = count * Ne + # final_parameters[kk, ee, rr, :] = ens_mean + # final_model_output[kk, ee, rr, :] = G_ens_mean + # break + # end + end + + final_ensemble = get_ϕ_final(prior, ekpobj) + + # save ekp files + per_method_dir = joinpath(output_dir,"$(nameof(typeof(method)))_$(today())") + if rr == 1 && ee == 1 + if !isdir(per_method_dir) + mkpath(per_method_dir) + end + prior_filename = joinpath(per_method_dir,"l63_priors.jld2") + JLD2.save(prior_filename,"prior", prior) + end + if save_all_ekp + # JLD2 + ekp_filename = joinpath(per_method_dir, "l63_ekp_$(N_ens)_$(rr).jld2") + JLD2.save( + ekp_filename, + "N_ens", N_ens, + "method", method, + "ekpobj", ekpobj, + ) + results_filename = joinpath(per_method_dir, "l63_calibrate_results_$(N_ens)_$(rr).jld2") + u_stored = get_u(ekpobj, return_array = false) + g_stored = get_g(ekpobj, return_array = false) + + JLD2.save( + results_filename, + "y", y, + "R", R, + "inputs", u_stored, + "outputs", g_stored, + "truth_params_structure", truth_params, # EnsembleMemberConfig + ) + end + + end + end +end + +# Saving data: +data_filename = joinpath(output_dir, "l63_output_$(today()).jld2") +JLD2.save( + data_filename, + "configuration", + configuration, + "method_names", + method_names, + "conv_alg_iters", + conv_alg_iters, + "final_parameters", + final_parameters, + "final_model_output", + final_model_output, +) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index efdf9f47d..ef92a7e1e 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -27,14 +27,16 @@ save_all_ekp = true cases = ["const-force", "vec-force", "flux-force"] # problem types # User specifications -#case = cases[1] # choose problem type -#N_ens_sizes = [5, 15, 30] # e.g. for case 1 - -#case = cases[2] # choose problem type -#N_ens_sizes = [50, 75, 100] - -case = cases[3] # choose problem type -N_ens_sizes = [50, 75, 100] # e.g. for case 1 +case = cases[1] # choose problem type +if case == "const-force" + N_ens_sizes = [5, 15, 30] +elseif case == "vec-force" + N_ens_sizes = [50, 75, 100] +elseif case == "flux-force" + N_ens_sizes = [50, 75, 100] +else + throw(ArgumentError("Expected case ∈ $(cases). Got case = $case.")) +end N_iter = 20 # maximum number of EKI iterations allowed diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl new file mode 100644 index 000000000..063634979 --- /dev/null +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -0,0 +1,224 @@ +# Import modules +include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) + +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +ENV["GKSwstype"] = "100" +using StatsPlots +using Plots +using Random +using JLD2 +using Dates + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers + +include("Lorenz63.jl") # Contains Lorenz 96 source code + + +function main() + + method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] + + #### CHOOSE YOUR CASE: (todo: looped over in some fashion) + # calib_data_dir + method=method_cases[1] + calibrate_date=Date("2026-05-28", "yyyy-mm-dd") + calib_directory="$(method)_$(calibrate_date)" + + # calib_filename_suffix items to loop over + N_enss = [20,25,30] + rng_idxs=[1,2,3,4] + + # emulate_sample cases + for N_ens in N_enss + for rng_idx in rng_idxs + calib_filename_suffix = "$(N_ens)_$(rng_idx)" + @info("Perform emulate-sample for L63: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") + min_iter = 1 + skip_iter = 1 + max_iter = 15 # number of EKP iterations to use data from is at most this + + #### + + # need to loop over + rng_seed = 44011 + rng = Random.MersenneTwister(rng_seed) + + # loading relevant data + homedir = pwd() + println(homedir) + figure_save_directory = joinpath(homedir, "output/", calib_directory) + data_save_directory = joinpath(homedir, "output", calib_directory) + loaded_calib_files = [ + joinpath(data_save_directory, "l63_ekp_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, "l63_priors.jld2"), + joinpath(data_save_directory, "l63_calibrate_results_$(calib_filename_suffix).jld2"), + ] + if !isfile(loaded_calib_files[1]) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + + ekpobj = load(loaded_calib_files[1])["ekpobj"] + priors = load(loaded_calib_files[2])["prior"] + truth_sample = load(loaded_calib_files[3])["y"] + truth_params_obj = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space + truth_params_constrained = truth_params_obj.u + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = get_obs_noise_cov(ekpobj) + + n_params = length(truth_params) # "input dim" + output_dim = size(Γy, 1) + ### + ### Emulate: + ### + # choice of machine-learning tool in the emulation stage + nugget = 1e-6 + overrides = Dict( + # "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 20.0, + "n_iteration" => 10, + "n_features_opt" => 160, + ) + n_features = 200 + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) + mlt = ScalarRandomFeatureInterface( + n_features, + n_params, + rng = rng, + kernel_structure = kernel_structure, + optimizer_options = overrides, + ) + + + # Get training points from the EKP iteration number in the second input term + N_iter = min(max_iter, length(get_u(ekpobj))-1) # number of paired iterations taken from EKP + min_iter = min(max_iter, max(1, min_iter)) + if N_iter < 1 + @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") + continue + elseif N_iter == 1 + @warn("Convergence in only 1 iteration, removing train-test split based on iterations") + @info "Train iterations: $(min_iter:skip_iter:N_iter)" + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter+1) + else + @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" + @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + + end + + # Save data + @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs + + # data processing configuration + retain_var = 0.95 + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # Check how well the Gaussian Process regression predicts on the + # true parameters + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + if !(N_iter<2) + y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) + end + println("ML prediction on true parameters: ") + println(vec(y_mean)) + println("true data: ") + println(truth_sample) # what was used as truth + println(" ML predicted standard deviation") + println(sqrt.(diag(y_var[1], 0))) + println("ML MSE (truth): ") + println(mean((truth_sample - vec(y_mean)) .^ 2)) + if !(N_iter<2) + println("ML MSE (test ensemble): ") + println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) + end + + #end + ### + ### Sample: Markov Chain Monte Carlo + ### + # initial values + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + println("initial parameters: ", u0) + + # First let's run a short chain to determine a good step size + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + + # Now begin the actual MCMC + println("Begin MCMC - with step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + + post_mean = mean(posterior) + post_cov = cov(posterior) + println("post_mean") + println(post_mean) + println("post_cov") + println(post_cov) + println("D util") + println(det(inv(post_cov))) + println(" ") + + param_names = get_name(posterior) + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + + # marginal histogram + pp = plot(priors, c = :gray) + plot!(pp, posterior) + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + + # Save data + save( + joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2"), + "posterior", + posterior, + "priors", + priors, + "truth_params_constrained", + truth_params_constrained, + "final_params_constrained", + final_params_constrained, + "input_output_pairs", + input_output_pairs, + "truth_params", + truth_params, + "encoder_schedule", + get_encoder_schedule(emulator), + ) + end + end +end + + +main() diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl index b8b5aad38..93c245e95 100644 --- a/examples/EKIRace/plot_l96_forcing.jl +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -16,16 +16,16 @@ const EKP = EnsembleKalmanProcesses method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] -force_cases = ["const-force", "vec-force", "flux-force"] # problem types +forcing_cases = ["const-force", "vec-force", "flux-force"] # problem types -#### CHOOSE YOUR CASE: (todo: looped over in some fashion) +#### CHOOSE YOUR CASE: # calib_data_dir method=method_cases[1] calibrate_date=Date("2026-05-18", "yyyy-mm-dd") calib_directory="$(method)_$(calibrate_date)" # calib_filename_suffix items to loop over -force_cases=[force_cases[3]] +force_cases=[forcing_cases[3]] N_enss=[50,75,100] rng_idxs=[1,2,3,4] From 3728acec69f9356f7b831f140755e05ab9a76921 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 28 May 2026 16:58:12 -0700 Subject: [PATCH 11/86] improve l63 prior initialization, and create exp_to_leaderboard for it --- examples/EKIRace/calibrate_l63.jl | 26 ++- examples/EKIRace/emulate_sample_l63.jl | 22 +- examples/EKIRace/emulate_sample_l96.jl | 16 +- .../l63_exp_to_leaderboard_utilities.jl | 213 ++++++++++++++++++ ...jl => l96_exp_to_leaderboard_utilities.jl} | 7 +- 5 files changed, 251 insertions(+), 33 deletions(-) create mode 100644 examples/EKIRace/l63_exp_to_leaderboard_utilities.jl rename examples/EKIRace/{exp_to_leaderboard_utilities.jl => l96_exp_to_leaderboard_utilities.jl} (97%) diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index 61be6b6c8..9b69300b9 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -37,16 +37,23 @@ nx = 3 # dimensions of parameter vector nu = 2 truth_params = EnsembleMemberConfig([28.0, 8.0 / 3.0]) +#= +the following better encoder this prior: prior_mean = [3.3, 1.2] prior_cov = [ 0.15^2 0 0 0.5^2 ] -#Creating prior distribution distribution = Parameterized(MvNormal(prior_mean, prior_cov)) -constraint = repeat([no_constraint()], 2) +constraint = repeat([no_constraint()], 2) # TODO: fix this... name = "l63_prior" prior = ParameterDistribution(distribution, constraint, name) +=# + +prior_r = constrained_gaussian("rho", exp(3.3), 4.153, 0, Inf) +prior_b = constrained_gaussian("beta", exp(1.2), 2.016, 0 ,Inf) +prior = combine_distributions([prior_r, prior_b]) +#Creating prior distribution T = 40.0 @@ -138,6 +145,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) @info "Ensemble size: $(N_ens)" for (kk, method) in enumerate(methods) + @info "Method: $(nameof(typeof(method)))" if isa(method, Unscented) ekpobj = EKP.EnsembleKalmanProcess( y, @@ -145,9 +153,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, - accelerator = DefaultAccelerator(), +# accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(), + scheduler = DataMisfitController(terminate_at=100), ) else ekpobj = EKP.EnsembleKalmanProcess( @@ -157,9 +165,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, - accelerator = DefaultAccelerator(), +# accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(), + scheduler = DataMisfitController(terminate_at=100), ) end Ne = get_N_ens(ekpobj) @@ -169,9 +177,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) params_i = get_ϕ_final(prior, ekpobj) # Calculating RMSE_e - ens_mean = mean(params_i, dims = 2)[:] + ens_mean = mean(params_i, dims = 2)[:] # in constrained_space G_ens_mean = lorenz_forward( - EnsembleMemberConfig(exp.(ens_mean)), + EnsembleMemberConfig(ens_mean), x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, 1), lorenz_config_settings, observation_config, @@ -190,7 +198,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) G_ens = hcat( [ lorenz_forward( - EnsembleMemberConfig(exp.(params_i[:, j])), + EnsembleMemberConfig(params_i[:, j]), (x0 .+ ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), nx, Ne))[:, j], lorenz_config_settings, observation_config, diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl index 063634979..32c1aabff 100644 --- a/examples/EKIRace/emulate_sample_l63.jl +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -44,7 +44,7 @@ function main() @info("Perform emulate-sample for L63: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") min_iter = 1 skip_iter = 1 - max_iter = 15 # number of EKP iterations to use data from is at most this + max_iter = 10 # number of EKP iterations to use data from is at most this #### @@ -70,7 +70,7 @@ function main() priors = load(loaded_calib_files[2])["prior"] truth_sample = load(loaded_calib_files[3])["y"] truth_params_obj = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space - truth_params_constrained = truth_params_obj.u + truth_params_constrained = truth_params_obj.u # the parameters here have been exponentiated not using our in-built constraints. truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) Γy = get_obs_noise_cov(ekpobj) @@ -84,7 +84,7 @@ function main() overrides = Dict( # "verbose" => true, "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 20.0, + "cov_sample_multiplier" => 2.0, "n_iteration" => 10, "n_features_opt" => 160, ) @@ -100,7 +100,8 @@ function main() # Get training points from the EKP iteration number in the second input term - N_iter = min(max_iter, length(get_u(ekpobj))-1) # number of paired iterations taken from EKP + iter_end = length(get_u(ekpobj))-1 + N_iter = min(max_iter, iter_end) # number of paired iterations taken from EKP min_iter = min(max_iter, max(1, min_iter)) if N_iter < 1 @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") @@ -108,20 +109,17 @@ function main() elseif N_iter == 1 @warn("Convergence in only 1 iteration, removing train-test split based on iterations") @info "Train iterations: $(min_iter:skip_iter:N_iter)" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter+1) + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) else @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:iter_end) # "next" iterations end - # Save data - @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs - # data processing configuration - retain_var = 0.95 + retain_var = 0.99 encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] encoder_kwargs = encoder_kwargs_from(ekpobj, priors) @@ -182,7 +180,7 @@ function main() param_names = get_name(posterior) - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + posterior_samples = reduce(vcat,[get_distribution(posterior)[name] for name in get_name(posterior)]) #samples are columns constrained_posterior_samples = mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) @@ -200,7 +198,7 @@ function main() # Save data save( - joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, "l63_posterior_$(calib_filename_suffix).jld2"), "posterior", posterior, "priors", diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index 82884a3da..157343ce6 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -111,8 +111,9 @@ function main() ) - # Get training points from the EKP iteration number in the second input term - N_iter = min(max_iter, length(get_u(ekpobj))-1) # number of paired iterations taken from EKP + # Get training points from the EKP iteration number in the second input term + iter_end = length(get_u(ekpobj))-1 + N_iter = min(max_iter, iter_end) # number of paired iterations taken from EKP min_iter = min(max_iter, max(1, min_iter)) if N_iter < 1 @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") @@ -120,18 +121,15 @@ function main() elseif N_iter == 1 @warn("Convergence in only 1 iteration, removing train-test split based on iterations") @info "Train iterations: $(min_iter:skip_iter:N_iter)" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter+1) + input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) else @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:(length(get_u(ekpobj)) - 1)) # "next" iterations + input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:iter_end) # "next" iteration end - - # Save data - @save joinpath(data_save_directory, "input_output_pairs_$(calib_filename_suffix).jld2") input_output_pairs - + # data processing configuration retain_var = 0.95 encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] @@ -194,7 +192,7 @@ function main() param_names = get_name(posterior) - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + posterior_samples = reduce(vcat,[get_distribution(posterior)[name] for name in get_name(posterior)]) #samples are columns constrained_posterior_samples = mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl new file mode 100644 index 000000000..f94be253f --- /dev/null +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -0,0 +1,213 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +# Step 1, Read and extract the available experiments +method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] + +#### CHOOSE YOUR CASE: +# calib_data_dir +method = method_cases[1] +calibrate_date = Date("2026-05-28", "yyyy-mm-dd") +calib_directory = "$(method)_$(calibrate_date)" + +# calib_filename_suffix items to loop over +N_enss = [20,25,30] +rng_idxs = [1, 2, 3, 4] + +### For saving the cases in the leaderboard style + +method_cases_key = Dict( + "Inversion" => "ces-eki-dmc", + "TransformInversion" => "ces-etki-dmc", + "Unscented" => "ces-uki-dmc", + "GaussNewtonInversion" => "ces-iekf", +) +nc_save_filename = "$(method_cases_key[method])_l63_ensemble_results_$(calibrate_date).nc" + +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_directory) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + calib_filename_suffix = "$(N_ens)_$(rng_idx)" + data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + if isfile(data_file) + push!(valid_files, calib_filename_suffix) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l63.jl first.") +end + +### Load data + +# Load first valid file to determine parameter dimension +first_loaded = JLD2.load(joinpath(data_save_directory, "posterior_$(valid_files[1]).jld2")) +n_params = length(vec(mean(first_loaded["posterior"]))) + +targets = [1.0] + +n_rng = length(rng_idxs) +n_ens = length(N_enss) +n_targets = length(targets) + +# Pre-allocate storage arrays indexed as (random_seed, ensemble_size, rmse_target, algorithm_type, ...) +true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) + +for (N_ens, rng_idx) in valid_file_items + calib_filename_suffix = "$(N_ens)_$(rng_idx)" + post_filename = "posterior_$(calib_filename_suffix).jld2" + @info "loading case $(post_filename)" + loaded = JLD2.load(joinpath(data_save_directory, post_filename)) + + post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) + truth_params = loaded["truth_params"] # truth in unconstrained space + post_name = get_name(post_dist) + + pm = vec(mean(post_dist)) # posterior mean, shape (n_params,) + pc = cov(post_dist) # posterior covariance, shape (n_params, n_params) + + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat,[get_distribution(post_dist)[name] for name in get_name(post_dist)]) #samples are columns + + pmode_idx = argmax(logpdf(post_normal, post_samples)) + pmode = post_samples[:, pmode_idx] + + + # Load ekpobj to compute calibration evaluation count + ekp_loaded = JLD2.load(joinpath(data_save_directory, "l63_ekp_$(calib_filename_suffix).jld2")) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps + + # Array indices + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + diff = pm - truth_params + + for l in 1:n_targets + true_param_arr[i, j, l, 1, :] = truth_params + post_mean_arr[i, j, l, 1, :] = pm + post_cov_arr[i, j, l, 1, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, l, 1] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + # total forward model evaluations to reach target + n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens + end + +end + + +### Saving the data in nc + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "rmse_target", n_targets) +defDim(ds, "algorithm_type", 1) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) # second param axis for covariance matrix + +# Coordinate variables +alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) +alg_var[1] = method_cases_key[method] + +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) +rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" +rmse_var[:] = targets + +# Data variables + +n_evals_v = defVar( + ds, + "n_evals_to_target", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +n_evals_v.attrib["description"] = "Number of forward model evaluations (i.e., algorithm cost): conv_alg_iters * ensemble_size. Stored as float in case of NaN values." +n_evals_v[:, :, :, :] = n_evals_arr + +true_param_v = defVar( + ds, + "true_param", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + fillvalue = NaN, +) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :, :, :] = true_param_arr + +post_mean_v = defVar( + ds, + "post_mean", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + fillvalue = NaN, +) +post_mean_v.attrib["description"] = "mean(posterior)" +post_mean_v[:, :, :, :, :] = post_mean_arr + +post_cov_v = defVar( + ds, + "post_cov", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim", "param_dim_2"); + fillvalue = NaN, +) +post_cov_v.attrib["description"] = "cov(posterior)" +post_cov_v[:, :, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar( + ds, + "mahalanobis", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" +mahalanobis_v[:, :, :, :] = mahal_arr + +posterior_logpdf_true_v_map_v = defVar( + ds, + "posterior_logpdf_true_v_map", + Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + fillvalue = NaN, +) +posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" +posterior_logpdf_true_v_map_v[:, :, :, :] = logpdf_true_v_map_arr + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl similarity index 97% rename from examples/EKIRace/exp_to_leaderboard_utilities.jl rename to examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index be717b479..fe0151c42 100644 --- a/examples/EKIRace/exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -97,9 +97,10 @@ for (fc, N_ens, rng_idx) in valid_file_items C_reg = Symmetric(pc + 1e-10 * I) post_normal = MvNormal(pm, C_reg) - post_array = get_distribution(post_dist)[get_name(post_dist)[1]] # safe as posterior never has constraints - just samples - pmode_idx = argmax(logpdf(post_normal, post_array)) - pmode = post_array[:, pmode_idx] + post_samples = reduce(vcat,[get_distribution(post_dist)[name] for name in get_name(post_dist)]) #samples are columns + + pmode_idx = argmax(logpdf(post_normal, post_samples)) + pmode = post_samples[:, pmode_idx] # Load ekpobj to compute calibration evaluation count From b596d10087c9e7c104694a4cdebefeb99be71381 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 10:57:06 -0700 Subject: [PATCH 12/86] slow down timestep for l63 --- examples/EKIRace/calibrate_l63.jl | 8 ++++---- examples/EKIRace/calibrate_l96.jl | 6 ++++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index 9b69300b9..1f3cbe44b 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -22,10 +22,10 @@ save_all_ekp = true ############### Choose problem type and structure ###################### ######################################################################## -N_ens_sizes = [20, 25, 30] # list of number of ensemble members (should be problem dependent) +N_ens_sizes = [10, 25, 40] # list of number of ensemble members (should be problem dependent) N_iter = 20 # maximum number of EKI iterations allowed target_rmse = 1.0 # target RMSE -n_repeats = 4 +n_repeats = 20 rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running Lorenz 63 problem" @info "Maximum number of EKI iterations: $N_iter" @@ -155,7 +155,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) verbose = verbose_flag, # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DataMisfitController(terminate_at=100), + scheduler = DefaultScheduler(0.1), ) else ekpobj = EKP.EnsembleKalmanProcess( @@ -167,7 +167,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) verbose = verbose_flag, # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DataMisfitController(terminate_at=100), + scheduler = DefaultScheduler(0.1), ) end Ne = get_N_ens(ekpobj) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index ef92a7e1e..816b741b4 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -27,7 +27,7 @@ save_all_ekp = true cases = ["const-force", "vec-force", "flux-force"] # problem types # User specifications -case = cases[1] # choose problem type +case = cases[2] # choose problem type if case == "const-force" N_ens_sizes = [5, 15, 30] elseif case == "vec-force" @@ -41,7 +41,7 @@ end N_iter = 20 # maximum number of EKI iterations allowed target_rmse = 1.0 # target RMSE -n_repeats = 4 +n_repeats = 20 rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running $case case" @info "Maximum number of EKI iterations: $N_iter" @@ -272,6 +272,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) verbose = verbose_flag, # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), + #scheduler = DefaultScheduler(0.1), scheduler = DataMisfitController(terminate_at=100), ) else @@ -285,6 +286,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), scheduler = DataMisfitController(terminate_at=100), +# scheduler = DefaultScheduler(0.1), ) end Ne = get_N_ens(ekpobj) From 0d225d43bfb2583cbd5d6041de5c841e19f4b315 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 11:00:08 -0700 Subject: [PATCH 13/86] weighted mle --- examples/EKIRace/emulate_sample_l63.jl | 21 ++++++++++-------- examples/EKIRace/emulate_sample_l96.jl | 30 ++++++++++++++++---------- 2 files changed, 31 insertions(+), 20 deletions(-) diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl index 32c1aabff..bd94c5507 100644 --- a/examples/EKIRace/emulate_sample_l63.jl +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -34,8 +34,8 @@ function main() calib_directory="$(method)_$(calibrate_date)" # calib_filename_suffix items to loop over - N_enss = [20,25,30] - rng_idxs=[1,2,3,4] + N_enss = [10,25,40] + rng_idxs=collect(1:20) # emulate_sample cases for N_ens in N_enss @@ -84,11 +84,11 @@ function main() overrides = Dict( # "verbose" => true, "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 2.0, + "cov_sample_multiplier" => 5.0, "n_iteration" => 10, - "n_features_opt" => 160, + "n_features_opt" => 60, ) - n_features = 200 + n_features = 100 kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) mlt = ScalarRandomFeatureInterface( n_features, @@ -143,11 +143,14 @@ function main() println(truth_sample) # what was used as truth println(" ML predicted standard deviation") println(sqrt.(diag(y_var[1], 0))) - println("ML MSE (truth): ") - println(mean((truth_sample - vec(y_mean)) .^ 2)) + println("ML weighted-MSE (truth): ") + diff = (truth_sample - vec(y_mean)) + Γypluscov = Symmetric(mean(y_var) + Γy) + println(mean(diff' * (Γypluscov\ diff))) if !(N_iter<2) - println("ML MSE (test ensemble): ") - println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) + println("ML weighted-MSE (test ensemble): ") + difftest = get_outputs(input_output_pairs_test) - y_mean_test + println(mean(difftest' * (Γypluscov \ difftest))) end #end diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index 157343ce6..c10bef139 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -31,13 +31,18 @@ function main() #### CHOOSE YOUR CASE: (todo: looped over in some fashion) # calib_data_dir method=method_cases[1] - calibrate_date=Date("2026-05-18", "yyyy-mm-dd") + calibrate_date=Date("2026-05-28", "yyyy-mm-dd") calib_directory="$(method)_$(calibrate_date)" # calib_filename_suffix items to loop over - force_cases=[force_cases[3]] - N_enss = [100] # [5,15,30] #[50,75,100] - rng_idxs=[2,3,4] + force_cases=[force_cases[2]] + N_ens_cases =[ + [5, 15, 30], + [50, 75, 100], + [50, 75, 100], + ] + N_enss = N_ens_cases[2] + rng_idxs=collect(1:20) # emulate_sample cases for force_case in force_cases @@ -155,13 +160,16 @@ function main() println(truth_sample) # what was used as truth println(" ML predicted standard deviation") println(sqrt.(diag(y_var[1], 0))) - println("ML MSE (truth): ") - println(mean((truth_sample - vec(y_mean)) .^ 2)) + println("ML weighted-MSE (truth): ") + diff = (truth_sample - vec(y_mean)) + Γypluscov = Symmetric(mean(y_var) + Γy) + println(mean(diff' * (Γypluscov\ diff))) if !(N_iter<2) - println("ML MSE (test ensemble): ") - println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 2)) + println("ML weighted-MSE (test ensemble): ") + difftest = get_outputs(input_output_pairs_test) - y_mean_test + println(mean(difftest' * (Γypluscov \ difftest))) end - + #end ### ### Sample: Markov Chain Monte Carlo @@ -204,13 +212,13 @@ function main() vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) end - figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)") + figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)") savefig(figpath * ".png") savefig(figpath * ".pdf") # Save data save( - joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, "l96_posterior_$(calib_filename_suffix).jld2"), "posterior", posterior, "priors", From 3a4ec8e8f027f2e7b5aa41a16d7c4becfc40d963 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 11:00:47 -0700 Subject: [PATCH 14/86] update exp_to_lb --- examples/EKIRace/l63_exp_to_leaderboard_utilities.jl | 10 +++++----- examples/EKIRace/l96_exp_to_leaderboard_utilities.jl | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index f94be253f..012c1ec83 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -17,8 +17,8 @@ calibrate_date = Date("2026-05-28", "yyyy-mm-dd") calib_directory = "$(method)_$(calibrate_date)" # calib_filename_suffix items to loop over -N_enss = [20,25,30] -rng_idxs = [1, 2, 3, 4] +N_enss = [10,25,40] +rng_idxs = collect(1:20) ### For saving the cases in the leaderboard style @@ -40,7 +40,7 @@ valid_files = [] for N_ens in N_enss for rng_idx in rng_idxs calib_filename_suffix = "$(N_ens)_$(rng_idx)" - data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + data_file = joinpath(data_save_directory, "l63_posterior_$(calib_filename_suffix).jld2") if isfile(data_file) push!(valid_files, calib_filename_suffix) push!(valid_file_items, (N_ens, rng_idx)) @@ -58,7 +58,7 @@ end ### Load data # Load first valid file to determine parameter dimension -first_loaded = JLD2.load(joinpath(data_save_directory, "posterior_$(valid_files[1]).jld2")) +first_loaded = JLD2.load(joinpath(data_save_directory, "l63_posterior_$(valid_files[1]).jld2")) n_params = length(vec(mean(first_loaded["posterior"]))) targets = [1.0] @@ -77,7 +77,7 @@ n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) for (N_ens, rng_idx) in valid_file_items calib_filename_suffix = "$(N_ens)_$(rng_idx)" - post_filename = "posterior_$(calib_filename_suffix).jld2" + post_filename = "l63_posterior_$(calib_filename_suffix).jld2" @info "loading case $(post_filename)" loaded = JLD2.load(joinpath(data_save_directory, post_filename)) diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index fe0151c42..b002ec957 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -47,7 +47,7 @@ valid_files = [] for N_ens in N_enss for rng_idx in rng_idxs calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + data_file = joinpath(data_save_directory, "l96_posterior_$(calib_filename_suffix).jld2") if isfile(data_file) push!(valid_files, calib_filename_suffix) push!(valid_file_items, (force_case, N_ens, rng_idx)) @@ -65,7 +65,7 @@ end ### Load data # Load first valid file to determine parameter dimension -first_loaded = JLD2.load(joinpath(data_save_directory, "posterior_$(valid_files[1]).jld2")) +first_loaded = JLD2.load(joinpath(data_save_directory, "l96_posterior_$(valid_files[1]).jld2")) n_params = length(vec(mean(first_loaded["posterior"]))) targets = [1.0] @@ -84,7 +84,7 @@ n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) for (fc, N_ens, rng_idx) in valid_file_items calib_filename_suffix = "$(fc)_$(N_ens)_$(rng_idx)" - post_filename = "posterior_$(calib_filename_suffix).jld2" + post_filename = "l96_posterior_$(calib_filename_suffix).jld2" @info "loading case $(post_filename)" loaded = JLD2.load(joinpath(data_save_directory, post_filename)) From 60614e0ea3652627246d1ea77172d3aca6d363ce Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 12:22:10 -0700 Subject: [PATCH 15/86] add filename helpers and shared config --- examples/EKIRace/calibrate_l63.jl | 35 ++--- examples/EKIRace/calibrate_l96.jl | 50 ++---- examples/EKIRace/emulate_sample_l63.jl | 44 +++--- examples/EKIRace/emulate_sample_l96.jl | 62 ++++---- examples/EKIRace/experiment_config.jl | 143 ++++++++++++++++++ .../l63_exp_to_leaderboard_utilities.jl | 46 +++--- .../l96_exp_to_leaderboard_utilities.jl | 57 +++---- 7 files changed, 259 insertions(+), 178 deletions(-) create mode 100644 examples/EKIRace/experiment_config.jl diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index 1f3cbe44b..a4f416710 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -15,6 +15,7 @@ using EnsembleKalmanProcesses.Localizers const EKP = EnsembleKalmanProcesses include("Lorenz63.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") verbose_flag = false save_all_ekp = true @@ -22,11 +23,12 @@ save_all_ekp = true ############### Choose problem type and structure ###################### ######################################################################## -N_ens_sizes = [10, 25, 40] # list of number of ensemble members (should be problem dependent) -N_iter = 20 # maximum number of EKI iterations allowed -target_rmse = 1.0 # target RMSE -n_repeats = 20 -rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds +cfg = experiment_config(:l63) +N_ens_sizes = cfg.N_ens_sizes +N_iter = cfg.N_iter +target_rmse = cfg.target_rmse +n_repeats = cfg.n_repeats +rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running Lorenz 63 problem" @info "Maximum number of EKI iterations: $N_iter" @info "RMSE target: $target_rmse" @@ -121,12 +123,7 @@ conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how man final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) -method_names = [ - ("Inversion(prior)", "TEKI"), - ("TransformInversion(prior)", "ETKI"), - ("GaussNewtonInversion(prior)", "GNKI"), - ("Unscented(prior; impose_prior=true)", "UKI"), -] +# method_names defined in experiment_config.jl for (rr, rng_seed) in enumerate(rng_seeds) @@ -224,29 +221,25 @@ for (rr, rng_seed) in enumerate(rng_seeds) final_ensemble = get_ϕ_final(prior, ekpobj) # save ekp files - per_method_dir = joinpath(output_dir,"$(nameof(typeof(method)))_$(today())") - if rr == 1 && ee == 1 + per_method_dir = joinpath(output_dir, calib_directory(nameof(typeof(method)), cfg)) + if rr == 1 && ee == 1 if !isdir(per_method_dir) mkpath(per_method_dir) end - prior_filename = joinpath(per_method_dir,"l63_priors.jld2") - JLD2.save(prior_filename,"prior", prior) + JLD2.save(joinpath(per_method_dir, prior_filename(cfg)), "prior", prior) end if save_all_ekp # JLD2 - ekp_filename = joinpath(per_method_dir, "l63_ekp_$(N_ens)_$(rr).jld2") JLD2.save( - ekp_filename, + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rr)), "N_ens", N_ens, "method", method, "ekpobj", ekpobj, ) - results_filename = joinpath(per_method_dir, "l63_calibrate_results_$(N_ens)_$(rr).jld2") u_stored = get_u(ekpobj, return_array = false) g_stored = get_g(ekpobj, return_array = false) - JLD2.save( - results_filename, + joinpath(per_method_dir, results_filename(cfg, N_ens, rr)), "y", y, "R", R, "inputs", u_stored, @@ -260,7 +253,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) end # Saving data: -data_filename = joinpath(output_dir, "l63_output_$(today()).jld2") +data_filename = joinpath(output_dir, summary_filename(cfg)) JLD2.save( data_filename, "configuration", diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index 816b741b4..d741dcc2f 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -18,6 +18,7 @@ using EnsembleKalmanProcesses.Localizers const EKP = EnsembleKalmanProcesses include("Lorenz96.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") verbose_flag = false save_all_ekp = true @@ -25,24 +26,14 @@ save_all_ekp = true ############### Choose problem type and structure ###################### ######################################################################## -cases = ["const-force", "vec-force", "flux-force"] # problem types -# User specifications -case = cases[2] # choose problem type -if case == "const-force" - N_ens_sizes = [5, 15, 30] -elseif case == "vec-force" - N_ens_sizes = [50, 75, 100] -elseif case == "flux-force" - N_ens_sizes = [50, 75, 100] -else - throw(ArgumentError("Expected case ∈ $(cases). Got case = $case.")) -end - - -N_iter = 20 # maximum number of EKI iterations allowed -target_rmse = 1.0 # target RMSE -n_repeats = 20 -rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +case = cfg.force_case # "const-force", "vec-force", or "flux-force" +N_ens_sizes = cfg.N_ens_sizes +N_iter = cfg.N_iter +target_rmse = cfg.target_rmse +n_repeats = cfg.n_repeats +rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running $case case" @info "Maximum number of EKI iterations: $N_iter" @info "RMSE target: $target_rmse" @@ -238,12 +229,7 @@ conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how man final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) -method_names = [ - ("Inversion(prior)", "TEKI"), - ("TransformInversion(prior)", "ETKI"), - ("GaussNewtonInversion(prior)", "GNKI"), - ("Unscented(prior; impose_prior=true)", "UKI"), -] +# method_names defined in experiment_config.jl for (rr, rng_seed) in enumerate(rng_seeds) @info "Random seed: $(rng_seed)" @@ -336,29 +322,25 @@ for (rr, rng_seed) in enumerate(rng_seeds) final_ensemble = get_ϕ_final(prior, ekpobj) # save ekp files - per_method_dir = joinpath(output_dir,"$(nameof(typeof(method)))_$(today())") - if rr == 1 && ee == 1 + per_method_dir = joinpath(output_dir, calib_directory(nameof(typeof(method)), cfg)) + if rr == 1 && ee == 1 if !isdir(per_method_dir) mkpath(per_method_dir) end - prior_filename = joinpath(per_method_dir,"l96_priors_$(case).jld2") - JLD2.save(prior_filename,"prior", prior) + JLD2.save(joinpath(per_method_dir, prior_filename(cfg)), "prior", prior) end if save_all_ekp # JLD2 - ekp_filename = joinpath(per_method_dir, "l96_ekp_$(case)_$(N_ens)_$(rr).jld2") JLD2.save( - ekp_filename, + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rr)), "N_ens", N_ens, "method", method, "ekpobj", ekpobj, ) - results_filename = joinpath(per_method_dir, "l96_calibrate_results_$(case)_$(N_ens)_$(rr).jld2") u_stored = get_u(ekpobj, return_array = false) g_stored = get_g(ekpobj, return_array = false) - JLD2.save( - results_filename, + joinpath(per_method_dir, results_filename(cfg, N_ens, rr)), "y", y, "R", R, "inputs", u_stored, @@ -374,7 +356,7 @@ end # Saving summary data: -data_filename = joinpath(output_dir, "l96_calibrate_$(case)_$(today()).jld2") +data_filename = joinpath(output_dir, summary_filename(cfg)) JLD2.save( data_filename, "configuration", diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl index bd94c5507..73dbe87e5 100644 --- a/examples/EKIRace/emulate_sample_l63.jl +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -20,31 +20,27 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers -include("Lorenz63.jl") # Contains Lorenz 96 source code +include("Lorenz63.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") function main() - method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] - - #### CHOOSE YOUR CASE: (todo: looped over in some fashion) - # calib_data_dir - method=method_cases[1] - calibrate_date=Date("2026-05-28", "yyyy-mm-dd") - calib_directory="$(method)_$(calibrate_date)" - - # calib_filename_suffix items to loop over - N_enss = [10,25,40] - rng_idxs=collect(1:20) + #### CHOOSE YOUR CASE: + cfg = experiment_config(:l63) + method = method_cases[1] # method_cases defined in experiment_config.jl + calib_dir = calib_directory(method, cfg) + N_enss = cfg.N_ens_sizes + rng_idxs = collect(1:cfg.n_repeats) # emulate_sample cases for N_ens in N_enss for rng_idx in rng_idxs - calib_filename_suffix = "$(N_ens)_$(rng_idx)" - @info("Perform emulate-sample for L63: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") + calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L63: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") min_iter = 1 skip_iter = 1 - max_iter = 10 # number of EKP iterations to use data from is at most this + max_iter = cfg.max_iter #### @@ -55,12 +51,12 @@ function main() # loading relevant data homedir = pwd() println(homedir) - figure_save_directory = joinpath(homedir, "output/", calib_directory) - data_save_directory = joinpath(homedir, "output", calib_directory) + figure_save_directory = joinpath(homedir, "output/", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) loaded_calib_files = [ - joinpath(data_save_directory, "l63_ekp_$(calib_filename_suffix).jld2"), - joinpath(data_save_directory, "l63_priors.jld2"), - joinpath(data_save_directory, "l63_calibrate_results_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)), + joinpath(data_save_directory, prior_filename(cfg)), + joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)), ] if !isfile(loaded_calib_files[1]) throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) @@ -86,9 +82,9 @@ function main() "scheduler" => DataMisfitController(terminate_at = 1000.0), "cov_sample_multiplier" => 5.0, "n_iteration" => 10, - "n_features_opt" => 60, + "n_features_opt" => cfg.n_features_opt, ) - n_features = 100 + n_features = cfg.n_features kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) mlt = ScalarRandomFeatureInterface( n_features, @@ -119,7 +115,7 @@ function main() end # data processing configuration - retain_var = 0.99 + retain_var = cfg.retain_var encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] encoder_kwargs = encoder_kwargs_from(ekpobj, priors) @@ -201,7 +197,7 @@ function main() # Save data save( - joinpath(data_save_directory, "l63_posterior_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), "posterior", posterior, "priors", diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index c10bef139..8c8694e7e 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -20,39 +20,30 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers -include("Lorenz96.jl") # Contains Lorenz 96 source code +include("Lorenz96.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") function main() - method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] - force_cases = ["const-force", "vec-force", "flux-force"] # problem types - - #### CHOOSE YOUR CASE: (todo: looped over in some fashion) - # calib_data_dir - method=method_cases[1] - calibrate_date=Date("2026-05-28", "yyyy-mm-dd") - calib_directory="$(method)_$(calibrate_date)" - - # calib_filename_suffix items to loop over - force_cases=[force_cases[2]] - N_ens_cases =[ - [5, 15, 30], - [50, 75, 100], - [50, 75, 100], - ] - N_enss = N_ens_cases[2] - rng_idxs=collect(1:20) + #### CHOOSE YOUR CASE: + @assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For emulate_sample_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" + cfg = experiment_config(EXPERIMENT) + method = method_cases[1] # method_cases defined in experiment_config.jl + calib_dir = calib_directory(method, cfg) + force_cases = [cfg.force_case] + N_enss = cfg.N_ens_sizes + rng_idxs = collect(1:cfg.n_repeats) # emulate_sample cases for force_case in force_cases for N_ens in N_enss for rng_idx in rng_idxs - calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - @info("Perform emulate-sample for L96: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") + calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") min_iter = 1 skip_iter = 1 - max_iter = 15 # number of EKP iterations to use data from is at most this + max_iter = cfg.max_iter #### @@ -63,12 +54,12 @@ function main() # loading relevant data homedir = pwd() println(homedir) - figure_save_directory = joinpath(homedir, "output/", calib_directory) - data_save_directory = joinpath(homedir, "output", calib_directory) + figure_save_directory = joinpath(homedir, "output/", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) loaded_calib_files = [ - joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2"), - joinpath(data_save_directory, "l96_priors_$(force_case).jld2"), - joinpath(data_save_directory, "l96_calibrate_results_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)), + joinpath(data_save_directory, prior_filename(cfg)), + joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)), ] if !isfile(loaded_calib_files[1]) throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) @@ -98,14 +89,19 @@ function main() ### # choice of machine-learning tool in the emulation stage nugget = 1e-6 + if N_ens * length(get_u(ekpobj)) < 200 + csm = 5.0 + else + csm = 1.0 + end overrides = Dict( # "verbose" => true, "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 20.0, + "cov_sample_multiplier" => csm, "n_iteration" => 10, - "n_features_opt" => 160, + "n_features_opt" => cfg.n_features_opt, ) - n_features = 200 + n_features = cfg.n_features kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) mlt = ScalarRandomFeatureInterface( n_features, @@ -136,7 +132,7 @@ function main() end # data processing configuration - retain_var = 0.95 + retain_var = cfg.retain_var encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] encoder_kwargs = encoder_kwargs_from(ekpobj, priors) @@ -215,10 +211,10 @@ function main() figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)") savefig(figpath * ".png") savefig(figpath * ".pdf") - + # Save data save( - joinpath(data_save_directory, "l96_posterior_$(calib_filename_suffix).jld2"), + joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), "posterior", posterior, "priors", diff --git a/examples/EKIRace/experiment_config.jl b/examples/EKIRace/experiment_config.jl new file mode 100644 index 000000000..483dff27f --- /dev/null +++ b/examples/EKIRace/experiment_config.jl @@ -0,0 +1,143 @@ +using Dates + +######################################################################## +############### USER TOGGLE ######################################### +######################################################################## +# Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux +const EXPERIMENT = :l63 + +# Date identifying this calibration run — written by calibrate, read by +# emulate_sample and exp_to_leaderboard (so all stages stay in sync). +const calibrate_date = Date("2026-05-28", "yyyy-mm-dd") + +######################################################################## +############### SHARED CONSTANTS #################################### +######################################################################## +const method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] + +const method_names = [ + ("Inversion(prior)", "TEKI"), + ("TransformInversion(prior)", "ETKI"), + ("GaussNewtonInversion(prior)", "GNKI"), + ("Unscented(prior; impose_prior=true)", "UKI"), +] + +const method_cases_key = Dict( + "Inversion" => "ces-eki-dmc", + "TransformInversion" => "ces-etki-dmc", + "Unscented" => "ces-uki-dmc", + "GaussNewtonInversion" => "ces-iekf", +) + +const forcing_cases_key = Dict( + "const-force" => "ensemble_results", + "vec-force" => "spatial_forcing_ensemble_results", + "flux-force" => "nn_forcing_ensemble_results", +) + +######################################################################## +############### PER-CASE CONFIG ##################################### +######################################################################## +function experiment_config(case::Symbol) + if case == :l63 + return ( + model = "l63", + force_case = nothing, + N_ens_sizes = [10, 25, 40], + N_iter = 20, + target_rmse = 1.0, + n_repeats = 20, + max_iter = 10, + retain_var = 0.99, + n_features = 100, + n_features_opt = 60, + calibrate_date = calibrate_date, + ) + elseif case == :l96_const + return ( + model = "l96", + force_case = "const-force", + N_ens_sizes = [5, 15, 30], + N_iter = 20, + target_rmse = 1.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_vec + return ( + model = "l96", + force_case = "vec-force", + N_ens_sizes = [50, 75, 100], + N_iter = 20, + target_rmse = 1.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_flux + return ( + model = "l96", + force_case = "flux-force", + N_ens_sizes = [50, 75, 100], + N_iter = 20, + target_rmse = 1.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + else + throw(ArgumentError("Unknown experiment: $case. Expected one of :l63, :l96_const, :l96_vec, :l96_flux")) + end +end + +######################################################################## +############### FILENAME BUILDERS ################################### +######################################################################## +function case_suffix(cfg, N_ens, rng_idx) + if cfg.force_case === nothing + return "$(N_ens)_$(rng_idx)" + else + return "$(cfg.force_case)_$(N_ens)_$(rng_idx)" + end +end + +calib_directory(method, cfg) = "$(method)_$(cfg.calibrate_date)" + +function prior_filename(cfg) + if cfg.force_case === nothing + return "$(cfg.model)_priors.jld2" + else + return "$(cfg.model)_priors_$(cfg.force_case).jld2" + end +end + +ekp_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_ekp_$(case_suffix(cfg, N_ens, rng_idx)).jld2" +results_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_calibrate_results_$(case_suffix(cfg, N_ens, rng_idx)).jld2" +posterior_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_posterior_$(case_suffix(cfg, N_ens, rng_idx)).jld2" + +function summary_filename(cfg) + if cfg.force_case === nothing + return "$(cfg.model)_output_$(cfg.calibrate_date).jld2" + else + return "$(cfg.model)_calibrate_$(cfg.force_case)_$(cfg.calibrate_date).jld2" + end +end + +function nc_filename(cfg, method) + key = method_cases_key[method] + if cfg.force_case === nothing + return "$(key)_$(cfg.model)_ensemble_results_$(cfg.calibrate_date).nc" + else + return "$(key)_$(cfg.model)_$(forcing_cases_key[cfg.force_case])_$(cfg.calibrate_date).nc" + end +end diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index 012c1ec83..6b05fceb1 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -7,42 +7,31 @@ using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses +include("experiment_config.jl") + # Step 1, Read and extract the available experiments -method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] #### CHOOSE YOUR CASE: -# calib_data_dir -method = method_cases[1] -calibrate_date = Date("2026-05-28", "yyyy-mm-dd") -calib_directory = "$(method)_$(calibrate_date)" - -# calib_filename_suffix items to loop over -N_enss = [10,25,40] -rng_idxs = collect(1:20) - -### For saving the cases in the leaderboard style - -method_cases_key = Dict( - "Inversion" => "ces-eki-dmc", - "TransformInversion" => "ces-etki-dmc", - "Unscented" => "ces-uki-dmc", - "GaussNewtonInversion" => "ces-iekf", -) -nc_save_filename = "$(method_cases_key[method])_l63_ensemble_results_$(calibrate_date).nc" +cfg = experiment_config(:l63) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +nc_save_filename = nc_filename(cfg, method) ### determine valid files before we try to load homedir = joinpath(pwd()) -data_save_directory = joinpath(homedir, "output", calib_directory) +data_save_directory = joinpath(homedir, "output", calib_dir) valid_file_items = [] valid_files = [] for N_ens in N_enss for rng_idx in rng_idxs - calib_filename_suffix = "$(N_ens)_$(rng_idx)" - data_file = joinpath(data_save_directory, "l63_posterior_$(calib_filename_suffix).jld2") + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) if isfile(data_file) - push!(valid_files, calib_filename_suffix) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) push!(valid_file_items, (N_ens, rng_idx)) end end @@ -58,7 +47,7 @@ end ### Load data # Load first valid file to determine parameter dimension -first_loaded = JLD2.load(joinpath(data_save_directory, "l63_posterior_$(valid_files[1]).jld2")) +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posterior"]))) targets = [1.0] @@ -76,10 +65,9 @@ logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) for (N_ens, rng_idx) in valid_file_items - calib_filename_suffix = "$(N_ens)_$(rng_idx)" - post_filename = "l63_posterior_$(calib_filename_suffix).jld2" - @info "loading case $(post_filename)" - loaded = JLD2.load(joinpath(data_save_directory, post_filename)) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) truth_params = loaded["truth_params"] # truth in unconstrained space @@ -97,7 +85,7 @@ for (N_ens, rng_idx) in valid_file_items # Load ekpobj to compute calibration evaluation count - ekp_loaded = JLD2.load(joinpath(data_save_directory, "l63_ekp_$(calib_filename_suffix).jld2")) + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps # Array indices diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index b002ec957..a02880f4c 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -7,50 +7,34 @@ using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses +include("experiment_config.jl") + # Step 1, Read and extract the available experiments -method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] -force_cases = ["const-force", "vec-force", "flux-force"] # problem types #### CHOOSE YOUR CASE: -# calib_data_dir -method = method_cases[1] -calibrate_date = Date("2026-05-18", "yyyy-mm-dd") -calib_directory = "$(method)_$(calibrate_date)" - -# calib_filename_suffix items to loop over -force_case = force_cases[2] -N_enss = [50,75,100] #[5, 15, 30] -rng_idxs = [1, 2, 3, 4] - -### For saving the cases in the leaderboard style -forcing_cases_key = Dict( - "const-force" => "ensemble_results", - "vec-force" => "spatial_forcing_ensemble_results", - "flux-force" => "nn_forcing_ensemble_results", -) +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For l96_exp_to_leaderboard_utilities.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +force_case = cfg.force_case +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) -method_cases_key = Dict( - "Inversion" => "ces-eki-dmc", - "TransformInversion" => "ces-etki-dmc", - "Unscented" => "ces-uki-dmc", - "GaussNewtonInversion" => "ces-iekf", -) -nc_save_filename = "$(method_cases_key[method])_l96_$(forcing_cases_key[force_case])_$(calibrate_date).nc" +nc_save_filename = nc_filename(cfg, method) ### determine valid files before we try to load homedir = joinpath(pwd()) -data_save_directory = joinpath(homedir, "output", calib_directory) +data_save_directory = joinpath(homedir, "output", calib_dir) valid_file_items = [] valid_files = [] for N_ens in N_enss for rng_idx in rng_idxs - calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - data_file = joinpath(data_save_directory, "l96_posterior_$(calib_filename_suffix).jld2") + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) if isfile(data_file) - push!(valid_files, calib_filename_suffix) - push!(valid_file_items, (force_case, N_ens, rng_idx)) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) end end end @@ -65,7 +49,7 @@ end ### Load data # Load first valid file to determine parameter dimension -first_loaded = JLD2.load(joinpath(data_save_directory, "l96_posterior_$(valid_files[1]).jld2")) +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posterior"]))) targets = [1.0] @@ -82,11 +66,10 @@ mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -for (fc, N_ens, rng_idx) in valid_file_items - calib_filename_suffix = "$(fc)_$(N_ens)_$(rng_idx)" - post_filename = "l96_posterior_$(calib_filename_suffix).jld2" - @info "loading case $(post_filename)" - loaded = JLD2.load(joinpath(data_save_directory, post_filename)) +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) truth_params = loaded["truth_params"] # truth in unconstrained space @@ -104,7 +87,7 @@ for (fc, N_ens, rng_idx) in valid_file_items # Load ekpobj to compute calibration evaluation count - ekp_loaded = JLD2.load(joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2")) + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps # Array indices From 03b9ee25b6d30da2f6d84c8e096ff19214bf446b Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 12:58:51 -0700 Subject: [PATCH 16/86] move to DMC terminations time not RMSE, which makes CES much more reliable --- examples/EKIRace/calibrate_l63.jl | 62 +++++++++++---------------- examples/EKIRace/calibrate_l96.jl | 51 +++++++++++----------- examples/EKIRace/experiment_config.jl | 34 +++++++++------ 3 files changed, 69 insertions(+), 78 deletions(-) diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index a4f416710..93acf6de6 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -26,14 +26,13 @@ save_all_ekp = true cfg = experiment_config(:l63) N_ens_sizes = cfg.N_ens_sizes N_iter = cfg.N_iter -target_rmse = cfg.target_rmse n_repeats = cfg.n_repeats +terminate_at = cfg.terminate_at rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running Lorenz 63 problem" @info "Maximum number of EKI iterations: $N_iter" -@info "RMSE target: $target_rmse" configuration = - Dict("N_iter" => N_iter, "N_ens_sizes" => N_ens_sizes, "target_rmse" => target_rmse, "rng_seeds" => rng_seeds) + Dict("N_iter" => N_iter, "N_ens_sizes" => N_ens_sizes, "terminate_at"=> terminate_at, "rng_seeds" => rng_seeds) nx = 3 # dimensions of parameter vector nu = 2 @@ -118,8 +117,7 @@ ic_cov_sqrt = sqrt(ic_cov) ########################### Running EKI Race ########################### ######################################################################## -# Counters -conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how many iterations it takes to converge (per algorithm, per rand seed, per ense size) +conv_alg_iters = fill(NaN, 4, length(N_ens_sizes), length(rng_seeds)) final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) @@ -134,10 +132,10 @@ for (rr, rng_seed) in enumerate(rng_seeds) # initial parameters: N_params x N_ens initial_params = construct_initial_ensemble(rng, prior, N_ens) methods = [ - Inversion(prior), - TransformInversion(prior), + Inversion(), + TransformInversion(), GaussNewtonInversion(prior), - Unscented(prior; impose_prior = true), + Unscented(prior), ] @info "Ensemble size: $(N_ens)" @@ -150,9 +148,8 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, -# accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(0.1), + scheduler = DataMisfitController(terminate_at = terminate_at), ) else ekpobj = EKP.EnsembleKalmanProcess( @@ -162,18 +159,18 @@ for (rr, rng_seed) in enumerate(rng_seeds) deepcopy(method); rng = copy(rng), verbose = verbose_flag, -# accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DefaultScheduler(0.1), + scheduler = DataMisfitController(terminate_at = terminate_at), ) end Ne = get_N_ens(ekpobj) - count = 0 + ens_mean_final = zeros(nu) + G_ens_mean_final = zeros(ny) for i in 1:N_iter params_i = get_ϕ_final(prior, ekpobj) - # Calculating RMSE_e + # Calculating RMSE_e (diagnostic; not used as stopping criterion) ens_mean = mean(params_i, dims = 2)[:] # in constrained_space G_ens_mean = lorenz_forward( EnsembleMemberConfig(ens_mean), @@ -183,15 +180,10 @@ for (rr, rng_seed) in enumerate(rng_seeds) ) RMSE_e = norm(R_inv_var * (y - G_ens_mean[:])) / sqrt(size(y, 1)) @info "RMSE (at G(u_mean)): $(RMSE_e)" - # Convergence criteria - if RMSE_e < target_rmse - conv_alg_iters[kk, ee, rr] = count * Ne - final_parameters[kk, ee, rr, :] = ens_mean - final_model_output[kk, ee, rr, :] = G_ens_mean - break - end - # If RMSE convergence criteria is not satisfied + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] + G_ens = hcat( [ lorenz_forward( @@ -202,22 +194,18 @@ for (rr, rng_seed) in enumerate(rng_seeds) ) for j in 1:Ne ]..., ) - # Update - EKP.update_ensemble!(ekpobj, G_ens) - count = count + 1 - - # # Calculate RMSE_f to the mean G(u) if desired - # RMSE_f = sqrt(get_error_metrics(ekpobj)["loss"][end]) # equivalently - # @info "RMSE (at mean(G(u)): $(RMSE_f)" - # # Convergence criteria - # if RMSE_f < target_rmse - # conv_alg_iters[kk, ee, rr] = count * Ne - # final_parameters[kk, ee, rr, :] = ens_mean - # final_model_output[kk, ee, rr, :] = G_ens_mean - # break - # end + terminated = EKP.update_ensemble!(ekpobj, G_ens) + if !isnothing(terminated) + conv_alg_iters[kk, ee, rr] = i * Ne + break + end end - + final_parameters[kk, ee, rr, :] = ens_mean_final + final_model_output[kk, ee, rr, :] = G_ens_mean_final + if isnan(conv_alg_iters[kk, ee, rr]) # if didnt terminate + conv_alg_iters[kk, ee, rr] = N_iter * Ne + end + final_ensemble = get_ϕ_final(prior, ekpobj) # save ekp files diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index d741dcc2f..fa957df62 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -31,19 +31,18 @@ cfg = experiment_config(EXPERIMENT) case = cfg.force_case # "const-force", "vec-force", or "flux-force" N_ens_sizes = cfg.N_ens_sizes N_iter = cfg.N_iter -target_rmse = cfg.target_rmse n_repeats = cfg.n_repeats +terminate_at = cfg.terminate_at rng_seeds = randperm(1_000_000)[1:n_repeats] # list of random seeds @info "Running $case case" @info "Maximum number of EKI iterations: $N_iter" -@info "RMSE target: $target_rmse" # saved and loaded for plotting etc. configuration = Dict( "case" => case, "N_iter" => N_iter, "N_ens_sizes" => N_ens_sizes, - "target_rmse" => target_rmse, "rng_seeds" => rng_seeds, + "terminate_at" => terminate_at ) if case == "const-force" @@ -224,8 +223,7 @@ end ########################### Running EKI Race ########################### ######################################################################## -# Counters -conv_alg_iters = zeros(4, length(N_ens_sizes), length(rng_seeds)) #count how many iterations it takes to converge (per algorithm, per rand seed, per ense size) +conv_alg_iters = fill(NaN, 4, length(N_ens_sizes), length(rng_seeds)) final_parameters = zeros(4, length(N_ens_sizes), length(rng_seeds), nu) final_model_output = zeros(4, length(N_ens_sizes), length(rng_seeds), ny) @@ -239,10 +237,10 @@ for (rr, rng_seed) in enumerate(rng_seeds) # initial parameters: N_params x N_ens initial_params = construct_initial_ensemble(rng, prior, N_ens) methods = [ - Inversion(prior), - TransformInversion(prior), + Inversion(), + TransformInversion(), GaussNewtonInversion(prior), - Unscented(prior; impose_prior = true), + Unscented(prior), ] @info "Ensemble size: $(N_ens)" @@ -259,7 +257,7 @@ for (rr, rng_seed) in enumerate(rng_seeds) # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), #scheduler = DefaultScheduler(0.1), - scheduler = DataMisfitController(terminate_at=100), + scheduler = DataMisfitController(terminate_at = terminate_at), ) else ekpobj = EKP.EnsembleKalmanProcess( @@ -271,19 +269,19 @@ for (rr, rng_seed) in enumerate(rng_seeds) verbose = verbose_flag, # accelerator = DefaultAccelerator(), localization_method = NoLocalization(), - scheduler = DataMisfitController(terminate_at=100), + scheduler = DataMisfitController(terminate_at = terminate_at), # scheduler = DefaultScheduler(0.1), ) end Ne = get_N_ens(ekpobj) - count = 0 + ens_mean_final = zeros(nu) + G_ens_mean_final = zeros(ny) for i in 1:N_iter params_i = get_ϕ_final(prior, ekpobj) - # Calculating RMSE_e + # Calculating RMSE_e (diagnostic; not used as stopping criterion) ens_mean = mean(params_i, dims = 2)[:] - # forcing = build_forcing(ens_mean, phi_structure) forcing = build_forcing(phi, ens_mean, phi_structure, sample_range) G_ens_mean = lorenz_forward( forcing, @@ -293,15 +291,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) ) RMSE_e = norm(R_inv_var * (y - G_ens_mean[:])) / sqrt(size(y, 1)) @info "RMSE (at G(u_mean)): $(RMSE_e)" - # Convergence criteria - if RMSE_e < target_rmse - conv_alg_iters[kk, ee, rr] = count * Ne - final_parameters[kk, ee, rr, :] = ens_mean - final_model_output[kk, ee, rr, :] = G_ens_mean - break - end + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] - # If RMSE convergence criteria is not satisfied G_ens = hcat( [ lorenz_forward( @@ -312,13 +304,18 @@ for (rr, rng_seed) in enumerate(rng_seeds) ) for j in 1:Ne ]..., ) - # Update - EKP.update_ensemble!(ekpobj, G_ens) - count = count + 1 - - - + terminated = EKP.update_ensemble!(ekpobj, G_ens) + if !isnothing(terminated) + conv_alg_iters[kk, ee, rr] = i * Ne + break + end + end + final_parameters[kk, ee, rr, :] = ens_mean_final + final_model_output[kk, ee, rr, :] = G_ens_mean_final + if isnan(conv_alg_iters[kk, ee, rr]) # if didnt terminate + conv_alg_iters[kk, ee, rr] = N_iter * Ne end + final_ensemble = get_ϕ_final(prior, ekpobj) # save ekp files diff --git a/examples/EKIRace/experiment_config.jl b/examples/EKIRace/experiment_config.jl index 483dff27f..2018238bd 100644 --- a/examples/EKIRace/experiment_config.jl +++ b/examples/EKIRace/experiment_config.jl @@ -4,32 +4,32 @@ using Dates ############### USER TOGGLE ######################################### ######################################################################## # Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux -const EXPERIMENT = :l63 +EXPERIMENT = :l63 # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). -const calibrate_date = Date("2026-05-28", "yyyy-mm-dd") +calibrate_date = Date("2026-05-28", "yyyy-mm-dd") ######################################################################## ############### SHARED CONSTANTS #################################### ######################################################################## -const method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] +method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] -const method_names = [ - ("Inversion(prior)", "TEKI"), - ("TransformInversion(prior)", "ETKI"), - ("GaussNewtonInversion(prior)", "GNKI"), - ("Unscented(prior; impose_prior=true)", "UKI"), +method_names = [ + ("Inversion()", "EKI"), + ("TransformInversion()", "ETKI"), + ("GaussNewtonInversion()", "GNKI"), + ("Unscented(prior)", "UKI"), ] -const method_cases_key = Dict( +method_cases_key = Dict( "Inversion" => "ces-eki-dmc", "TransformInversion" => "ces-etki-dmc", "Unscented" => "ces-uki-dmc", "GaussNewtonInversion" => "ces-iekf", ) -const forcing_cases_key = Dict( +forcing_cases_key = Dict( "const-force" => "ensemble_results", "vec-force" => "spatial_forcing_ensemble_results", "flux-force" => "nn_forcing_ensemble_results", @@ -38,6 +38,12 @@ const forcing_cases_key = Dict( ######################################################################## ############### PER-CASE CONFIG ##################################### ######################################################################## +# Important dials: +# terminate_at: psuedotime to terminate. (T=1 is approx posterior for the finite-time method variants) +# N_iter: otherwise, maximum iterations +# N_ens_sizes: Vector of experiments +# n_repeats: number of rng seeds + function experiment_config(case::Symbol) if case == :l63 return ( @@ -45,7 +51,7 @@ function experiment_config(case::Symbol) force_case = nothing, N_ens_sizes = [10, 25, 40], N_iter = 20, - target_rmse = 1.0, + terminate_at = 2.0, # DataMisfitController end time n_repeats = 20, max_iter = 10, retain_var = 0.99, @@ -59,7 +65,7 @@ function experiment_config(case::Symbol) force_case = "const-force", N_ens_sizes = [5, 15, 30], N_iter = 20, - target_rmse = 1.0, + terminate_at = 2.0, n_repeats = 20, max_iter = 15, retain_var = 0.95, @@ -73,7 +79,7 @@ function experiment_config(case::Symbol) force_case = "vec-force", N_ens_sizes = [50, 75, 100], N_iter = 20, - target_rmse = 1.0, + terminate_at = 2.0, n_repeats = 20, max_iter = 15, retain_var = 0.95, @@ -87,7 +93,7 @@ function experiment_config(case::Symbol) force_case = "flux-force", N_ens_sizes = [50, 75, 100], N_iter = 20, - target_rmse = 1.0, + terminate_at = 2.0, n_repeats = 20, max_iter = 15, retain_var = 0.95, From f368cccdd2e8bcc89332ceb3b0bd87cb86bcbfc8 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 13:31:24 -0700 Subject: [PATCH 17/86] running ES for loops over (1:k) k=1,...,max_iter, so get monotone posteriors for each series --- examples/EKIRace/emulate_sample_l63.jl | 241 +++++++---------- examples/EKIRace/emulate_sample_l96.jl | 243 +++++++----------- examples/EKIRace/experiment_config.jl | 5 +- .../l63_exp_to_leaderboard_utilities.jl | 134 +++++----- .../l96_exp_to_leaderboard_utilities.jl | 134 +++++----- 5 files changed, 306 insertions(+), 451 deletions(-) diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl index 73dbe87e5..f423d86db 100644 --- a/examples/EKIRace/emulate_sample_l63.jl +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -38,19 +38,9 @@ function main() for rng_idx in rng_idxs calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) @info("Perform emulate-sample for L63: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - min_iter = 1 - skip_iter = 1 - max_iter = cfg.max_iter - - #### - - # need to loop over - rng_seed = 44011 - rng = Random.MersenneTwister(rng_seed) - + # loading relevant data homedir = pwd() - println(homedir) figure_save_directory = joinpath(homedir, "output/", calib_dir) data_save_directory = joinpath(homedir, "output", calib_dir) loaded_calib_files = [ @@ -59,159 +49,106 @@ function main() joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)), ] if !isfile(loaded_calib_files[1]) - throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l63.jl")) end - + ekpobj = load(loaded_calib_files[1])["ekpobj"] priors = load(loaded_calib_files[2])["prior"] truth_sample = load(loaded_calib_files[3])["y"] - truth_params_obj = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space - truth_params_constrained = truth_params_obj.u # the parameters here have been exponentiated not using our in-built constraints. + truth_params_obj = load(loaded_calib_files[3])["truth_params_structure"] + truth_params_constrained = truth_params_obj.u truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) Γy = get_obs_noise_cov(ekpobj) - - n_params = length(truth_params) # "input dim" - output_dim = size(Γy, 1) - ### - ### Emulate: - ### - # choice of machine-learning tool in the emulation stage + n_params = length(truth_params) + + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + retain_var = cfg.retain_var nugget = 1e-6 - overrides = Dict( - # "verbose" => true, - "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => 5.0, - "n_iteration" => 10, - "n_features_opt" => cfg.n_features_opt, - ) - n_features = cfg.n_features - kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) - mlt = ScalarRandomFeatureInterface( - n_features, - n_params, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - - - # Get training points from the EKP iteration number in the second input term - iter_end = length(get_u(ekpobj))-1 - N_iter = min(max_iter, iter_end) # number of paired iterations taken from EKP - min_iter = min(max_iter, max(1, min_iter)) - if N_iter < 1 - @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") + + iter_end = length(get_u(ekpobj)) - 1 + K = min(iter_end - 1, cfg.max_iter) + if K < 1 + @error("Not enough EKP iterations for emulate-sample loop. Skipping $(calib_filename_suffix)...") continue - elseif N_iter == 1 - @warn("Convergence in only 1 iteration, removing train-test split based on iterations") - @info "Train iterations: $(min_iter:skip_iter:N_iter)" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) - else - @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" - @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:iter_end) # "next" iterations - end - - # data processing configuration - retain_var = cfg.retain_var - encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] - encoder_kwargs = encoder_kwargs_from(ekpobj, priors) - - emulator = Emulator( - mlt, - input_output_pairs; - encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, - ) - optimize_hyperparameters!(emulator) - - # Check how well the Gaussian Process regression predicts on the - # true parameters - y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) - if !(N_iter<2) - y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) - end - println("ML prediction on true parameters: ") - println(vec(y_mean)) - println("true data: ") - println(truth_sample) # what was used as truth - println(" ML predicted standard deviation") - println(sqrt.(diag(y_var[1], 0))) - println("ML weighted-MSE (truth): ") - diff = (truth_sample - vec(y_mean)) - Γypluscov = Symmetric(mean(y_var) + Γy) - println(mean(diff' * (Γypluscov\ diff))) - if !(N_iter<2) - println("ML weighted-MSE (test ensemble): ") - difftest = get_outputs(input_output_pairs_test) - y_mean_test - println(mean(difftest' * (Γypluscov \ difftest))) - end - - #end - ### - ### Sample: Markov Chain Monte Carlo - ### - # initial values - u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) - println("initial parameters: ", u0) - - # First let's run a short chain to determine a good step size - mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) - new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) - - # Now begin the actual MCMC - println("Begin MCMC - with step size ", new_step) - chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) - - posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) - - post_mean = mean(posterior) - post_cov = cov(posterior) - println("post_mean") - println(post_mean) - println("post_cov") - println(post_cov) - println("D util") - println(det(inv(post_cov))) - println(" ") - - param_names = get_name(posterior) - - posterior_samples = reduce(vcat,[get_distribution(posterior)[name] for name in get_name(posterior)]) #samples are columns - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - - # marginal histogram - pp = plot(priors, c = :gray) - plot!(pp, posterior) - final_params_constrained = get_ϕ_mean_final(priors, ekpobj) - for (i, sp) in enumerate(pp.subplots) - vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) - vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + @info "Running emulate-sample for k = 1:$(K) training iterations" + + posteriors_by_k = Dict{Int, Any}() + iop_by_k = Dict{Int, Any}() + encoder_schedule_by_k = Dict{Int, Any}() + + for k in 1:K + rng = Random.MersenneTwister(44011) + input_output_pairs = Utilities.get_training_points(ekpobj, 1:k) + @info "k = $(k): training on $(N_ens * k) samples (iterations 1:$(k))" + + ### Emulate + overrides = Dict( + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 5.0, + "n_iteration" => 10, + "n_features_opt" => cfg.n_features_opt, + ) + mlt = ScalarRandomFeatureInterface( + cfg.n_features, + n_params, + rng = rng, + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), + optimizer_options = overrides, + ) + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # diagnostics + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + diff = truth_sample - vec(y_mean) + Γypluscov = Symmetric(mean(y_var) + Γy) + println("k=$(k) | ML weighted-MSE (truth): ", mean(diff' * (Γypluscov \ diff))) + + ### Sample: MCMC + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + println("k=$(k) | Begin MCMC - step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + println("k=$(k) | post_mean: ", mean(posterior)) + println("k=$(k) | D util: ", det(inv(cov(posterior)))) + + # figure + pp = plot(priors, c = :gray) + plot!(pp, posterior) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + + posteriors_by_k[k] = posterior + iop_by_k[k] = input_output_pairs + encoder_schedule_by_k[k] = get_encoder_schedule(emulator) end - figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)") - savefig(figpath * ".png") - savefig(figpath * ".pdf") - - # Save data + + # one file per EKI experiment, containing all k posteriors save( joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), - "posterior", - posterior, - "priors", - priors, - "truth_params_constrained", - truth_params_constrained, - "final_params_constrained", - final_params_constrained, - "input_output_pairs", - input_output_pairs, - "truth_params", - truth_params, - "encoder_schedule", - get_encoder_schedule(emulator), + "posteriors_by_k", posteriors_by_k, + "input_output_pairs_by_k", iop_by_k, + "encoder_schedules_by_k", encoder_schedule_by_k, + "priors", priors, + "truth_params_constrained", truth_params_constrained, + "final_params_constrained", final_params_constrained, + "truth_params", truth_params, + "k_values", collect(1:K), ) end end diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index 8c8694e7e..9f9e6bf8b 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -41,19 +41,9 @@ function main() for rng_idx in rng_idxs calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) @info("Perform emulate-sample for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - min_iter = 1 - skip_iter = 1 - max_iter = cfg.max_iter - - #### - - # need to loop over - rng_seed = 44011 - rng = Random.MersenneTwister(rng_seed) - + # loading relevant data homedir = pwd() - println(homedir) figure_save_directory = joinpath(homedir, "output/", calib_dir) data_save_directory = joinpath(homedir, "output", calib_dir) loaded_calib_files = [ @@ -64,14 +54,14 @@ function main() if !isfile(loaded_calib_files[1]) throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) end - + ekpobj = load(loaded_calib_files[1])["ekpobj"] priors = load(loaded_calib_files[2])["prior"] truth_sample = load(loaded_calib_files[3])["y"] - (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] #true parameters in constrained space - + (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] + truth_params_constrained = if force_case == "const-force" - [truth_params_obj.val] #1d + [truth_params_obj.val] elseif force_case == "vec-force" truth_params_obj.val elseif force_case == "flux-force" @@ -81,154 +71,97 @@ function main() end truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) Γy = get_obs_noise_cov(ekpobj) - - n_params = length(truth_params) # "input dim" - output_dim = size(Γy, 1) - ### - ### Emulate: - ### - # choice of machine-learning tool in the emulation stage + n_params = length(truth_params) + + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + retain_var = cfg.retain_var nugget = 1e-6 - if N_ens * length(get_u(ekpobj)) < 200 - csm = 5.0 - else - csm = 1.0 - end - overrides = Dict( -# "verbose" => true, - "scheduler" => DataMisfitController(terminate_at = 1000.0), - "cov_sample_multiplier" => csm, - "n_iteration" => 10, - "n_features_opt" => cfg.n_features_opt, - ) - n_features = cfg.n_features - kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()) - mlt = ScalarRandomFeatureInterface( - n_features, - n_params, - rng = rng, - kernel_structure = kernel_structure, - optimizer_options = overrides, - ) - - - # Get training points from the EKP iteration number in the second input term - iter_end = length(get_u(ekpobj))-1 - N_iter = min(max_iter, iter_end) # number of paired iterations taken from EKP - min_iter = min(max_iter, max(1, min_iter)) - if N_iter < 1 - @error("No iterations found in ekpobj, perhaps a calibration issue? Skipping case $(calib_filename_suffix)...") + + iter_end = length(get_u(ekpobj)) - 1 + K = min(iter_end - 1, cfg.max_iter) + if K < 1 + @error("Not enough EKP iterations for emulate-sample loop. Skipping $(calib_filename_suffix)...") continue - elseif N_iter == 1 - @warn("Convergence in only 1 iteration, removing train-test split based on iterations") - @info "Train iterations: $(min_iter:skip_iter:N_iter)" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:N_iter) - else - @info "Train iterations: $(min_iter:skip_iter:(N_iter-1))" - @info "Test iterations: $(N_iter:(length(get_u(ekpobj)) - 1))" - input_output_pairs = Utilities.get_training_points(ekpobj, min_iter:skip_iter:(N_iter - 1)) - input_output_pairs_test = Utilities.get_training_points(ekpobj, N_iter:iter_end) # "next" iteration - - end - - # data processing configuration - retain_var = cfg.retain_var - encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] - encoder_kwargs = encoder_kwargs_from(ekpobj, priors) - - emulator = Emulator( - mlt, - input_output_pairs; - encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, - ) - optimize_hyperparameters!(emulator) - - # Check how well the Gaussian Process regression predicts on the - # true parameters - y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) - if !(N_iter<2) - y_mean_test, y_var_test = Emulators.predict(emulator, get_inputs(input_output_pairs_test)) - end - println("ML prediction on true parameters: ") - println(vec(y_mean)) - println("true data: ") - println(truth_sample) # what was used as truth - println(" ML predicted standard deviation") - println(sqrt.(diag(y_var[1], 0))) - println("ML weighted-MSE (truth): ") - diff = (truth_sample - vec(y_mean)) - Γypluscov = Symmetric(mean(y_var) + Γy) - println(mean(diff' * (Γypluscov\ diff))) - if !(N_iter<2) - println("ML weighted-MSE (test ensemble): ") - difftest = get_outputs(input_output_pairs_test) - y_mean_test - println(mean(difftest' * (Γypluscov \ difftest))) end + @info "Running emulate-sample for k = 1:$(K) training iterations" - #end - ### - ### Sample: Markov Chain Monte Carlo - ### - # initial values - u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) - println("initial parameters: ", u0) - - # First let's run a short chain to determine a good step size - mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) - new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) - - # Now begin the actual MCMC - println("Begin MCMC - with step size ", new_step) - chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) - - posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) - - post_mean = mean(posterior) - post_cov = cov(posterior) - println("post_mean") - println(post_mean) - println("post_cov") - println(post_cov) - println("D util") - println(det(inv(post_cov))) - println(" ") - - param_names = get_name(posterior) - - posterior_samples = reduce(vcat,[get_distribution(posterior)[name] for name in get_name(posterior)]) #samples are columns - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - - # marginal histogram - pp = plot(priors, c = :gray) - plot!(pp, posterior) - final_params_constrained = get_ϕ_mean_final(priors, ekpobj) - for (i, sp) in enumerate(pp.subplots) - vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) - vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + posteriors_by_k = Dict{Int, Any}() + iop_by_k = Dict{Int, Any}() + encoder_schedule_by_k = Dict{Int, Any}() + + for k in 1:K + rng = Random.MersenneTwister(44011) + input_output_pairs = Utilities.get_training_points(ekpobj, 1:k) + @info "k = $(k): training on $(N_ens * k) samples (iterations 1:$(k))" + + ### Emulate + csm = N_ens * k < 200 ? 5.0 : 1.0 + overrides = Dict( + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => csm, + "n_iteration" => 10, + "n_features_opt" => cfg.n_features_opt, + ) + mlt = ScalarRandomFeatureInterface( + cfg.n_features, + n_params, + rng = rng, + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), + optimizer_options = overrides, + ) + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # diagnostics + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + diff = truth_sample - vec(y_mean) + Γypluscov = Symmetric(mean(y_var) + Γy) + println("k=$(k) | ML weighted-MSE (truth): ", mean(diff' * (Γypluscov \ diff))) + + ### Sample: MCMC + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + println("k=$(k) | Begin MCMC - step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + println("k=$(k) | post_mean: ", mean(posterior)) + println("k=$(k) | D util: ", det(inv(cov(posterior)))) + + # figure + pp = plot(priors, c = :gray) + plot!(pp, posterior) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + + posteriors_by_k[k] = posterior + iop_by_k[k] = input_output_pairs + encoder_schedule_by_k[k] = get_encoder_schedule(emulator) end - figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)") - savefig(figpath * ".png") - savefig(figpath * ".pdf") - # Save data + # one file per EKI experiment, containing all k posteriors save( joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), - "posterior", - posterior, - "priors", - priors, - "truth_params_constrained", - truth_params_constrained, - "final_params_constrained", - final_params_constrained, - "input_output_pairs", - input_output_pairs, - "truth_params", - truth_params, - "encoder_schedule", - get_encoder_schedule(emulator), + "posteriors_by_k", posteriors_by_k, + "input_output_pairs_by_k", iop_by_k, + "encoder_schedules_by_k", encoder_schedule_by_k, + "priors", priors, + "truth_params_constrained", truth_params_constrained, + "final_params_constrained", final_params_constrained, + "truth_params", truth_params, + "k_values", collect(1:K), ) end end diff --git a/examples/EKIRace/experiment_config.jl b/examples/EKIRace/experiment_config.jl index 2018238bd..e6edba6e2 100644 --- a/examples/EKIRace/experiment_config.jl +++ b/examples/EKIRace/experiment_config.jl @@ -4,11 +4,12 @@ using Dates ############### USER TOGGLE ######################################### ######################################################################## # Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux -EXPERIMENT = :l63 +EXPERIMENT = :l96_const # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). -calibrate_date = Date("2026-05-28", "yyyy-mm-dd") +#calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +calibrate_date = today() ######################################################################## ############### SHARED CONSTANTS #################################### diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index 6b05fceb1..e4c1e4d23 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -46,66 +46,65 @@ end ### Load data -# Load first valid file to determine parameter dimension +# Load first valid file to determine parameter dimension and max k first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posterior"]))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) targets = [1.0] +n_k = cfg.max_iter # maximum number of training-iteration slices n_rng = length(rng_idxs) n_ens = length(N_enss) n_targets = length(targets) -# Pre-allocate storage arrays indexed as (random_seed, ensemble_size, rmse_target, algorithm_type, ...) -true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params, n_params) -mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +# Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not +true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) - - post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) - truth_params = loaded["truth_params"] # truth in unconstrained space - post_name = get_name(post_dist) - - pm = vec(mean(post_dist)) # posterior mean, shape (n_params,) - pc = cov(post_dist) # posterior covariance, shape (n_params, n_params) - - C_reg = Symmetric(pc + 1e-10 * I) - post_normal = MvNormal(pm, C_reg) - post_samples = reduce(vcat,[get_distribution(post_dist)[name] for name in get_name(post_dist)]) #samples are columns - - pmode_idx = argmax(logpdf(post_normal, post_samples)) - pmode = post_samples[:, pmode_idx] - - - # Load ekpobj to compute calibration evaluation count - ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) - conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps - - # Array indices + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + # Load ekpobj for total calibration cost + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - diff = pm - truth_params - for l in 1:n_targets - true_param_arr[i, j, l, 1, :] = truth_params - post_mean_arr[i, j, l, 1, :] = pm - post_cov_arr[i, j, l, 1, :, :] = pc - # (m - truth)' C^{-1} (m - truth) - mahal_arr[i, j, l, 1] = diff' * (C_reg \ diff) - # log p(truth | posterior) - log p(mode | posterior) - logpdf_true_v_map_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # total forward model evaluations to reach target - n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens + true_param_arr[i, j, l, 1, :] = truth_params + n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens end + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + for l in 1:n_targets + post_mean_arr[i, j, l, 1, k, :] = pm + post_cov_arr[i, j, l, 1, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, l, 1, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, l, 1, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + end + end end @@ -117,8 +116,9 @@ defDim(ds, "random_seed", n_rng) defDim(ds, "ensemble_size", n_ens) defDim(ds, "rmse_target", n_targets) defDim(ds, "algorithm_type", 1) +defDim(ds, "k_iter", n_k) defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) # second param axis for covariance matrix +defDim(ds, "param_dim_2", n_params) # Coordinate variables alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) @@ -135,22 +135,22 @@ rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" rmse_var[:] = targets +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + # Data variables n_evals_v = defVar( - ds, - "n_evals_to_target", - Float64, + ds, "n_evals_to_target", Float64, ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); fillvalue = NaN, ) -n_evals_v.attrib["description"] = "Number of forward model evaluations (i.e., algorithm cost): conv_alg_iters * ensemble_size. Stored as float in case of NaN values." +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." n_evals_v[:, :, :, :] = n_evals_arr true_param_v = defVar( - ds, - "true_param", - Float64, + ds, "true_param", Float64, ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); fillvalue = NaN, ) @@ -158,44 +158,36 @@ true_param_v.attrib["description"] = "truth_param" true_param_v[:, :, :, :, :] = true_param_arr post_mean_v = defVar( - ds, - "post_mean", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + ds, "post_mean", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim"); fillvalue = NaN, ) -post_mean_v.attrib["description"] = "mean(posterior)" -post_mean_v[:, :, :, :, :] = post_mean_arr +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :, :, :] = post_mean_arr post_cov_v = defVar( - ds, - "post_cov", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim", "param_dim_2"); + ds, "post_cov", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim", "param_dim_2"); fillvalue = NaN, ) -post_cov_v.attrib["description"] = "cov(posterior)" -post_cov_v[:, :, :, :, :, :] = post_cov_arr +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :, :, :] = post_cov_arr mahalanobis_v = defVar( - ds, - "mahalanobis", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ds, "mahalanobis", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); fillvalue = NaN, ) mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" -mahalanobis_v[:, :, :, :] = mahal_arr +mahalanobis_v[:, :, :, :, :] = mahal_arr posterior_logpdf_true_v_map_v = defVar( - ds, - "posterior_logpdf_true_v_map", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ds, "posterior_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); fillvalue = NaN, ) posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" -posterior_logpdf_true_v_map_v[:, :, :, :] = logpdf_true_v_map_arr +posterior_logpdf_true_v_map_v[:, :, :, :, :] = logpdf_true_v_map_arr close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index a02880f4c..d5fac44fc 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -48,66 +48,65 @@ end ### Load data -# Load first valid file to determine parameter dimension +# Load first valid file to determine parameter dimension and max k first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posterior"]))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) targets = [1.0] +n_k = cfg.max_iter # maximum number of training-iteration slices n_rng = length(rng_idxs) n_ens = length(N_enss) n_targets = length(targets) -# Pre-allocate storage arrays indexed as (random_seed, ensemble_size, rmse_target, algorithm_type, ...) -true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params, n_params) -mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +# Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not +true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) +post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) - - post_dist = loaded["posterior"] # ParameterDistribution from MCMC chain (unconstrained space) - truth_params = loaded["truth_params"] # truth in unconstrained space - post_name = get_name(post_dist) - - pm = vec(mean(post_dist)) # posterior mean, shape (n_params,) - pc = cov(post_dist) # posterior covariance, shape (n_params, n_params) - - C_reg = Symmetric(pc + 1e-10 * I) - post_normal = MvNormal(pm, C_reg) - post_samples = reduce(vcat,[get_distribution(post_dist)[name] for name in get_name(post_dist)]) #samples are columns - - pmode_idx = argmax(logpdf(post_normal, post_samples)) - pmode = post_samples[:, pmode_idx] - - - # Load ekpobj to compute calibration evaluation count - ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) - conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) # number of completed EKP update steps - - # Array indices + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + # Load ekpobj for total calibration cost + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - diff = pm - truth_params - for l in 1:n_targets - true_param_arr[i, j, l, 1, :] = truth_params - post_mean_arr[i, j, l, 1, :] = pm - post_cov_arr[i, j, l, 1, :, :] = pc - # (m - truth)' C^{-1} (m - truth) - mahal_arr[i, j, l, 1] = diff' * (C_reg \ diff) - # log p(truth | posterior) - log p(mode | posterior) - logpdf_true_v_map_arr[i, j, l, 1] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # total forward model evaluations to reach target - n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens + true_param_arr[i, j, l, 1, :] = truth_params + n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens end + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + for l in 1:n_targets + post_mean_arr[i, j, l, 1, k, :] = pm + post_cov_arr[i, j, l, 1, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, l, 1, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, l, 1, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + end + end end @@ -119,8 +118,9 @@ defDim(ds, "random_seed", n_rng) defDim(ds, "ensemble_size", n_ens) defDim(ds, "rmse_target", n_targets) defDim(ds, "algorithm_type", 1) +defDim(ds, "k_iter", n_k) defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) # second param axis for covariance matrix +defDim(ds, "param_dim_2", n_params) # Coordinate variables alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) @@ -137,22 +137,22 @@ rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" rmse_var[:] = targets +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + # Data variables n_evals_v = defVar( - ds, - "n_evals_to_target", - Float64, + ds, "n_evals_to_target", Float64, ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); fillvalue = NaN, ) -n_evals_v.attrib["description"] = "Number of forward model evaluations (i.e., algorithm cost): conv_alg_iters * ensemble_size. Stored as float in case of NaN values." +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." n_evals_v[:, :, :, :] = n_evals_arr true_param_v = defVar( - ds, - "true_param", - Float64, + ds, "true_param", Float64, ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); fillvalue = NaN, ) @@ -160,44 +160,36 @@ true_param_v.attrib["description"] = "truth_param" true_param_v[:, :, :, :, :] = true_param_arr post_mean_v = defVar( - ds, - "post_mean", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + ds, "post_mean", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim"); fillvalue = NaN, ) -post_mean_v.attrib["description"] = "mean(posterior)" -post_mean_v[:, :, :, :, :] = post_mean_arr +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :, :, :] = post_mean_arr post_cov_v = defVar( - ds, - "post_cov", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim", "param_dim_2"); + ds, "post_cov", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim", "param_dim_2"); fillvalue = NaN, ) -post_cov_v.attrib["description"] = "cov(posterior)" -post_cov_v[:, :, :, :, :, :] = post_cov_arr +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :, :, :] = post_cov_arr mahalanobis_v = defVar( - ds, - "mahalanobis", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ds, "mahalanobis", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); fillvalue = NaN, ) mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" -mahalanobis_v[:, :, :, :] = mahal_arr +mahalanobis_v[:, :, :, :, :] = mahal_arr posterior_logpdf_true_v_map_v = defVar( - ds, - "posterior_logpdf_true_v_map", - Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ds, "posterior_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); fillvalue = NaN, ) posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" -posterior_logpdf_true_v_map_v[:, :, :, :] = logpdf_true_v_map_arr +posterior_logpdf_true_v_map_v[:, :, :, :, :] = logpdf_true_v_map_arr close(ds) @info "Saved leaderboard data to $(nc_save_filename)" From 8fa3c9aedff9e83d90f9af0e6ac9e3f376e423d4 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 14:52:31 -0700 Subject: [PATCH 18/86] rm targets and alg types in nc --- .../l63_exp_to_leaderboard_utilities.jl | 38 ++++++++----------- .../l96_exp_to_leaderboard_utilities.jl | 38 ++++++++----------- 2 files changed, 32 insertions(+), 44 deletions(-) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index e4c1e4d23..71bf731f0 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -50,20 +50,18 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -targets = [1.0] n_k = cfg.max_iter # maximum number of training-iteration slices -n_rng = length(rng_idxs) -n_ens = length(N_enss) -n_targets = length(targets) +n_rng = length(rng_idxs) +n_ens = length(N_enss) # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not -true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params) -post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params, n_params) -mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) -logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -81,10 +79,8 @@ for (N_ens, rng_idx) in valid_file_items i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - for l in 1:n_targets - true_param_arr[i, j, l, 1, :] = truth_params - n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens - end + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens for k in k_values post_dist = posteriors_by_k[k] @@ -96,14 +92,12 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params - for l in 1:n_targets - post_mean_arr[i, j, l, 1, k, :] = pm - post_cov_arr[i, j, l, 1, k, :, :] = pc - # (m - truth)' C^{-1} (m - truth) - mahal_arr[i, j, l, 1, k] = diff' * (C_reg \ diff) - # log p(truth | posterior) - log p(mode | posterior) - logpdf_true_v_map_arr[i, j, l, 1, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - end + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) end end diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index d5fac44fc..150c9d5f2 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -52,20 +52,18 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -targets = [1.0] n_k = cfg.max_iter # maximum number of training-iteration slices -n_rng = length(rng_idxs) -n_ens = length(N_enss) -n_targets = length(targets) +n_rng = length(rng_idxs) +n_ens = length(N_enss) # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not -true_param_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_params) -n_evals_arr = fill(NaN, n_rng, n_ens, n_targets, 1) -post_mean_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params) -post_cov_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k, n_params, n_params) -mahal_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) -logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_targets, 1, n_k) +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -83,10 +81,8 @@ for (N_ens, rng_idx) in valid_file_items i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - for l in 1:n_targets - true_param_arr[i, j, l, 1, :] = truth_params - n_evals_arr[i, j, l, 1] = conv_alg_iters * N_ens - end + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens for k in k_values post_dist = posteriors_by_k[k] @@ -98,14 +94,12 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params - for l in 1:n_targets - post_mean_arr[i, j, l, 1, k, :] = pm - post_cov_arr[i, j, l, 1, k, :, :] = pc - # (m - truth)' C^{-1} (m - truth) - mahal_arr[i, j, l, 1, k] = diff' * (C_reg \ diff) - # log p(truth | posterior) - log p(mode | posterior) - logpdf_true_v_map_arr[i, j, l, 1, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - end + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) end end From 6d38f021b7144ca9457a783a6f88c7911698b47a Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 14:53:05 -0700 Subject: [PATCH 19/86] rm targets and alg types in nc --- .../l63_exp_to_leaderboard_utilities.jl | 43 ++++++++----------- .../l96_exp_to_leaderboard_utilities.jl | 43 ++++++++----------- 2 files changed, 34 insertions(+), 52 deletions(-) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index 71bf731f0..3be38b873 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -106,18 +106,13 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "rmse_target", n_targets) -defDim(ds, "algorithm_type", 1) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) # Coordinate variables -alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) -alg_var[1] = method_cases_key[method] - rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) rng_var[:] = rng_idxs @@ -125,10 +120,6 @@ ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) ens_var.attrib["description"] = "Number of ensemble members" ens_var[:] = N_enss -rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) -rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" -rmse_var[:] = targets - k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) @@ -137,51 +128,51 @@ k_var[:] = collect(1:n_k) n_evals_v = defVar( ds, "n_evals_to_target", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ("random_seed", "ensemble_size"); fillvalue = NaN, ) n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." -n_evals_v[:, :, :, :] = n_evals_arr +n_evals_v[:, :] = n_evals_arr true_param_v = defVar( ds, "true_param", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + ("random_seed", "ensemble_size", "param_dim"); fillvalue = NaN, ) true_param_v.attrib["description"] = "truth_param" -true_param_v[:, :, :, :, :] = true_param_arr +true_param_v[:, :, :] = true_param_arr post_mean_v = defVar( ds, "post_mean", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim"); + ("random_seed", "ensemble_size", "k_iter", "param_dim"); fillvalue = NaN, ) post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" -post_mean_v[:, :, :, :, :, :] = post_mean_arr +post_mean_v[:, :, :, :] = post_mean_arr post_cov_v = defVar( ds, "post_cov", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim", "param_dim_2"); + ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); fillvalue = NaN, ) post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" -post_cov_v[:, :, :, :, :, :, :] = post_cov_arr +post_cov_v[:, :, :, :, :] = post_cov_arr mahalanobis_v = defVar( ds, "mahalanobis", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); + ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" -mahalanobis_v[:, :, :, :, :] = mahal_arr +mahalanobis_v[:, :, :] = mahal_arr posterior_logpdf_true_v_map_v = defVar( ds, "posterior_logpdf_true_v_map", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); + ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" -posterior_logpdf_true_v_map_v[:, :, :, :, :] = logpdf_true_v_map_arr +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 150c9d5f2..b875a5d2b 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -108,18 +108,13 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "rmse_target", n_targets) -defDim(ds, "algorithm_type", 1) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) # Coordinate variables -alg_var = defVar(ds, "algorithm_type", String, ("algorithm_type",)) -alg_var[1] = method_cases_key[method] - rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) rng_var[:] = rng_idxs @@ -127,10 +122,6 @@ ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) ens_var.attrib["description"] = "Number of ensemble members" ens_var[:] = N_enss -rmse_var = defVar(ds, "rmse_target", Float64, ("rmse_target",); fillvalue = NaN) -rmse_var.attrib["description"] = "Target accuracy level (root mean square error)" -rmse_var[:] = targets - k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) @@ -139,51 +130,51 @@ k_var[:] = collect(1:n_k) n_evals_v = defVar( ds, "n_evals_to_target", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type"); + ("random_seed", "ensemble_size"); fillvalue = NaN, ) n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." -n_evals_v[:, :, :, :] = n_evals_arr +n_evals_v[:, :] = n_evals_arr true_param_v = defVar( ds, "true_param", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "param_dim"); + ("random_seed", "ensemble_size", "param_dim"); fillvalue = NaN, ) true_param_v.attrib["description"] = "truth_param" -true_param_v[:, :, :, :, :] = true_param_arr +true_param_v[:, :, :] = true_param_arr post_mean_v = defVar( ds, "post_mean", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim"); + ("random_seed", "ensemble_size", "k_iter", "param_dim"); fillvalue = NaN, ) post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" -post_mean_v[:, :, :, :, :, :] = post_mean_arr +post_mean_v[:, :, :, :] = post_mean_arr post_cov_v = defVar( ds, "post_cov", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter", "param_dim", "param_dim_2"); + ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); fillvalue = NaN, ) post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" -post_cov_v[:, :, :, :, :, :, :] = post_cov_arr +post_cov_v[:, :, :, :, :] = post_cov_arr mahalanobis_v = defVar( ds, "mahalanobis", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); + ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" -mahalanobis_v[:, :, :, :, :] = mahal_arr +mahalanobis_v[:, :, :] = mahal_arr posterior_logpdf_true_v_map_v = defVar( ds, "posterior_logpdf_true_v_map", Float64, - ("random_seed", "ensemble_size", "rmse_target", "algorithm_type", "k_iter"); + ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" -posterior_logpdf_true_v_map_v[:, :, :, :, :] = logpdf_true_v_map_arr +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr close(ds) @info "Saved leaderboard data to $(nc_save_filename)" From 12f02ba460b265b591b286415bd9c687e994b405 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 15:47:38 -0700 Subject: [PATCH 20/86] compute metrics from the MH and logpdf metrics --- .../EKIRace/compute_leaderboard_metrics.jl | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 examples/EKIRace/compute_leaderboard_metrics.jl diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl new file mode 100644 index 000000000..b17711eac --- /dev/null +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -0,0 +1,32 @@ +using NCDatasets +using Distributions +using Statistics + +filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" + +# load +ncd = NCDataset(filename) +mh = ncd[:mahalanobis] +lp = ncd[:posterior_logpdf_true_v_map] +# Gaussian: mh=-2lp +(n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) +# assume dimnames(mh) = "random_seed", "ensemble_size", "k_iter" + +# count missings +mh1_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp1_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] + +# criteria of "good" posterior +qq_good = quantile(Chisq(n_par), [0.1, 0.5, 0.9]) +# now we compare over the rng dimension +mh_scores = fill(NaN, length(qq_good), size(mh,2), size(mh,3)) +lp_scores = fill(NaN, length(qq_good), size(lp,2), size(lp,3)) +for (idx,qg) in enumerate(qq_good) + mh_scores[idx,:,:] = [sum(skipmissing(mh[:,i,j]) .<= qg; init=0) for i in axes(mh,2), j in axes(mh,3)] + lp_scores[idx,:,:] = [sum(skipmissing(-2*lp[:,i,j]) .<= qg; init=0) for i in axes(lp,2), j in axes(lp,3)] + # divide by the # non-missing values. In case of all missing we divide by 1 + mh_scores[idx,:,:] ./= max.(mh1_nonmiss,1) + lp_scores[idx,:,:] ./= max.(lp1_nonmiss,1) +end + + From a9de9c4c564799d364d3a7b840ec448e0f81442f Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 15:48:32 -0700 Subject: [PATCH 21/86] remove so many figures --- examples/EKIRace/emulate_sample_l63.jl | 3 +-- examples/EKIRace/emulate_sample_l96.jl | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/EKIRace/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl index f423d86db..622499a87 100644 --- a/examples/EKIRace/emulate_sample_l63.jl +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -129,9 +129,8 @@ function main() vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) end - figpath = joinpath(figure_save_directory, "posterior_hist_$(calib_filename_suffix)_k$(k)") + figpath = joinpath(figure_save_directory, "l63_posterior_hist_$(calib_filename_suffix)_k$(k)") savefig(figpath * ".png") - savefig(figpath * ".pdf") posteriors_by_k[k] = posterior iop_by_k[k] = input_output_pairs diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index 9f9e6bf8b..30011025d 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -144,7 +144,6 @@ function main() end figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)_k$(k)") savefig(figpath * ".png") - savefig(figpath * ".pdf") posteriors_by_k[k] = posterior iop_by_k[k] = input_output_pairs From 7f37c384c5e10cedef9da31bc29f6bdf96ec3bd7 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 16:02:00 -0700 Subject: [PATCH 22/86] trim to max k_iter --- examples/EKIRace/l63_exp_to_leaderboard_utilities.jl | 12 +++++++++++- examples/EKIRace/l96_exp_to_leaderboard_utilities.jl | 12 +++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index 3be38b873..fd388e348 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -50,7 +50,7 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -n_k = cfg.max_iter # maximum number of training-iteration slices +n_k = cfg.max_iter # upper bound for pre-allocation; trimmed after loading n_rng = length(rng_idxs) n_ens = length(N_enss) @@ -63,6 +63,8 @@ post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +actual_n_k = 0 # track the maximum k actually written across all files + for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" @@ -71,6 +73,7 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] + actual_n_k = max(actual_n_k, maximum(k_values)) # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -102,6 +105,13 @@ for (N_ens, rng_idx) in valid_file_items end +# Trim k_iter dimension to the maximum k actually present in the data +n_k = actual_n_k +post_mean_arr = post_mean_arr[:, :, 1:n_k, :] +post_cov_arr = post_cov_arr[:, :, 1:n_k, :, :] +mahal_arr = mahal_arr[:, :, 1:n_k] +logpdf_true_v_map_arr = logpdf_true_v_map_arr[:, :, 1:n_k] + ### Saving the data in nc ds = NCDataset(nc_save_filename, "c") diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index b875a5d2b..d55ba6bf0 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -52,7 +52,7 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -n_k = cfg.max_iter # maximum number of training-iteration slices +n_k = cfg.max_iter # upper bound for pre-allocation; trimmed after loading n_rng = length(rng_idxs) n_ens = length(N_enss) @@ -65,6 +65,8 @@ post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +actual_n_k = 0 # track the maximum k actually written across all files + for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" @@ -73,6 +75,7 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] + actual_n_k = max(actual_n_k, maximum(k_values)) # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -104,6 +107,13 @@ for (N_ens, rng_idx) in valid_file_items end +# Trim k_iter dimension to the maximum k actually present in the data +n_k = actual_n_k +post_mean_arr = post_mean_arr[:, :, 1:n_k, :] +post_cov_arr = post_cov_arr[:, :, 1:n_k, :, :] +mahal_arr = mahal_arr[:, :, 1:n_k] +logpdf_true_v_map_arr = logpdf_true_v_map_arr[:, :, 1:n_k] + ### Saving the data in nc ds = NCDataset(nc_save_filename, "c") From 29103368c400852c4b90ef288cc299ff84fe7db6 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 29 May 2026 16:09:13 -0700 Subject: [PATCH 23/86] get the right n_k --- .../EKIRace/l63_exp_to_leaderboard_utilities.jl | 15 ++++----------- .../EKIRace/l96_exp_to_leaderboard_utilities.jl | 15 ++++----------- 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index fd388e348..d9a64582f 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -50,7 +50,10 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -n_k = cfg.max_iter # upper bound for pre-allocation; trimmed after loading +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) n_rng = length(rng_idxs) n_ens = length(N_enss) @@ -63,8 +66,6 @@ post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) -actual_n_k = 0 # track the maximum k actually written across all files - for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" @@ -73,7 +74,6 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] - actual_n_k = max(actual_n_k, maximum(k_values)) # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -105,13 +105,6 @@ for (N_ens, rng_idx) in valid_file_items end -# Trim k_iter dimension to the maximum k actually present in the data -n_k = actual_n_k -post_mean_arr = post_mean_arr[:, :, 1:n_k, :] -post_cov_arr = post_cov_arr[:, :, 1:n_k, :, :] -mahal_arr = mahal_arr[:, :, 1:n_k] -logpdf_true_v_map_arr = logpdf_true_v_map_arr[:, :, 1:n_k] - ### Saving the data in nc ds = NCDataset(nc_save_filename, "c") diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index d55ba6bf0..a69532fe7 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -52,7 +52,10 @@ end first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) -n_k = cfg.max_iter # upper bound for pre-allocation; trimmed after loading +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) n_rng = length(rng_idxs) n_ens = length(N_enss) @@ -65,8 +68,6 @@ post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) -actual_n_k = 0 # track the maximum k actually written across all files - for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" @@ -75,7 +76,6 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] - actual_n_k = max(actual_n_k, maximum(k_values)) # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -107,13 +107,6 @@ for (N_ens, rng_idx) in valid_file_items end -# Trim k_iter dimension to the maximum k actually present in the data -n_k = actual_n_k -post_mean_arr = post_mean_arr[:, :, 1:n_k, :] -post_cov_arr = post_cov_arr[:, :, 1:n_k, :, :] -mahal_arr = mahal_arr[:, :, 1:n_k] -logpdf_true_v_map_arr = logpdf_true_v_map_arr[:, :, 1:n_k] - ### Saving the data in nc ds = NCDataset(nc_save_filename, "c") From f7b96e4480761edc83968a61731167f9fc6adaef Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 01:39:09 -0700 Subject: [PATCH 24/86] added initial hpc scripts --- examples/EKIRace/hpc-variant/Lorenz63.jl | 138 +++++++ examples/EKIRace/hpc-variant/Lorenz96.jl | 192 ++++++++++ examples/EKIRace/hpc-variant/Project.toml | 13 + examples/EKIRace/hpc-variant/README.md | 116 ++++++ .../hpc-variant/calibrate_array.sbatch | 41 +++ examples/EKIRace/hpc-variant/calibrate_l63.jl | 274 ++++++++++++++ examples/EKIRace/hpc-variant/calibrate_l96.jl | 348 ++++++++++++++++++ .../hpc-variant/emulate_sample_array.sbatch | 38 ++ .../EKIRace/hpc-variant/emulate_sample_l63.jl | 166 +++++++++ .../EKIRace/hpc-variant/emulate_sample_l96.jl | 180 +++++++++ .../EKIRace/hpc-variant/experiment_config.jl | 182 +++++++++ examples/EKIRace/hpc-variant/submit_l63.sh | 36 ++ .../EKIRace/hpc-variant/submit_l96_const.sh | 36 ++ .../EKIRace/hpc-variant/submit_l96_flux.sh | 36 ++ .../EKIRace/hpc-variant/submit_l96_vec.sh | 36 ++ 15 files changed, 1832 insertions(+) create mode 100644 examples/EKIRace/hpc-variant/Lorenz63.jl create mode 100644 examples/EKIRace/hpc-variant/Lorenz96.jl create mode 100644 examples/EKIRace/hpc-variant/Project.toml create mode 100644 examples/EKIRace/hpc-variant/README.md create mode 100644 examples/EKIRace/hpc-variant/calibrate_array.sbatch create mode 100644 examples/EKIRace/hpc-variant/calibrate_l63.jl create mode 100644 examples/EKIRace/hpc-variant/calibrate_l96.jl create mode 100644 examples/EKIRace/hpc-variant/emulate_sample_array.sbatch create mode 100644 examples/EKIRace/hpc-variant/emulate_sample_l63.jl create mode 100644 examples/EKIRace/hpc-variant/emulate_sample_l96.jl create mode 100644 examples/EKIRace/hpc-variant/experiment_config.jl create mode 100755 examples/EKIRace/hpc-variant/submit_l63.sh create mode 100755 examples/EKIRace/hpc-variant/submit_l96_const.sh create mode 100755 examples/EKIRace/hpc-variant/submit_l96_flux.sh create mode 100755 examples/EKIRace/hpc-variant/submit_l96_vec.sh diff --git a/examples/EKIRace/hpc-variant/Lorenz63.jl b/examples/EKIRace/hpc-variant/Lorenz63.jl new file mode 100644 index 000000000..1cff36af5 --- /dev/null +++ b/examples/EKIRace/hpc-variant/Lorenz63.jl @@ -0,0 +1,138 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using StatsPlots +using Plots +using Random +using JLD2 +using Statistics + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +# This will change for different Lorenz simulators +struct LorenzConfig{FT1 <: Real, FT2 <: Real} + "Length of a fixed integration timestep" + dt::FT1 + "Total duration of integration (T = N*dt)" + T::FT2 +end + +# This will change for each ensemble member +struct EnsembleMemberConfig{VV <: AbstractVector} + "rho, beta (unknowns)" + u::VV +end + +# This will change for different "Observations" of Lorenz +struct ObservationConfig{FT1 <: Real, FT2 <: Real} + "initial time to gather statistics (T_start = N_start*dt)" + T_start::FT1 + "end time to gather statistics (T_end = N_end*dt)" + T_end::FT2 +end + +######################################################################### +############################ Model Functions ############################ +######################################################################### + +# Forward pass of forward model +# Inputs: +# - params: structure with u (unknowns vector) +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_forward( + params::EnsembleMemberConfig, + x0::VorM, + config::LorenzConfig, + observation_config::ObservationConfig, +) where {VorM <: AbstractVecOrMat} + # run the Lorenz simulation + xn = lorenz_solve(params, x0, config) + # Get statistics + gt = stats(xn, config, observation_config) + return gt +end + +#Calculates statistics for forward model output +# Inputs: +# - xn: timeseries of states for length of simulation through Lorenz63 +function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat} + T_start = observation_config.T_start + T_end = observation_config.T_end + dt = config.dt + N_start = Int(ceil(T_start / dt)) + N_end = Int(ceil(T_end / dt)) + xn_stat = xn[:, N_start:N_end] + N_state = size(xn_stat, 1) + gt = zeros(9) # Might want to switch to more general statement? + gt[1:3] = mean(xn_stat, dims = 2) + xn_stat_cov = cov(xn_stat, dims = 2) + gt[4:6] = diag(xn_stat_cov) + gt[7:8] = xn_stat_cov[1, 2:3] + gt[9] = xn_stat_cov[2, 3] + return gt +end + +# Forward pass of the Lorenz 96 model +# Inputs: +# - params: structure with u (unknowns vector) +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + # Initialize + nstep = Int(ceil(config.T / config.dt)) + state_dim = isa(x0, AbstractVector) ? length(x0) : size(x0, 1) + xn = zeros(size(x0, 1), nstep + 1) + xn[:, 1] = x0 + + # March forward in time + for j in 1:nstep + xn[:, j + 1] = RK4(params, xn[:, j], config) + end + # Output + return xn +end + +# Lorenz 96 system +# f = dx/dt +# Inputs: +# - params: structure with u (unknowns vector) +# - x: current state +function f(params::EnsembleMemberConfig, x::VorM) where {VorM <: AbstractVecOrMat} + u = params.u + N = length(x) + f = zeros(N) + + f[1] = 10.0 * (x[2] - x[1]) + f[2] = x[1] * (u[1] - x[3]) - x[2] + f[3] = x[1] * x[2] - u[2] * x[3] + + # Output + return f +end + +# RK4 solve +# Inputs: +# - params: structure with F (state-dependent-forcing vector) +# - xold: current state +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function RK4(params::EnsembleMemberConfig, xold::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + N = length(xold) + dt = config.dt + + # Predictor steps (note no time-dependence is needed here) + k1 = f(params, xold) + k2 = f(params, xold + k1 * dt / 2.0) + k3 = f(params, xold + k2 * dt / 2.0) + k4 = f(params, xold + k3 * dt) + # Step + xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + # Output + return xnew +end diff --git a/examples/EKIRace/hpc-variant/Lorenz96.jl b/examples/EKIRace/hpc-variant/Lorenz96.jl new file mode 100644 index 000000000..da75d4f93 --- /dev/null +++ b/examples/EKIRace/hpc-variant/Lorenz96.jl @@ -0,0 +1,192 @@ +# Import modules +using LinearAlgebra +using Statistics +using Random, Distributions +using Flux + + + +# This will change for different Lorenz simulators +struct LorenzConfig{FT1 <: Real, FT2 <: Real} + "Length of a fixed integration timestep" + dt::FT1 + "Total duration of integration (T = N*dt)" + T::FT2 +end + +# This will change for each ensemble member +abstract type EnsembleMemberConfig end +# struct EnsembleMemberConfig{FT} +# val::FT +# end + +# Sub-type of ensemble config for constant forcing +struct ConstantEMC{FT <: Real} <: EnsembleMemberConfig + val::FT +end +build_forcing(::T, val::FT, args...) where {T <: ConstantEMC, FT <: Real} = ConstantEMC(val) +build_forcing(::T, val::FT, args...) where {T <: ConstantEMC, FT <: AbstractVector} = ConstantEMC(val[1]) + +# Sub-type of ensemble config for spatially-dependent forcing +struct VectorEMC{VV <: AbstractVector} <: EnsembleMemberConfig + val::VV +end +build_forcing(::T, val::VV, args...) where {T <: VectorEMC, VV <: AbstractVector} = VectorEMC(val) + +# Sub-type of ensemble config for spatially-dependent forcing with neural network approximation +struct FluxEMC{FC <: Flux.Chain, VV <: AbstractVector} <: EnsembleMemberConfig + model::FC + sample_range::VV +end +function build_forcing(::T, params, model, sample_range) where {T <: FluxEMC} + _, reconstructor = Flux.destructure(model) + return FluxEMC(reconstructor(params), Float32.(sample_range)) +end + +# Constant-global +forcing(params::ConstantEMC, x, i) = params.val +forcing(params::ConstantEMC, x) = repeat([params.val], length(x)) + +# Constant-vector +forcing(params::VectorEMC, x, i) = params.val[i] +forcing(params::VectorEMC, x) = params.val + +# Flux +forcing(params::FluxEMC, x, i) = Float64(params.model([params.sample_range[i]])[1]) +forcing(params::FluxEMC, x) = Float64.(params.model([sr])[1] for sr in params.sample_range) + + + + +# This will change for different "Observations" of Lorenz +struct ObservationConfig{FT1 <: Real, FT2 <: Real} + "initial time to gather statistics (T_start = N_start*dt)" + T_start::FT1 + "end time to gather statistics (T_end = N_end*dt)" + T_end::FT2 +end + +######################################################################### +############################ Model Functions ############################ +######################################################################### + +# Forward pass of forward model +# Inputs: +# - params: structure with F +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_forward( + params::EnsembleMemberConfig, + x0::VorM, + config::LorenzConfig, + observation_config::ObservationConfig, +) where {VorM <: AbstractVecOrMat} + # run the Lorenz simulation + xn = lorenz_solve(params, x0, config) + # Get statistics + gt = stats(xn, config, observation_config) + return gt +end + +#Calculates statistics for forward model output +# Inputs: +# - xn: timeseries of states for length of simulation through Lorenz96 +function stats(xn::VorM, config::LorenzConfig, observation_config::ObservationConfig) where {VorM <: AbstractVecOrMat} + T_start = observation_config.T_start + T_end = observation_config.T_end + dt = config.dt + N_start = Int(ceil(T_start / dt)) + N_end = Int(ceil(T_end / dt)) + xn_stat = xn[:, N_start:N_end] + N_state = size(xn_stat, 1) + gt = zeros(2 * N_state) + gt[1:N_state] = mean(xn_stat, dims = 2) + gt[(N_state + 1):(2 * N_state)] = std(xn_stat, dims = 2) + return gt +end + +# Forward pass of the Lorenz 96 model +# Inputs: +# - params: structure with F +# - x0: initial condition vector +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function lorenz_solve(params::EnsembleMemberConfig, x0::VorM, config::LorenzConfig) where {VorM <: AbstractVecOrMat} + # Initialize + nstep = Int(ceil(config.T / config.dt)) + state_dim = isa(x0, AbstractVector) ? length(x0) : size(x0, 1) + xn = zeros(size(x0, 1), nstep + 1) + xn[:, 1] = x0 + + # March forward in time + forcing_vec = forcing(params, x0) # not state dependent so evaluate once here + for j in 1:nstep + xn[:, j + 1] = RK4(forcing_vec, xn[:, j], config) + end + # Output + return xn +end + +# Lorenz 96 system +# f = dx/dt +# Inputs: +# - params: structure with F +# - x: current state +function f(forcing_vec::VV, x::VorM) where {VV <: AbstractVector, VorM <: AbstractVecOrMat} + N = length(x) + f = zeros(N) + # Loop over N positions + for i in 3:(N - 1) + f[i] = -x[i - 2] * x[i - 1] + x[i - 1] * x[i + 1] - x[i] + forcing_vec[i] + end + # Periodic boundary conditions + f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + forcing_vec[1] + f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + forcing_vec[2] + f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + forcing_vec[N] + + # Output + return f +end + +# RK4 solve +# Inputs: +# - params: structure with F +# - xold: current state +# - config: structure including dt (timestep Float64(1)) and T (total time Float64(1)) +function RK4(forcing_vec::VV, xold::VorM, config::LorenzConfig) where {VV <: AbstractVector, VorM <: AbstractVecOrMat} + N = length(xold) + dt = config.dt + + # Predictor steps (note no time-dependence is needed here) + k1 = f(forcing_vec, xold) + k2 = f(forcing_vec, xold + k1 * dt / 2.0) + k3 = f(forcing_vec, xold + k2 * dt / 2.0) + k4 = f(forcing_vec, xold + k3 * dt) + # Step + xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + # Output + return xnew +end + + +# Neural network functions +function train_network(model, x_train, y_train) + loss(model, x, y) = Flux.Losses.mse(model(x), y) + + # Reshape x_train and y_train for Flux compatibility + x_train = reshape(x_train, 1, :) + y_train = reshape(y_train, 1, :) + x_train = Float32.(x_train) + y_train = Float32.(y_train) + + opt = Flux.setup(Adam(), model) + data = Flux.DataLoader((x_train, y_train), batchsize = 32, shuffle = true) # train the model + + # Train the model over multiple epochs + epochs = 5000 + for epoch in 1:epochs + Flux.train!(loss, model, data, opt) + end + + params, _ = Flux.destructure(model) + return model, params +end diff --git a/examples/EKIRace/hpc-variant/Project.toml b/examples/EKIRace/hpc-variant/Project.toml new file mode 100644 index 000000000..39c1c6f84 --- /dev/null +++ b/examples/EKIRace/hpc-variant/Project.toml @@ -0,0 +1,13 @@ +[deps] +BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" +CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd" diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md new file mode 100644 index 000000000..9634a8302 --- /dev/null +++ b/examples/EKIRace/hpc-variant/README.md @@ -0,0 +1,116 @@ +# EKIRace HPC variant + +Parallelised versions of the `calibrate` and `emulate_sample` scripts. +Each script can run serially (identical behaviour to the originals) or as a +SLURM job array where every `(N_ens, rng_idx)` cell is an independent task. + +## One-time setup + +In `experiment_config.jl`, pin the calibrate date to today's date before +starting a run, e.g. + +```julia +calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +``` + +All subsequent stages (emulate_sample, leaderboard) read this value to locate +the right output directory; keeping it fixed avoids mismatches when jobs run +past midnight or across days. + +The submit scripts (`submit_*.sh`) handle precompilation automatically before +queuing any jobs. + +## Standalone (serial) + +Run with no arguments to sweep all `(N_ens, rng_idx)` cells sequentially, +exactly as the original scripts did. + +```bash +# L63 +julia --project=. calibrate_l63.jl +julia --project=. emulate_sample_l63.jl + +# L96 — set EXPERIMENT or edit the toggle in experiment_config.jl +EXPERIMENT=l96_const julia --project=. calibrate_l96.jl +EXPERIMENT=l96_const julia --project=. emulate_sample_l96.jl +``` + +You can also run a single cell by passing its 1-based task index: + +```bash +julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only +julia --project=. emulate_sample_l63.jl 5 # fifth cell only +``` + +## HPC (Caltech Resnick cluster, SLURM) + +### Submission scripts (recommended) + +One script per case handles precompilation and chains calibrate → emulate_sample +automatically. All four can be launched simultaneously — output files are +case-specific so there are no write conflicts. The optional `EXP_ID` argument +suffixes SLURM job names to keep the queue readable: + +```bash +bash submit_l63.sh [EXP_ID] +bash submit_l96_const.sh [EXP_ID] +bash submit_l96_vec.sh [EXP_ID] +bash submit_l96_flux.sh [EXP_ID] + +# example: run all four with a shared label +for s in submit_l63.sh submit_l96_const.sh submit_l96_vec.sh submit_l96_flux.sh; do + bash "$s" run1 & +done +wait +``` + +### Manual submission + +```bash +# L63 +sbatch --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch + +# L96 (submit once per forcing case) +sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch +sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec calibrate_array.sbatch +sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux calibrate_array.sbatch +``` + +Each task writes its per-method `ekp` and `results` files, which are consumed +directly by the emulate_sample stage. + +### Emulate-sample + +Submit after the calibrate array is complete. Use `--dependency` to chain +automatically: + +```bash +# L63 +cid=$(sbatch --parsable --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) +sbatch --dependency=afterok:$cid \ + --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch + +# L96 +cid=$(sbatch --parsable --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ + calibrate_array.sbatch) +sbatch --dependency=afterok:$cid \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ + emulate_sample_array.sbatch +``` + +### Adjusting array size + +The sbatch files default to `--array=1-60` (3 ensemble sizes × 20 repeats). +If you change `N_ens_sizes` or `n_repeats` in `experiment_config.jl`, update +the upper bound to `length(N_ens_sizes) * n_repeats`. The `%20` suffix caps +concurrent tasks to 20 at a time as a cluster-courtesy limit; raise or remove +it if you want faster turnaround. + +### Smoke test + +Before a full submission, run a single-task array to verify the job finds its +input files and writes output correctly: + +```bash +sbatch --array=1-1 --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch +``` diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch new file mode 100644 index 000000000..0392cf232 --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -0,0 +1,41 @@ +#!/bin/bash +# SLURM array job for the calibrate stage. +# array=1-60%20 means run array 1-60, with max of 20 at a time. +# +# L63 (60 tasks = 3 N_ens × 20 repeats): +# sbatch --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch +# +# L96 (same count; adjust SCRIPT and set EXPERIMENT): +# sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch +# sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec calibrate_array.sbatch +# sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux calibrate_array.sbatch +# +# If N_ens_sizes or n_repeats change, update --array upper bound to match +# length(N_ens_sizes) * n_repeats. +# +# One-time precompile before submitting (avoids 60 simultaneous precompiles): +# mkdir -p output/slurm +# julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + +#SBATCH --job-name=calib +#SBATCH --output=output/slurm/calib_%A_%a.out +#SBATCH --error=output/slurm/calib_%A_%a.err +#SBATCH --array=1-60%20 +#SBATCH --time=04:00:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia # pin version as needed, e.g.: module load julia/1.10.4 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +# Default to L63 if SCRIPT not set via --export. +SCRIPT=${SCRIPT:-calibrate_l63.jl} + +cd "$(dirname "$0")" + +julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl new file mode 100644 index 000000000..53538adb5 --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -0,0 +1,274 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Dates + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +include("Lorenz63.jl") +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true + +######################################################################## +############### Deterministic problem setup ############################ +######################################################################## + +function build_setup(cfg) + N_iter = cfg.N_iter + N_ens_sizes = cfg.N_ens_sizes + terminate_at = cfg.terminate_at + # Seeded so every array task sees the same rng_seeds list. + rng_seeds = randperm(MersenneTwister(20260529), 1_000_000)[1:cfg.n_repeats] + + nx = 3 + nu = 2 + truth_params = EnsembleMemberConfig([28.0, 8.0 / 3.0]) + + prior_r = constrained_gaussian("rho", exp(3.3), 4.153, 0, Inf) + prior_b = constrained_gaussian("beta", exp(1.2), 2.016, 0, Inf) + prior = combine_distributions([prior_r, prior_b]) + T = 40.0 + + rng_seed_init = 11 + rng_i = MersenneTwister(rng_seed_init) + + t = 0.01 + T_long = 1000.0 + picking_initial_condition = LorenzConfig(t, T_long) + x_initial = rand(rng_i, Normal(0.0, 1.0), nx) + x_spun_up = lorenz_solve(truth_params, x_initial, picking_initial_condition) + x0 = x_spun_up[:, end] + + ny = 9 + lorenz_config_settings = LorenzConfig(t, T) + T_start = 30.0 + T_end = T + observation_config = ObservationConfig(T_start, T_end) + y = lorenz_forward(truth_params, x0, lorenz_config_settings, observation_config) + + multiple = 36 + window = T_end - T_start + T_R = multiple * window + T_start + R_config = LorenzConfig(t, T_R) + R_run = lorenz_solve(truth_params, x_initial, R_config) + R_sample_size = Int(ceil(multiple)) + R_samples = zeros(ny, R_sample_size) + for ii in 1:R_sample_size + local_obs_config = ObservationConfig(T_start + (ii - 1) * window, T_start + ii * window) + R_samples[:, ii] = stats(R_run, R_config, local_obs_config) + end + R = cov(R_samples, dims = 2) + R_inv_var = sqrt(inv(R)) + + covT = 2000.0 + cov_solve = lorenz_solve(truth_params, x0, LorenzConfig(t, covT)) + ic_cov = 0.1 * cov(cov_solve, dims = 2) + ic_cov_sqrt = sqrt(ic_cov) + + configuration = Dict( + "N_iter" => N_iter, + "N_ens_sizes" => N_ens_sizes, + "terminate_at" => terminate_at, + "rng_seeds" => rng_seeds, + ) + + return (; + nx, nu, ny, + truth_params, prior, + x0, x_initial, + y, R, R_inv_var, + lorenz_config_settings, observation_config, ic_cov_sqrt, + rng_seeds, configuration, + N_iter, terminate_at, + ) +end + +# Write prior files for all 4 method dirs, idempotent. +function write_priors(cfg, setup, output_dir) + for method_name in method_cases + per_method_dir = joinpath(output_dir, calib_directory(method_name, cfg)) + mkpath(per_method_dir) + pf = joinpath(per_method_dir, prior_filename(cfg)) + if !isfile(pf) + JLD2.save(pf, "prior", setup.prior) + end + end +end + +######################################################################## +############### Per-cell calibration (all 4 methods) ################## +######################################################################## + +function calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + rng = MersenneTwister(setup.rng_seeds[rng_idx]) + @info "Calibrating (N_ens=$(N_ens), rng_idx=$(rng_idx))" + + initial_params = construct_initial_ensemble(rng, setup.prior, N_ens) + methods = [ + Inversion(), + TransformInversion(), + GaussNewtonInversion(setup.prior), + Unscented(setup.prior), + ] + + conv_cell = fill(NaN, 4) + params_cell = zeros(4, setup.nu) + output_cell = zeros(4, setup.ny) + + for (kk, method) in enumerate(methods) + @info "Method: $(nameof(typeof(method)))" + if isa(method, Unscented) + ekpobj = EKP.EnsembleKalmanProcess( + setup.y, + setup.R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = setup.terminate_at), + ) + else + ekpobj = EKP.EnsembleKalmanProcess( + initial_params, + setup.y, + setup.R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = setup.terminate_at), + ) + end + Ne = get_N_ens(ekpobj) + + ens_mean_final = zeros(setup.nu) + G_ens_mean_final = zeros(setup.ny) + for i in 1:setup.N_iter + params_i = get_ϕ_final(setup.prior, ekpobj) + ens_mean = mean(params_i, dims = 2)[:] + G_ens_mean = lorenz_forward( + EnsembleMemberConfig(ens_mean), + setup.x0 .+ setup.ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), setup.nx, 1), + setup.lorenz_config_settings, + setup.observation_config, + ) + RMSE_e = norm(setup.R_inv_var * (setup.y - G_ens_mean[:])) / sqrt(setup.ny) + @info "RMSE (at G(u_mean)): $(RMSE_e)" + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] + G_ens = hcat( + [ + lorenz_forward( + EnsembleMemberConfig(params_i[:, j]), + (setup.x0 .+ setup.ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), setup.nx, Ne))[:, j], + setup.lorenz_config_settings, + setup.observation_config, + ) for j in 1:Ne + ]..., + ) + terminated = EKP.update_ensemble!(ekpobj, G_ens) + if !isnothing(terminated) + conv_cell[kk] = i * Ne + break + end + end + params_cell[kk, :] = ens_mean_final + output_cell[kk, :] = G_ens_mean_final + if isnan(conv_cell[kk]) + conv_cell[kk] = setup.N_iter * Ne + end + + if save_all_ekp + per_method_dir = joinpath(output_dir, calib_directory(nameof(typeof(method)), cfg)) + JLD2.save( + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rng_idx)), + "N_ens", N_ens, + "method", method, + "ekpobj", ekpobj, + ) + u_stored = get_u(ekpobj, return_array = false) + g_stored = get_g(ekpobj, return_array = false) + JLD2.save( + joinpath(per_method_dir, results_filename(cfg, N_ens, rng_idx)), + "y", setup.y, + "R", setup.R, + "inputs", u_stored, + "outputs", g_stored, + "truth_params_structure", setup.truth_params, + ) + end + end + + return (conv_cell, params_cell, output_cell) +end + +######################################################################## +############### Summary helpers ######################################## +######################################################################## + +function save_combined_summary(cfg, setup, output_dir, conv, pars, outs) + data_filename = joinpath(output_dir, summary_filename(cfg)) + JLD2.save( + data_filename, + "configuration", setup.configuration, + "method_names", method_names, + "conv_alg_iters", conv, + "final_parameters", pars, + "final_model_output", outs, + ) + @info "Saved summary to $(data_filename)" +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + cfg = experiment_config(:l63) + output_dir = joinpath(@__DIR__, "output") + mkpath(output_dir) + + setup = build_setup(cfg) + write_priors(cfg, setup, output_dir) + + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + # Serial mode: run every cell, assemble and write combined summary. + n_ens = length(cfg.N_ens_sizes) + n_rep = cfg.n_repeats + conv = fill(NaN, 4, n_ens, n_rep) + pars = zeros(4, n_ens, n_rep, setup.nu) + outs = zeros(4, n_ens, n_rep, setup.ny) + for (t, (N_ens, rng_idx)) in enumerate(tasks) + ee = (t - 1) ÷ n_rep + 1 + rr = (t - 1) % n_rep + 1 + c, p, o = calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + conv[:, ee, rr] = c + pars[:, ee, rr, :] = p + outs[:, ee, rr, :] = o + end + save_combined_summary(cfg, setup, output_dir, conv, pars, outs) + else + # Array mode: run one cell; ekp/results files are written inside calibrate_one. + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl new file mode 100644 index 000000000..43afe35d6 --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -0,0 +1,348 @@ +# Import modules +using Distributions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates + +# CES +using EnsembleKalmanProcesses +using EnsembleKalmanProcesses.DataContainers +using EnsembleKalmanProcesses.ParameterDistributions +using EnsembleKalmanProcesses.Localizers + +const EKP = EnsembleKalmanProcesses + +include("Lorenz96.jl") +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true + +######################################################################## +############### Deterministic problem setup ############################ +######################################################################## + +function build_setup(cfg, output_dir) + case = cfg.force_case + N_iter = cfg.N_iter + N_ens_sizes = cfg.N_ens_sizes + terminate_at = cfg.terminate_at + # Seeded so every array task sees the same rng_seeds list. + rng_seeds = randperm(MersenneTwister(20260529), 1_000_000)[1:cfg.n_repeats] + + if case == "const-force" + nx = 40 + nu = 1 + phi = ConstantEMC(8.0) + phi_structure = nothing + sample_range = nothing + prior = constrained_gaussian("φ", 10.0, 4.0, 0, Inf) + T = 14.0 + inff = 2 + + elseif case == "vec-force" + nx = 40 + nu = nx + sinusoid = 8 .+ 6 * sin.((4 * pi * range(0, stop = nx - 1, step = 1)) / nx) + phi = VectorEMC(sinusoid) + phi_structure = nothing + sample_range = nothing + pl, psig = 2.0, 3.0 + prior_cov = zeros(nx, nx) + for ii in 1:nx, jj in 1:nx + prior_cov[ii, jj] = psig^2 * exp(-abs(ii - jj) / pl) + end + prior_mean = 8.0 * ones(nx) + distribution = Parameterized(MvNormal(prior_mean, prior_cov)) + constraint = repeat([no_constraint()], 40) + prior = ParameterDistribution(distribution, constraint, "ml96_prior") + T = 54.0 + inff = 2 + + elseif case == "flux-force" + nx = 100 + nu = 61 + true_sinusoid(x) = 8 .+ 6 * sin.((4 * pi * x) / 10) + x_train = collect(-5.0:0.01:5.0) + y_train = true_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) + phi_structure = Chain(Dense(1 => 20, tanh), Dense(20 => 1)) + true_model, _ = train_network(phi_structure, x_train, y_train) + sample_range = Float32.(collect(-5.0:0.1:4.9)) + phi = FluxEMC(true_model, sample_range) + prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) + prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) + prior_model, prior_mean = train_network(phi_structure, x_train, prior_train) + prior_cov = (0.1^2) * I(length(prior_mean)) + distribution = Parameterized(MvNormal(prior_mean, prior_cov)) + constraint = repeat([no_constraint()], 61) + prior = ParameterDistribution(distribution, constraint, "l96_nn_prior") + T = 54.0 + inff = 2.5 + else + throw(ArgumentError("Unknown force_case: $case")) + end + + ny = nx * 2 + rng_seed_init = 11 + rng_i = MersenneTwister(rng_seed_init) + + # Expensive computations are cached to a prelim file so all array tasks load the same data. + prelim_file = joinpath(output_dir, "l96_computed_preliminaries_$(case).jld2") + if isfile(prelim_file) + loaded = JLD2.load(prelim_file) + x0 = loaded["x0"] + y = loaded["y"] + ic_cov_sqrt = loaded["ic_cov_sqrt"] + R = loaded["R"] + R_inv_var = loaded["R_inv_var"] + lorenz_config_settings = loaded["lorenz_config_settings"] + observation_config = loaded["observation_config"] + @info "Loaded precomputed preliminaries from $(prelim_file)" + else + t = 0.01 + T_long = 1000.0 + lorenz_config_settings = LorenzConfig(t, T) + @info "Initial spin-up: running $(T_long) time units" + picking_initial_condition = LorenzConfig(t, T_long) + x_initial = rand(rng_i, Normal(0.0, 1.0), nx) + x_spun_up = lorenz_solve(phi, x_initial, picking_initial_condition) + x0 = x_spun_up[:, end] + T_start = 4.0 + T_end = T + observation_config = ObservationConfig(T_start, T_end) + @info "Computing synthetic data: running $(T) time units" + y = lorenz_forward(phi, x0, lorenz_config_settings, observation_config) + window = T_end - T_start + T_R = 10 * window * ny + T_start + R_config = LorenzConfig(t, T_R) + @info "Estimating noise covariance: running $(T_R) time units" + R_run = lorenz_solve(phi, x_initial, R_config) + R_sample_size = Int(ceil(10 * ny)) + R_samples = zeros(ny, R_sample_size) + for ii in 1:R_sample_size + local_obs_config = ObservationConfig(T_start + (ii - 1) * window, T_start + ii * window) + R_samples[:, ii] = stats(R_run, R_config, local_obs_config) + end + R = cov(R_samples, dims = 2) * inff + R_inv_var = sqrt(inv(R)) + covT = 2000.0 + @info "Estimating IC covariance: running $(covT) time units" + cov_solve = lorenz_solve(phi, x0, LorenzConfig(t, covT)) + ic_cov = 0.1 * cov(cov_solve, dims = 2) + ic_cov_sqrt = sqrt(ic_cov) + JLD2.save( + prelim_file, + "x0", x0, "y", y, + "lorenz_config_settings", lorenz_config_settings, + "observation_config", observation_config, + "ic_cov_sqrt", ic_cov_sqrt, + "R", R, "R_inv_var", R_inv_var, + ) + @info "Saved computed quantities to $(prelim_file)" + end + + configuration = Dict( + "case" => case, + "N_iter" => N_iter, + "N_ens_sizes" => N_ens_sizes, + "rng_seeds" => rng_seeds, + "terminate_at" => terminate_at, + ) + + return (; + nx, nu, ny, + phi, phi_structure, sample_range, prior, + x0, y, R, R_inv_var, + lorenz_config_settings, observation_config, ic_cov_sqrt, + rng_seeds, configuration, + N_iter, terminate_at, + ) +end + +# Write prior files for all 4 method dirs, idempotent. +function write_priors(cfg, setup, output_dir) + for method_name in method_cases + per_method_dir = joinpath(output_dir, calib_directory(method_name, cfg)) + mkpath(per_method_dir) + pf = joinpath(per_method_dir, prior_filename(cfg)) + if !isfile(pf) + JLD2.save(pf, "prior", setup.prior) + end + end +end + +######################################################################## +############### Per-cell calibration (all 4 methods) ################## +######################################################################## + +function calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + rng = MersenneTwister(setup.rng_seeds[rng_idx]) + @info "Calibrating (N_ens=$(N_ens), rng_idx=$(rng_idx))" + + initial_params = construct_initial_ensemble(rng, setup.prior, N_ens) + methods = [ + Inversion(), + TransformInversion(), + GaussNewtonInversion(setup.prior), + Unscented(setup.prior), + ] + + conv_cell = fill(NaN, 4) + params_cell = zeros(4, setup.nu) + output_cell = zeros(4, setup.ny) + + for (kk, method) in enumerate(methods) + @info "Method: $(nameof(typeof(method)))" + if isa(method, Unscented) + ekpobj = EKP.EnsembleKalmanProcess( + setup.y, + setup.R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = setup.terminate_at), + ) + else + ekpobj = EKP.EnsembleKalmanProcess( + initial_params, + setup.y, + setup.R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = setup.terminate_at), + ) + end + Ne = get_N_ens(ekpobj) + + ens_mean_final = zeros(setup.nu) + G_ens_mean_final = zeros(setup.ny) + for i in 1:setup.N_iter + params_i = get_ϕ_final(setup.prior, ekpobj) + ens_mean = mean(params_i, dims = 2)[:] + forcing = build_forcing(setup.phi, ens_mean, setup.phi_structure, setup.sample_range) + G_ens_mean = lorenz_forward( + forcing, + setup.x0 .+ setup.ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), setup.nx, 1), + setup.lorenz_config_settings, + setup.observation_config, + ) + RMSE_e = norm(setup.R_inv_var * (setup.y - G_ens_mean[:])) / sqrt(setup.ny) + @info "RMSE (at G(u_mean)): $(RMSE_e)" + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] + G_ens = hcat( + [ + lorenz_forward( + build_forcing(setup.phi, params_i[:, j], setup.phi_structure, setup.sample_range), + (setup.x0 .+ setup.ic_cov_sqrt * rand(rng, Normal(0.0, 1.0), setup.nx, Ne))[:, j], + setup.lorenz_config_settings, + setup.observation_config, + ) for j in 1:Ne + ]..., + ) + terminated = EKP.update_ensemble!(ekpobj, G_ens) + if !isnothing(terminated) + conv_cell[kk] = i * Ne + break + end + end + params_cell[kk, :] = ens_mean_final + output_cell[kk, :] = G_ens_mean_final + if isnan(conv_cell[kk]) + conv_cell[kk] = setup.N_iter * Ne + end + + if save_all_ekp + per_method_dir = joinpath(output_dir, calib_directory(nameof(typeof(method)), cfg)) + JLD2.save( + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rng_idx)), + "N_ens", N_ens, + "method", method, + "ekpobj", ekpobj, + ) + u_stored = get_u(ekpobj, return_array = false) + g_stored = get_g(ekpobj, return_array = false) + JLD2.save( + joinpath(per_method_dir, results_filename(cfg, N_ens, rng_idx)), + "y", setup.y, + "R", setup.R, + "inputs", u_stored, + "outputs", g_stored, + "truth_params_structure", (setup.phi, setup.phi_structure), + ) + end + end + + return (conv_cell, params_cell, output_cell) +end + +######################################################################## +############### Summary helpers ######################################## +######################################################################## + +function save_combined_summary(cfg, setup, output_dir, conv, pars, outs) + data_filename = joinpath(output_dir, summary_filename(cfg)) + JLD2.save( + data_filename, + "configuration", setup.configuration, + "method_names", method_names, + "conv_alg_iters", conv, + "final_parameters", pars, + "final_model_output", outs, + ) + @info "Saved summary to $(data_filename)" +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + exp = l96_experiment() + @assert exp in (:l96_const, :l96_vec, :l96_flux) \ + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)" + + cfg = experiment_config(exp) + output_dir = joinpath(@__DIR__, "output") + mkpath(output_dir) + + setup = build_setup(cfg, output_dir) + write_priors(cfg, setup, output_dir) + + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + # Serial mode: run every cell, assemble and write combined summary. + n_ens = length(cfg.N_ens_sizes) + n_rep = cfg.n_repeats + conv = fill(NaN, 4, n_ens, n_rep) + pars = zeros(4, n_ens, n_rep, setup.nu) + outs = zeros(4, n_ens, n_rep, setup.ny) + for (t, (N_ens, rng_idx)) in enumerate(tasks) + ee = (t - 1) ÷ n_rep + 1 + rr = (t - 1) % n_rep + 1 + c, p, o = calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + conv[:, ee, rr] = c + pars[:, ee, rr, :] = p + outs[:, ee, rr, :] = o + end + save_combined_summary(cfg, setup, output_dir, conv, pars, outs) + else + # Array mode: run one cell; ekp/results files are written inside calibrate_one. + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch new file mode 100644 index 000000000..ca8788461 --- /dev/null +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -0,0 +1,38 @@ +#!/bin/bash +# SLURM array job for the emulate-sample stage. +# array=1-60%20 means run array 1-60, with max of 20 at a time. +# +# L63: +# sbatch --dependency=afterok: \ +# --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch +# +# L96: +# sbatch --dependency=afterok: \ +# --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const emulate_sample_array.sbatch +# +# Full chained submission example for L63: +# cid=$(sbatch --parsable --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) +# sbatch --dependency=afterok:$cid \ +# --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch + +#SBATCH --job-name=emusample +#SBATCH --output=output/slurm/emusample_%A_%a.out +#SBATCH --error=output/slurm/emusample_%A_%a.err +#SBATCH --array=1-60%20 +#SBATCH --time=12:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia # pin version as needed, e.g.: module load julia/1.10.4 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +SCRIPT=${SCRIPT:-emulate_sample_l63.jl} + +cd "$(dirname "$0")" + +julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l63.jl b/examples/EKIRace/hpc-variant/emulate_sample_l63.jl new file mode 100644 index 000000000..7e62d869b --- /dev/null +++ b/examples/EKIRace/hpc-variant/emulate_sample_l63.jl @@ -0,0 +1,166 @@ +# Import modules +include(joinpath(@__DIR__, "..", "..", "ci", "linkfig.jl")) + +using Distributions +using LinearAlgebra +ENV["GKSwstype"] = "100" +using StatsPlots +using Plots +using Random +using JLD2 +using Dates + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers + +include("Lorenz63.jl") +include("experiment_config.jl") + +######################################################################## +############### Per-cell emulate–sample ################################ +######################################################################## + +function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) + calib_dir = calib_directory(method, cfg) + calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L63: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + homedir = pwd() + figure_save_directory = joinpath(homedir, "output/", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) + loaded_calib_files = [ + joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)), + joinpath(data_save_directory, prior_filename(cfg)), + joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)), + ] + if !isfile(loaded_calib_files[1]) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l63.jl")) + end + + ekpobj = load(loaded_calib_files[1])["ekpobj"] + priors = load(loaded_calib_files[2])["prior"] + truth_sample = load(loaded_calib_files[3])["y"] + truth_params_obj = load(loaded_calib_files[3])["truth_params_structure"] + truth_params_constrained = truth_params_obj.u + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = get_obs_noise_cov(ekpobj) + n_params = length(truth_params) + + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + retain_var = cfg.retain_var + nugget = 1e-6 + + iter_end = length(get_u(ekpobj)) - 1 + K = min(iter_end - 1, cfg.max_iter) + if K < 1 + @error("Not enough EKP iterations for emulate-sample loop. Skipping $(calib_filename_suffix)...") + return + end + @info "Running emulate-sample for k = 1:$(K) training iterations" + + posteriors_by_k = Dict{Int, Any}() + iop_by_k = Dict{Int, Any}() + encoder_schedule_by_k = Dict{Int, Any}() + + for k in 1:K + rng = Random.MersenneTwister(44011) + input_output_pairs = Utilities.get_training_points(ekpobj, 1:k) + @info "k = $(k): training on $(N_ens * k) samples (iterations 1:$(k))" + + ### Emulate + overrides = Dict( + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => 5.0, + "n_iteration" => 10, + "n_features_opt" => cfg.n_features_opt, + ) + mlt = ScalarRandomFeatureInterface( + cfg.n_features, + n_params, + rng = rng, + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), + optimizer_options = overrides, + ) + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # diagnostics + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + diff = truth_sample - vec(y_mean) + Γypluscov = Symmetric(mean(y_var) + Γy) + println("k=$(k) | ML weighted-MSE (truth): ", mean(diff' * (Γypluscov \ diff))) + + ### Sample: MCMC + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + println("k=$(k) | Begin MCMC - step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + println("k=$(k) | post_mean: ", mean(posterior)) + println("k=$(k) | D util: ", det(inv(cov(posterior)))) + + # figure + pp = plot(priors, c = :gray) + plot!(pp, posterior) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "l63_posterior_hist_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + + posteriors_by_k[k] = posterior + iop_by_k[k] = input_output_pairs + encoder_schedule_by_k[k] = get_encoder_schedule(emulator) + end + + save( + joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), + "posteriors_by_k", posteriors_by_k, + "input_output_pairs_by_k", iop_by_k, + "encoder_schedules_by_k", encoder_schedule_by_k, + "priors", priors, + "truth_params_constrained", truth_params_constrained, + "final_params_constrained", final_params_constrained, + "truth_params", truth_params, + "k_values", collect(1:K), + ) +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + cfg = experiment_config(:l63) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + emulate_sample_one(cfg, N_ens, rng_idx) + end + else + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + emulate_sample_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl new file mode 100644 index 000000000..ab8507f6c --- /dev/null +++ b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl @@ -0,0 +1,180 @@ +# Import modules +include(joinpath(@__DIR__, "..", "..", "ci", "linkfig.jl")) + +using Distributions +using LinearAlgebra +ENV["GKSwstype"] = "100" +using StatsPlots +using Plots +using Random +using JLD2 +using Dates + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers + +include("Lorenz96.jl") +include("experiment_config.jl") + +######################################################################## +############### Per-cell emulate–sample ################################ +######################################################################## + +function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) + calib_dir = calib_directory(method, cfg) + calib_filename_suffix = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + homedir = pwd() + figure_save_directory = joinpath(homedir, "output/", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) + loaded_calib_files = [ + joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)), + joinpath(data_save_directory, prior_filename(cfg)), + joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)), + ] + if !isfile(loaded_calib_files[1]) + throw(ErrorException("data files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + + ekpobj = load(loaded_calib_files[1])["ekpobj"] + priors = load(loaded_calib_files[2])["prior"] + truth_sample = load(loaded_calib_files[3])["y"] + (truth_params_obj, truth_params_structure) = load(loaded_calib_files[3])["truth_params_structure"] + + force_case = cfg.force_case + truth_params_constrained = if force_case == "const-force" + [truth_params_obj.val] + elseif force_case == "vec-force" + truth_params_obj.val + elseif force_case == "flux-force" + truth_params_model = truth_params_obj.model + truth_params_constrained, rebuild = destructure(truth_params_model) + truth_params_constrained + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = get_obs_noise_cov(ekpobj) + n_params = length(truth_params) + + final_params_constrained = get_ϕ_mean_final(priors, ekpobj) + encoder_kwargs = encoder_kwargs_from(ekpobj, priors) + retain_var = cfg.retain_var + nugget = 1e-6 + + iter_end = length(get_u(ekpobj)) - 1 + K = min(iter_end - 1, cfg.max_iter) + if K < 1 + @error("Not enough EKP iterations for emulate-sample loop. Skipping $(calib_filename_suffix)...") + return + end + @info "Running emulate-sample for k = 1:$(K) training iterations" + + posteriors_by_k = Dict{Int, Any}() + iop_by_k = Dict{Int, Any}() + encoder_schedule_by_k = Dict{Int, Any}() + + for k in 1:K + rng = Random.MersenneTwister(44011) + input_output_pairs = Utilities.get_training_points(ekpobj, 1:k) + @info "k = $(k): training on $(N_ens * k) samples (iterations 1:$(k))" + + ### Emulate + csm = N_ens * k < 200 ? 5.0 : 1.0 + overrides = Dict( + "scheduler" => DataMisfitController(terminate_at = 1000.0), + "cov_sample_multiplier" => csm, + "n_iteration" => 10, + "n_features_opt" => cfg.n_features_opt, + ) + mlt = ScalarRandomFeatureInterface( + cfg.n_features, + n_params, + rng = rng, + kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), + optimizer_options = overrides, + ) + encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + emulator = Emulator( + mlt, + input_output_pairs; + encoder_schedule = deepcopy(encoder_schedule), + encoder_kwargs = encoder_kwargs, + ) + optimize_hyperparameters!(emulator) + + # diagnostics + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1)) + diff = truth_sample - vec(y_mean) + Γypluscov = Symmetric(mean(y_var) + Γy) + println("k=$(k) | ML weighted-MSE (truth): ", mean(diff' * (Γypluscov \ diff))) + + ### Sample: MCMC + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) + println("k=$(k) | Begin MCMC - step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + println("k=$(k) | post_mean: ", mean(posterior)) + println("k=$(k) | D util: ", det(inv(cov(posterior)))) + + # figure + pp = plot(priors, c = :gray) + plot!(pp, posterior) + for (i, sp) in enumerate(pp.subplots) + vline!(sp, [truth_params_constrained[i]], lc = :black, lw = 4) + vline!(sp, [final_params_constrained[i]], lc = :magenta, lw = 4) + end + figpath = joinpath(figure_save_directory, "l96_posterior_hist_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + + posteriors_by_k[k] = posterior + iop_by_k[k] = input_output_pairs + encoder_schedule_by_k[k] = get_encoder_schedule(emulator) + end + + save( + joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)), + "posteriors_by_k", posteriors_by_k, + "input_output_pairs_by_k", iop_by_k, + "encoder_schedules_by_k", encoder_schedule_by_k, + "priors", priors, + "truth_params_constrained", truth_params_constrained, + "final_params_constrained", final_params_constrained, + "truth_params", truth_params, + "k_values", collect(1:K), + ) +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + exp = l96_experiment() + @assert exp in (:l96_const, :l96_vec, :l96_flux) \ + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)" + cfg = experiment_config(exp) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + emulate_sample_one(cfg, N_ens, rng_idx) + end + else + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + emulate_sample_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl new file mode 100644 index 000000000..308d0e195 --- /dev/null +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -0,0 +1,182 @@ +using Dates + +######################################################################## +############### USER TOGGLE ######################################### +######################################################################## +# Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux +EXPERIMENT = :l96_const + +# Date identifying this calibration run — written by calibrate, read by +# emulate_sample and exp_to_leaderboard (so all stages stay in sync). +# PIN this before submitting an array job so all tasks use the same directory. +calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +#calibrate_date = today() + +######################################################################## +############### SHARED CONSTANTS #################################### +######################################################################## +method_cases = ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] + +method_names = [ + ("Inversion()", "EKI"), + ("TransformInversion()", "ETKI"), + ("GaussNewtonInversion()", "GNKI"), + ("Unscented(prior)", "UKI"), +] + +method_cases_key = Dict( + "Inversion" => "ces-eki-dmc", + "TransformInversion" => "ces-etki-dmc", + "Unscented" => "ces-uki-dmc", + "GaussNewtonInversion" => "ces-iekf", +) + +forcing_cases_key = Dict( + "const-force" => "ensemble_results", + "vec-force" => "spatial_forcing_ensemble_results", + "flux-force" => "nn_forcing_ensemble_results", +) + +######################################################################## +############### PER-CASE CONFIG ##################################### +######################################################################## +# Important dials: +# terminate_at: psuedotime to terminate. (T=1 is approx posterior for the finite-time method variants) +# N_iter: otherwise, maximum iterations +# N_ens_sizes: Vector of experiments +# n_repeats: number of rng seeds + +function experiment_config(case::Symbol) + if case == :l63 + return ( + model = "l63", + force_case = nothing, + N_ens_sizes = [10, 25, 40], + N_iter = 20, + terminate_at = 2.0, # DataMisfitController end time + n_repeats = 20, + max_iter = 10, + retain_var = 0.99, + n_features = 100, + n_features_opt = 60, + calibrate_date = calibrate_date, + ) + elseif case == :l96_const + return ( + model = "l96", + force_case = "const-force", + N_ens_sizes = [5, 15, 30], + N_iter = 20, + terminate_at = 2.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_vec + return ( + model = "l96", + force_case = "vec-force", + N_ens_sizes = [50, 75, 100], + N_iter = 20, + terminate_at = 2.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_flux + return ( + model = "l96", + force_case = "flux-force", + N_ens_sizes = [50, 75, 100], + N_iter = 20, + terminate_at = 2.0, + n_repeats = 20, + max_iter = 15, + retain_var = 0.95, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + else + throw(ArgumentError("Unknown experiment: $case. Expected one of :l63, :l96_const, :l96_vec, :l96_flux")) + end +end + +######################################################################## +############### FILENAME BUILDERS ################################### +######################################################################## +function case_suffix(cfg, N_ens, rng_idx) + if cfg.force_case === nothing + return "$(N_ens)_$(rng_idx)" + else + return "$(cfg.force_case)_$(N_ens)_$(rng_idx)" + end +end + +calib_directory(method, cfg) = "$(method)_$(cfg.calibrate_date)" + +function prior_filename(cfg) + if cfg.force_case === nothing + return "$(cfg.model)_priors.jld2" + else + return "$(cfg.model)_priors_$(cfg.force_case).jld2" + end +end + +ekp_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_ekp_$(case_suffix(cfg, N_ens, rng_idx)).jld2" +results_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_calibrate_results_$(case_suffix(cfg, N_ens, rng_idx)).jld2" +posterior_filename(cfg, N_ens, rng_idx) = "$(cfg.model)_posterior_$(case_suffix(cfg, N_ens, rng_idx)).jld2" + +function summary_filename(cfg) + if cfg.force_case === nothing + return "$(cfg.model)_output_$(cfg.calibrate_date).jld2" + else + return "$(cfg.model)_calibrate_$(cfg.force_case)_$(cfg.calibrate_date).jld2" + end +end + +function nc_filename(cfg, method) + key = method_cases_key[method] + if cfg.force_case === nothing + return "$(key)_$(cfg.model)_ensemble_results_$(cfg.calibrate_date).nc" + else + return "$(key)_$(cfg.model)_$(forcing_cases_key[cfg.force_case])_$(cfg.calibrate_date).nc" + end +end + +######################################################################## +############### ARRAY-JOB HELPERS ################################### +######################################################################## + +# Flattened (N_ens, rng_idx) task list; index t (1-based) == SLURM_ARRAY_TASK_ID. +# Outer loop = N_ens_sizes, inner = repeats, matching the original serial nesting. +flat_tasks(cfg) = + [(N_ens, rng_idx) for N_ens in cfg.N_ens_sizes for rng_idx in 1:cfg.n_repeats] + +# Task index: prefer SLURM_ARRAY_TASK_ID, then first CLI arg, else nothing (run all). +function task_index_from_args() + if haskey(ENV, "SLURM_ARRAY_TASK_ID") + return parse(Int, ENV["SLURM_ARRAY_TASK_ID"]) + elseif !isempty(ARGS) && !isempty(ARGS[1]) + return parse(Int, ARGS[1]) + else + return nothing + end +end + +# L96 experiment: prefer EXPERIMENT env var, then ARGS[2], else the manual toggle above. +function l96_experiment() + if haskey(ENV, "EXPERIMENT") + return Symbol(ENV["EXPERIMENT"]) + elseif length(ARGS) >= 2 && !isempty(ARGS[2]) + return Symbol(ARGS[2]) + else + return EXPERIMENT + end +end diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh new file mode 100755 index 000000000..f324cf1f7 --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Submit calibrate + emulate_sample for the L63 case. +# +# Usage: bash submit_l63.sh [EXP_ID] +# EXP_ID (optional): label appended to SLURM job names so the queue stays +# readable when all four cases run simultaneously, +# e.g. "run2" → jobs appear as "calib_l63_run2". + +set -euo pipefail + +EXP_ID=${1:-} +LABEL="l63${EXP_ID:+_${EXP_ID}}" +DIR="$(cd "$(dirname "$0")" && pwd)" + +cd "$DIR" +mkdir -p output/slurm + +echo "=== Precompiling ===" +julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + +echo "=== Submitting calibrate (L63) ===" +CALIB_JID=$(sbatch --parsable \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l63.jl \ + calibrate_array.sbatch) +echo " calibrate job ID: ${CALIB_JID}" + +echo "=== Submitting emulate_sample (L63, after ${CALIB_JID}) ===" +EMU_JID=$(sbatch --parsable \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l63.jl \ + emulate_sample_array.sbatch) +echo " emulate_sample job ID: ${EMU_JID}" + +echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh new file mode 100755 index 000000000..6f831e5a7 --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Submit calibrate + emulate_sample for the L96 const-force case. +# +# Usage: bash submit_l96_const.sh [EXP_ID] +# EXP_ID (optional): label appended to SLURM job names so the queue stays +# readable when all four cases run simultaneously, +# e.g. "run2" → jobs appear as "calib_l96c_run2". + +set -euo pipefail + +EXP_ID=${1:-} +LABEL="l96c${EXP_ID:+_${EXP_ID}}" +DIR="$(cd "$(dirname "$0")" && pwd)" + +cd "$DIR" +mkdir -p output/slurm + +echo "=== Precompiling ===" +julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + +echo "=== Submitting calibrate (L96 const-force) ===" +CALIB_JID=$(sbatch --parsable \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ + calibrate_array.sbatch) +echo " calibrate job ID: ${CALIB_JID}" + +echo "=== Submitting emulate_sample (L96 const-force, after ${CALIB_JID}) ===" +EMU_JID=$(sbatch --parsable \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ + emulate_sample_array.sbatch) +echo " emulate_sample job ID: ${EMU_JID}" + +echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh new file mode 100755 index 000000000..171020a0a --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Submit calibrate + emulate_sample for the L96 flux-force case. +# +# Usage: bash submit_l96_flux.sh [EXP_ID] +# EXP_ID (optional): label appended to SLURM job names so the queue stays +# readable when all four cases run simultaneously, +# e.g. "run2" → jobs appear as "calib_l96f_run2". + +set -euo pipefail + +EXP_ID=${1:-} +LABEL="l96f${EXP_ID:+_${EXP_ID}}" +DIR="$(cd "$(dirname "$0")" && pwd)" + +cd "$DIR" +mkdir -p output/slurm + +echo "=== Precompiling ===" +julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + +echo "=== Submitting calibrate (L96 flux-force) ===" +CALIB_JID=$(sbatch --parsable \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux \ + calibrate_array.sbatch) +echo " calibrate job ID: ${CALIB_JID}" + +echo "=== Submitting emulate_sample (L96 flux-force, after ${CALIB_JID}) ===" +EMU_JID=$(sbatch --parsable \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ + emulate_sample_array.sbatch) +echo " emulate_sample job ID: ${EMU_JID}" + +echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh new file mode 100755 index 000000000..923bb136a --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# Submit calibrate + emulate_sample for the L96 vec-force case. +# +# Usage: bash submit_l96_vec.sh [EXP_ID] +# EXP_ID (optional): label appended to SLURM job names so the queue stays +# readable when all four cases run simultaneously, +# e.g. "run2" → jobs appear as "calib_l96v_run2". + +set -euo pipefail + +EXP_ID=${1:-} +LABEL="l96v${EXP_ID:+_${EXP_ID}}" +DIR="$(cd "$(dirname "$0")" && pwd)" + +cd "$DIR" +mkdir -p output/slurm + +echo "=== Precompiling ===" +julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + +echo "=== Submitting calibrate (L96 vec-force) ===" +CALIB_JID=$(sbatch --parsable \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ + calibrate_array.sbatch) +echo " calibrate job ID: ${CALIB_JID}" + +echo "=== Submitting emulate_sample (L96 vec-force, after ${CALIB_JID}) ===" +EMU_JID=$(sbatch --parsable \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ + emulate_sample_array.sbatch) +echo " emulate_sample job ID: ${EMU_JID}" + +echo "=== Done. Monitor with: squeue -u \$USER ===" From b8493d6eff945390e36f66209a7c0e60a89162bf Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 12:56:52 -0700 Subject: [PATCH 25/86] add -A esm and julia version --- .../EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- .../hpc-variant/emulate_sample_array.sbatch | 2 +- examples/EKIRace/hpc-variant/submit_l63.sh | 16 +++++++++------- .../EKIRace/hpc-variant/submit_l96_const.sh | 18 ++++++++++-------- .../EKIRace/hpc-variant/submit_l96_flux.sh | 18 ++++++++++-------- examples/EKIRace/hpc-variant/submit_l96_vec.sh | 16 +++++++++------- 6 files changed, 40 insertions(+), 32 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 0392cf232..4165e8a43 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -28,7 +28,7 @@ set -euo pipefail -module load julia # pin version as needed, e.g.: module load julia/1.10.4 +module load julia/1.11.2 # pin version as needed, e.g.: module load julia/1.10.4 export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index ca8788461..545f18029 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -26,7 +26,7 @@ set -euo pipefail -module load julia # pin version as needed, e.g.: module load julia/1.10.4 +module load julia/1.12.2 # pin version as needed, e.g.: module load julia/1.10.4 export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index f324cf1f7..b794b6ca7 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -20,17 +20,19 @@ julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' echo "=== Submitting calibrate (L63) ===" CALIB_JID=$(sbatch --parsable \ - --job-name="calib_${LABEL}" \ - --export=ALL,SCRIPT=calibrate_l63.jl \ - calibrate_array.sbatch) + -A esm \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l63.jl \ + calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L63, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l63.jl \ - emulate_sample_array.sbatch) + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l63.jl \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 6f831e5a7..56eb77840 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -20,17 +20,19 @@ julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' echo "=== Submitting calibrate (L96 const-force) ===" CALIB_JID=$(sbatch --parsable \ - --job-name="calib_${LABEL}" \ - --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ - calibrate_array.sbatch) + -A esm \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ + calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 const-force, after ${CALIB_JID}) ===" -EMU_JID=$(sbatch --parsable \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ - emulate_sample_array.sbatch) +EMU_JID=$(sbatch --parsable \ + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 171020a0a..896944cea 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -20,17 +20,19 @@ julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' echo "=== Submitting calibrate (L96 flux-force) ===" CALIB_JID=$(sbatch --parsable \ - --job-name="calib_${LABEL}" \ - --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux \ - calibrate_array.sbatch) + -A esm \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux \ + calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 flux-force, after ${CALIB_JID}) ===" -EMU_JID=$(sbatch --parsable \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ - emulate_sample_array.sbatch) +EMU_JID=$(sbatch --parsable \ + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 923bb136a..966c408dd 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -20,17 +20,19 @@ julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' echo "=== Submitting calibrate (L96 vec-force) ===" CALIB_JID=$(sbatch --parsable \ - --job-name="calib_${LABEL}" \ - --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ - calibrate_array.sbatch) + -A esm \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ + calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 vec-force, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ - emulate_sample_array.sbatch) + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" From f049e65d215e8550e006930a72dfb967f47dd185 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 13:08:43 -0700 Subject: [PATCH 26/86] whitespace fix --- .../EKIRace/hpc-variant/submit_l96_const.sh | 12 ++++++------ .../EKIRace/hpc-variant/submit_l96_flux.sh | 12 ++++++------ examples/EKIRace/hpc-variant/submit_l96_vec.sh | 18 +++++++++--------- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 56eb77840..bae2be970 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -27,12 +27,12 @@ CALIB_JID=$(sbatch --parsable \ echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 const-force, after ${CALIB_JID}) ===" -EMU_JID=$(sbatch --parsable \ - -A esm \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ - emulate_sample_array.sbatch) +EMU_JID=$(sbatch --parsable \ + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 896944cea..cae234f2d 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -27,12 +27,12 @@ CALIB_JID=$(sbatch --parsable \ echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 flux-force, after ${CALIB_JID}) ===" -EMU_JID=$(sbatch --parsable \ - -A esm \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ - emulate_sample_array.sbatch) +EMU_JID=$(sbatch --parsable \ + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 966c408dd..d3715bde2 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -20,19 +20,19 @@ julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' echo "=== Submitting calibrate (L96 vec-force) ===" CALIB_JID=$(sbatch --parsable \ - -A esm \ - --job-name="calib_${LABEL}" \ - --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ - calibrate_array.sbatch) + -A esm \ + --job-name="calib_${LABEL}" \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ + calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" echo "=== Submitting emulate_sample (L96 vec-force, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ - -A esm \ - --job-name="emu_${LABEL}" \ - --dependency=afterok:${CALIB_JID} \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ - emulate_sample_array.sbatch) + -A esm \ + --job-name="emu_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ + emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" From fd203fbf892271b7a0d9549a40d82747ffd841b9 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 13:18:28 -0700 Subject: [PATCH 27/86] wrong julia --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 4165e8a43..1e4615da4 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -28,7 +28,7 @@ set -euo pipefail -module load julia/1.11.2 # pin version as needed, e.g.: module load julia/1.10.4 +module load julia/1.12.2 # pin version as needed, e.g.: module load julia/1.10.4 export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} From 109d1c61cd05d2095acd524b133473573e66272c Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 13:18:40 -0700 Subject: [PATCH 28/86] kill on invalid deps --- examples/EKIRace/hpc-variant/submit_l63.sh | 1 + examples/EKIRace/hpc-variant/submit_l96_const.sh | 1 + examples/EKIRace/hpc-variant/submit_l96_flux.sh | 1 + examples/EKIRace/hpc-variant/submit_l96_vec.sh | 1 + 4 files changed, 4 insertions(+) diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index b794b6ca7..345f6beaa 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -31,6 +31,7 @@ EMU_JID=$(sbatch --parsable \ -A esm \ --job-name="emu_${LABEL}" \ --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l63.jl \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index bae2be970..0c332a27c 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -31,6 +31,7 @@ EMU_JID=$(sbatch --parsable \ -A esm \ --job-name="emu_${LABEL}" \ --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index cae234f2d..cdd55eda7 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -31,6 +31,7 @@ EMU_JID=$(sbatch --parsable \ -A esm \ --job-name="emu_${LABEL}" \ --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_flux \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index d3715bde2..005716567 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -31,6 +31,7 @@ EMU_JID=$(sbatch --parsable \ -A esm \ --job-name="emu_${LABEL}" \ --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_vec \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" From 1716aa52102b1bf5e30205452480e5d32cbc43ad Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 13:43:37 -0700 Subject: [PATCH 29/86] directory linking --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- examples/EKIRace/hpc-variant/emulate_sample_array.sbatch | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 1e4615da4..20877dfee 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -36,6 +36,6 @@ export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} # Default to L63 if SCRIPT not set via --export. SCRIPT=${SCRIPT:-calibrate_l63.jl} -cd "$(dirname "$0")" +cd "${SLURM_SUBMIT_DIR}" julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index 545f18029..fb3a76f30 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -33,6 +33,6 @@ export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} SCRIPT=${SCRIPT:-emulate_sample_l63.jl} -cd "$(dirname "$0")" +cd "${SLURM_SUBMIT_DIR}" julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" From 05319366044adeac1c2aa1521a6bf06bcdf02775 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 13:49:49 -0700 Subject: [PATCH 30/86] turn off precimpiles for every member - causes breaks --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 ++ examples/EKIRace/hpc-variant/emulate_sample_array.sbatch | 2 ++ 2 files changed, 4 insertions(+) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 20877dfee..3586e9e64 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -38,4 +38,6 @@ SCRIPT=${SCRIPT:-calibrate_l63.jl} cd "${SLURM_SUBMIT_DIR}" +export JULIA_PKG_PRECOMPILE_AUTO=0 # already precompiled so should not get all jobs to re-precompile + julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index fb3a76f30..7f52bf86d 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -35,4 +35,6 @@ SCRIPT=${SCRIPT:-emulate_sample_l63.jl} cd "${SLURM_SUBMIT_DIR}" +export JULIA_PKG_PRECOMPILE_AUTO=0 # already precompiled so should not get all jobs to re-precompile + julia --project=. "${SCRIPT}" "${SLURM_ARRAY_TASK_ID}" From b377ca0792763b818a6147a8d54bf54a626f9651 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 14:05:51 -0700 Subject: [PATCH 31/86] add a precompile script --- examples/EKIRace/hpc-variant/README.md | 54 +++++++++---------- .../EKIRace/hpc-variant/precompile.sbatch | 20 +++++++ examples/EKIRace/hpc-variant/submit_l63.sh | 12 +++-- .../EKIRace/hpc-variant/submit_l96_const.sh | 12 +++-- .../EKIRace/hpc-variant/submit_l96_flux.sh | 12 +++-- .../EKIRace/hpc-variant/submit_l96_vec.sh | 12 +++-- 6 files changed, 81 insertions(+), 41 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/precompile.sbatch diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index 9634a8302..b801b9ba7 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -17,8 +17,8 @@ All subsequent stages (emulate_sample, leaderboard) read this value to locate the right output directory; keeping it fixed avoids mismatches when jobs run past midnight or across days. -The submit scripts (`submit_*.sh`) handle precompilation automatically before -queuing any jobs. +The submit scripts (`submit_*.sh`) submit a `precompile.sbatch` job first and +chain calibrate → emulate_sample behind it, so nothing runs on the login node. ## Standalone (serial) @@ -46,10 +46,11 @@ julia --project=. emulate_sample_l63.jl 5 # fifth cell only ### Submission scripts (recommended) -One script per case handles precompilation and chains calibrate → emulate_sample -automatically. All four can be launched simultaneously — output files are -case-specific so there are no write conflicts. The optional `EXP_ID` argument -suffixes SLURM job names to keep the queue readable: +One script per case submits three chained SLURM jobs: +`precompile.sbatch` → `calibrate_array.sbatch` → `emulate_sample_array.sbatch`. +All four cases can be launched simultaneously — output files are case-specific +so there are no write conflicts. The optional `EXP_ID` argument suffixes SLURM +job names to keep the queue readable: ```bash bash submit_l63.sh [EXP_ID] @@ -66,38 +67,33 @@ wait ### Manual submission -```bash -# L63 -sbatch --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch - -# L96 (submit once per forcing case) -sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch -sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec calibrate_array.sbatch -sbatch --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux calibrate_array.sbatch -``` - -Each task writes its per-method `ekp` and `results` files, which are consumed -directly by the emulate_sample stage. - -### Emulate-sample - -Submit after the calibrate array is complete. Use `--dependency` to chain -automatically: +Precompile first, then chain calibrate and emulate_sample behind it: ```bash +# precompile (shared across all cases — only one run needed per environment change) +pid=$(sbatch --parsable -A esm precompile.sbatch) + # L63 -cid=$(sbatch --parsable --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) -sbatch --dependency=afterok:$cid \ +cid=$(sbatch --parsable -A esm \ + --dependency=afterok:$pid --kill-on-invalid-dep=yes \ + --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) +sbatch -A esm \ + --dependency=afterok:$cid --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch -# L96 -cid=$(sbatch --parsable --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ - calibrate_array.sbatch) -sbatch --dependency=afterok:$cid \ +# L96 (repeat for each forcing case, reusing the same $pid) +cid=$(sbatch --parsable -A esm \ + --dependency=afterok:$pid --kill-on-invalid-dep=yes \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch) +sbatch -A esm \ + --dependency=afterok:$cid --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ emulate_sample_array.sbatch ``` +Each task writes its per-method `ekp` and `results` files, which are consumed +directly by the emulate_sample stage. + ### Adjusting array size The sbatch files default to `--array=1-60` (3 ensemble sizes × 20 repeats). diff --git a/examples/EKIRace/hpc-variant/precompile.sbatch b/examples/EKIRace/hpc-variant/precompile.sbatch new file mode 100644 index 000000000..9dfe4b49c --- /dev/null +++ b/examples/EKIRace/hpc-variant/precompile.sbatch @@ -0,0 +1,20 @@ +#!/bin/bash +# Single SLURM job that instantiates and precompiles the project. +# The submit_*.sh scripts depend on this job before launching array jobs, +# so that 60 array tasks never race to precompile simultaneously. + +#SBATCH --job-name=precompile +#SBATCH --output=output/slurm/precompile_%j.out +#SBATCH --error=output/slurm/precompile_%j.err +#SBATCH --time=00:30:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=1 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 # keep in sync with calibrate_array.sbatch + +cd "${SLURM_SUBMIT_DIR}" + +julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 345f6beaa..f276685e4 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -15,13 +15,19 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Precompiling ===" -julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' +echo "=== Submitting precompile ===" +PRECOMPILE_JID=$(sbatch --parsable \ + -A esm \ + --job-name="precompile_${LABEL}" \ + precompile.sbatch) +echo " precompile job ID: ${PRECOMPILE_JID}" -echo "=== Submitting calibrate (L63) ===" +echo "=== Submitting calibrate (L63, after ${PRECOMPILE_JID}) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ + --dependency=afterok:${PRECOMPILE_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l63.jl \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 0c332a27c..3f04132a3 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -15,13 +15,19 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Precompiling ===" -julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' +echo "=== Submitting precompile ===" +PRECOMPILE_JID=$(sbatch --parsable \ + -A esm \ + --job-name="precompile_${LABEL}" \ + precompile.sbatch) +echo " precompile job ID: ${PRECOMPILE_JID}" -echo "=== Submitting calibrate (L96 const-force) ===" +echo "=== Submitting calibrate (L96 const-force, after ${PRECOMPILE_JID}) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ + --dependency=afterok:${PRECOMPILE_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index cdd55eda7..4a7805f0d 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -15,13 +15,19 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Precompiling ===" -julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' +echo "=== Submitting precompile ===" +PRECOMPILE_JID=$(sbatch --parsable \ + -A esm \ + --job-name="precompile_${LABEL}" \ + precompile.sbatch) +echo " precompile job ID: ${PRECOMPILE_JID}" -echo "=== Submitting calibrate (L96 flux-force) ===" +echo "=== Submitting calibrate (L96 flux-force, after ${PRECOMPILE_JID}) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ + --dependency=afterok:${PRECOMPILE_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 005716567..18ee7f36d 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -15,13 +15,19 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Precompiling ===" -julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' +echo "=== Submitting precompile ===" +PRECOMPILE_JID=$(sbatch --parsable \ + -A esm \ + --job-name="precompile_${LABEL}" \ + precompile.sbatch) +echo " precompile job ID: ${PRECOMPILE_JID}" -echo "=== Submitting calibrate (L96 vec-force) ===" +echo "=== Submitting calibrate (L96 vec-force, after ${PRECOMPILE_JID}) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ + --dependency=afterok:${PRECOMPILE_JID} \ + --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" From ec66c6583635645fa41480838fbb52fb61de6b38 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 14:15:56 -0700 Subject: [PATCH 32/86] add precompile script to be run in advance of submit_l* scripts --- examples/EKIRace/hpc-variant/README.md | 36 +++++++++++-------- .../EKIRace/hpc-variant/precompile.sbatch | 8 +++-- examples/EKIRace/hpc-variant/submit_l63.sh | 12 ++----- .../EKIRace/hpc-variant/submit_l96_const.sh | 12 ++----- .../EKIRace/hpc-variant/submit_l96_flux.sh | 12 ++----- .../EKIRace/hpc-variant/submit_l96_vec.sh | 12 ++----- .../EKIRace/hpc-variant/submit_precompile.sh | 24 +++++++++++++ 7 files changed, 62 insertions(+), 54 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/submit_precompile.sh diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index b801b9ba7..05ba76478 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -17,8 +17,11 @@ All subsequent stages (emulate_sample, leaderboard) read this value to locate the right output directory; keeping it fixed avoids mismatches when jobs run past midnight or across days. -The submit scripts (`submit_*.sh`) submit a `precompile.sbatch` job first and -chain calibrate → emulate_sample behind it, so nothing runs on the login node. +Precompilation is handled by a dedicated `submit_precompile.sh` script that +queues `precompile.sbatch` as a compute job. Run it once before your experiments +and again whenever the environment changes (fresh checkout, package updates). +The `submit_l*.sh` scripts do not precompile — they will remind you at submission +time. ## Standalone (serial) @@ -46,19 +49,26 @@ julia --project=. emulate_sample_l63.jl 5 # fifth cell only ### Submission scripts (recommended) -One script per case submits three chained SLURM jobs: -`precompile.sbatch` → `calibrate_array.sbatch` → `emulate_sample_array.sbatch`. -All four cases can be launched simultaneously — output files are case-specific -so there are no write conflicts. The optional `EXP_ID` argument suffixes SLURM -job names to keep the queue readable: +Precompile once (or whenever the environment changes), then submit the cases: ```bash +bash submit_precompile.sh [EXP_ID] + bash submit_l63.sh [EXP_ID] bash submit_l96_const.sh [EXP_ID] bash submit_l96_vec.sh [EXP_ID] bash submit_l96_flux.sh [EXP_ID] +``` + +Each `submit_l*.sh` script chains `calibrate_array.sbatch` → `emulate_sample_array.sbatch` +for its case. All four cases can be launched simultaneously — output files are +case-specific so there are no write conflicts. The optional `EXP_ID` argument +suffixes SLURM job names to keep the queue readable. -# example: run all four with a shared label +If you are launching all four cases together you only need one precompile run: + +```bash +bash submit_precompile.sh run1 for s in submit_l63.sh submit_l96_const.sh submit_l96_vec.sh submit_l96_flux.sh; do bash "$s" run1 & done @@ -67,23 +77,19 @@ wait ### Manual submission -Precompile first, then chain calibrate and emulate_sample behind it: +Precompile via `submit_precompile.sh` (or directly), then submit calibrate and +emulate_sample: ```bash -# precompile (shared across all cases — only one run needed per environment change) -pid=$(sbatch --parsable -A esm precompile.sbatch) - # L63 cid=$(sbatch --parsable -A esm \ - --dependency=afterok:$pid --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) sbatch -A esm \ --dependency=afterok:$cid --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch -# L96 (repeat for each forcing case, reusing the same $pid) +# L96 cid=$(sbatch --parsable -A esm \ - --dependency=afterok:$pid --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch) sbatch -A esm \ --dependency=afterok:$cid --kill-on-invalid-dep=yes \ diff --git a/examples/EKIRace/hpc-variant/precompile.sbatch b/examples/EKIRace/hpc-variant/precompile.sbatch index 9dfe4b49c..f48c4bebf 100644 --- a/examples/EKIRace/hpc-variant/precompile.sbatch +++ b/examples/EKIRace/hpc-variant/precompile.sbatch @@ -1,14 +1,14 @@ #!/bin/bash # Single SLURM job that instantiates and precompiles the project. -# The submit_*.sh scripts depend on this job before launching array jobs, -# so that 60 array tasks never race to precompile simultaneously. +# Run via submit_precompile.sh before launching array jobs so that the 60+ +# array tasks start with a warm package cache and never race to precompile. #SBATCH --job-name=precompile #SBATCH --output=output/slurm/precompile_%j.out #SBATCH --error=output/slurm/precompile_%j.err #SBATCH --time=00:30:00 #SBATCH --mem=8G -#SBATCH --cpus-per-task=1 +#SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 set -euo pipefail @@ -17,4 +17,6 @@ module load julia/1.12.2 # keep in sync with calibrate_array.sbatch cd "${SLURM_SUBMIT_DIR}" +export JULIA_NUM_THREADS=$SLURM_CPUS_PER_TASK # parallelize + julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index f276685e4..ae0bc1365 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -15,19 +15,13 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Submitting precompile ===" -PRECOMPILE_JID=$(sbatch --parsable \ - -A esm \ - --job-name="precompile_${LABEL}" \ - precompile.sbatch) -echo " precompile job ID: ${PRECOMPILE_JID}" +echo "NOTE: This script does not precompile. Run bash submit_precompile.sh first" +echo " if you haven't done so recently (e.g. after a fresh checkout or package update)." -echo "=== Submitting calibrate (L63, after ${PRECOMPILE_JID}) ===" +echo "=== Submitting calibrate (L63) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ - --dependency=afterok:${PRECOMPILE_JID} \ - --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l63.jl \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 3f04132a3..30fc9e80f 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -15,19 +15,13 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Submitting precompile ===" -PRECOMPILE_JID=$(sbatch --parsable \ - -A esm \ - --job-name="precompile_${LABEL}" \ - precompile.sbatch) -echo " precompile job ID: ${PRECOMPILE_JID}" +echo "NOTE: This script does not precompile. Run bash submit_precompile.sh first" +echo " if you haven't done so recently (e.g. after a fresh checkout or package update)." -echo "=== Submitting calibrate (L96 const-force, after ${PRECOMPILE_JID}) ===" +echo "=== Submitting calibrate (L96 const-force) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ - --dependency=afterok:${PRECOMPILE_JID} \ - --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 4a7805f0d..867d9ae72 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -15,19 +15,13 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Submitting precompile ===" -PRECOMPILE_JID=$(sbatch --parsable \ - -A esm \ - --job-name="precompile_${LABEL}" \ - precompile.sbatch) -echo " precompile job ID: ${PRECOMPILE_JID}" +echo "NOTE: This script does not precompile. Run bash submit_precompile.sh first" +echo " if you haven't done so recently (e.g. after a fresh checkout or package update)." -echo "=== Submitting calibrate (L96 flux-force, after ${PRECOMPILE_JID}) ===" +echo "=== Submitting calibrate (L96 flux-force) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ - --dependency=afterok:${PRECOMPILE_JID} \ - --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_flux \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 18ee7f36d..289154ac0 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -15,19 +15,13 @@ DIR="$(cd "$(dirname "$0")" && pwd)" cd "$DIR" mkdir -p output/slurm -echo "=== Submitting precompile ===" -PRECOMPILE_JID=$(sbatch --parsable \ - -A esm \ - --job-name="precompile_${LABEL}" \ - precompile.sbatch) -echo " precompile job ID: ${PRECOMPILE_JID}" +echo "NOTE: This script does not precompile. Run bash submit_precompile.sh first" +echo " if you haven't done so recently (e.g. after a fresh checkout or package update)." -echo "=== Submitting calibrate (L96 vec-force, after ${PRECOMPILE_JID}) ===" +echo "=== Submitting calibrate (L96 vec-force) ===" CALIB_JID=$(sbatch --parsable \ -A esm \ --job-name="calib_${LABEL}" \ - --dependency=afterok:${PRECOMPILE_JID} \ - --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_vec \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_precompile.sh b/examples/EKIRace/hpc-variant/submit_precompile.sh new file mode 100644 index 000000000..eebbf4333 --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_precompile.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Submit the one-off precompile job. +# Run this once before (or instead of re-running) the submit_l*.sh scripts +# whenever the environment changes (new package versions, fresh checkout, etc.). +# +# Usage: bash submit_precompile.sh [EXP_ID] +# EXP_ID (optional): label appended to the SLURM job name. + +set -euo pipefail + +EXP_ID=${1:-} +LABEL="precompile${EXP_ID:+_${EXP_ID}}" +DIR="$(cd "$(dirname "$0")" && pwd)" + +cd "$DIR" +mkdir -p output/slurm + +echo "=== Submitting precompile ===" +PRECOMPILE_JID=$(sbatch --parsable \ + -A esm \ + --job-name="${LABEL}" \ + precompile.sbatch) +echo " precompile job ID: ${PRECOMPILE_JID}" +echo "=== Done. Monitor with: squeue -u \$USER ===" From d7cac06d26e8927a14f81b98efed1ecb4765dfd1 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 14:18:46 -0700 Subject: [PATCH 33/86] make precomp executeable --- examples/EKIRace/hpc-variant/submit_precompile.sh | 0 1 file changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 examples/EKIRace/hpc-variant/submit_precompile.sh diff --git a/examples/EKIRace/hpc-variant/submit_precompile.sh b/examples/EKIRace/hpc-variant/submit_precompile.sh old mode 100644 new mode 100755 From 40f44399cc4a418a3b53d9010ccdfc3dcc1ba092 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 14:29:55 -0700 Subject: [PATCH 34/86] typo --- examples/EKIRace/hpc-variant/calibrate_l96.jl | 6 ++++-- examples/EKIRace/hpc-variant/emulate_sample_l96.jl | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index 43afe35d6..e088d6eee 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -306,8 +306,10 @@ end function main() exp = l96_experiment() - @assert exp in (:l96_const, :l96_vec, :l96_flux) \ - "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)" + @assert( + exp in (:l96_const, :l96_vec, :l96_flux), + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)", + ) cfg = experiment_config(exp) output_dir = joinpath(@__DIR__, "output") diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl index ab8507f6c..4c3fcd5f9 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl +++ b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl @@ -158,8 +158,10 @@ end function main() exp = l96_experiment() - @assert exp in (:l96_const, :l96_vec, :l96_flux) \ - "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)" + @assert( + exp in (:l96_const, :l96_vec, :l96_flux), + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)", + ) cfg = experiment_config(exp) tasks = flat_tasks(cfg) idx = task_index_from_args() From ae86c2699d5248be0c911044024adb88b710fc34 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 14:29:31 -0700 Subject: [PATCH 35/86] increase concurrency --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- examples/EKIRace/hpc-variant/emulate_sample_array.sbatch | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 3586e9e64..f8265cbb1 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -20,7 +20,7 @@ #SBATCH --job-name=calib #SBATCH --output=output/slurm/calib_%A_%a.out #SBATCH --error=output/slurm/calib_%A_%a.err -#SBATCH --array=1-60%20 +#SBATCH --array=1-60%100 #SBATCH --time=04:00:00 #SBATCH --mem=8G #SBATCH --cpus-per-task=4 diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index 7f52bf86d..e4dd99503 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -18,7 +18,7 @@ #SBATCH --job-name=emusample #SBATCH --output=output/slurm/emusample_%A_%a.out #SBATCH --error=output/slurm/emusample_%A_%a.err -#SBATCH --array=1-60%20 +#SBATCH --array=1-60%100 #SBATCH --time=12:00:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 From 57b7df91e2a16268a520e8f7a7cd6cce9be54dea Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 15:23:46 -0700 Subject: [PATCH 36/86] remove race condition for calibrate --- examples/EKIRace/hpc-variant/calibrate_l63.jl | 8 +++++++- examples/EKIRace/hpc-variant/calibrate_l96.jl | 20 ++++++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl index 53538adb5..a0686dc33 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l63.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -101,7 +101,13 @@ function write_priors(cfg, setup, output_dir) mkpath(per_method_dir) pf = joinpath(per_method_dir, prior_filename(cfg)) if !isfile(pf) - JLD2.save(pf, "prior", setup.prior) + pf_tmp = pf * ".tmp.$(getpid())" + JLD2.save(pf_tmp, "prior", setup.prior) + try + mv(pf_tmp, pf) + catch + rm(pf_tmp; force = true) + end end end end diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index e088d6eee..f36a9be9c 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -134,15 +134,23 @@ function build_setup(cfg, output_dir) cov_solve = lorenz_solve(phi, x0, LorenzConfig(t, covT)) ic_cov = 0.1 * cov(cov_solve, dims = 2) ic_cov_sqrt = sqrt(ic_cov) + prelim_tmp = prelim_file * ".tmp.$(getpid())" JLD2.save( - prelim_file, + prelim_tmp, "x0", x0, "y", y, "lorenz_config_settings", lorenz_config_settings, "observation_config", observation_config, "ic_cov_sqrt", ic_cov_sqrt, "R", R, "R_inv_var", R_inv_var, ) - @info "Saved computed quantities to $(prelim_file)" + # Atomic rename: first writer wins; concurrent losers discard their copy. + try + mv(prelim_tmp, prelim_file) + @info "Saved computed quantities to $(prelim_file)" + catch + rm(prelim_tmp; force = true) + @info "Prelim file already written by another task; discarding duplicate" + end end configuration = Dict( @@ -170,7 +178,13 @@ function write_priors(cfg, setup, output_dir) mkpath(per_method_dir) pf = joinpath(per_method_dir, prior_filename(cfg)) if !isfile(pf) - JLD2.save(pf, "prior", setup.prior) + pf_tmp = pf * ".tmp.$(getpid())" + JLD2.save(pf_tmp, "prior", setup.prior) + try + mv(pf_tmp, pf) + catch + rm(pf_tmp; force = true) + end end end end From d5977bfc7e098f286c66c25bb5409803d66040b1 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 18:09:31 -0700 Subject: [PATCH 37/86] write-priors race --- examples/EKIRace/hpc-variant/calibrate_l63.jl | 2 +- examples/EKIRace/hpc-variant/calibrate_l96.jl | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl index a0686dc33..f44a0b706 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l63.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -101,7 +101,7 @@ function write_priors(cfg, setup, output_dir) mkpath(per_method_dir) pf = joinpath(per_method_dir, prior_filename(cfg)) if !isfile(pf) - pf_tmp = pf * ".tmp.$(getpid())" + pf_tmp = splitext(pf)[1] * ".tmp.$(getpid()).jld2" JLD2.save(pf_tmp, "prior", setup.prior) try mv(pf_tmp, pf) diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index f36a9be9c..fef4e5a05 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -134,7 +134,7 @@ function build_setup(cfg, output_dir) cov_solve = lorenz_solve(phi, x0, LorenzConfig(t, covT)) ic_cov = 0.1 * cov(cov_solve, dims = 2) ic_cov_sqrt = sqrt(ic_cov) - prelim_tmp = prelim_file * ".tmp.$(getpid())" + prelim_tmp = splitext(prelim_file)[1] * ".tmp.$(getpid()).jld2" JLD2.save( prelim_tmp, "x0", x0, "y", y, @@ -178,7 +178,7 @@ function write_priors(cfg, setup, output_dir) mkpath(per_method_dir) pf = joinpath(per_method_dir, prior_filename(cfg)) if !isfile(pf) - pf_tmp = pf * ".tmp.$(getpid())" + pf_tmp = splitext(pf)[1] * ".tmp.$(getpid()).jld2" JLD2.save(pf_tmp, "prior", setup.prior) try mv(pf_tmp, pf) From 48689c958ea04ef726f9d6a4628766481988e3c8 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 22:57:51 -0700 Subject: [PATCH 38/86] memory --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index f8265cbb1..36973b7ce 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -22,7 +22,7 @@ #SBATCH --error=output/slurm/calib_%A_%a.err #SBATCH --array=1-60%100 #SBATCH --time=04:00:00 -#SBATCH --mem=8G +#SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 From 856ac753952c0781125d3f23d9d032ff842b195e Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 22:59:05 -0700 Subject: [PATCH 39/86] add comp_metrics --- .../compute_leaderboard_metrics.jl | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl new file mode 100644 index 000000000..b17711eac --- /dev/null +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -0,0 +1,32 @@ +using NCDatasets +using Distributions +using Statistics + +filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" + +# load +ncd = NCDataset(filename) +mh = ncd[:mahalanobis] +lp = ncd[:posterior_logpdf_true_v_map] +# Gaussian: mh=-2lp +(n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) +# assume dimnames(mh) = "random_seed", "ensemble_size", "k_iter" + +# count missings +mh1_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp1_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] + +# criteria of "good" posterior +qq_good = quantile(Chisq(n_par), [0.1, 0.5, 0.9]) +# now we compare over the rng dimension +mh_scores = fill(NaN, length(qq_good), size(mh,2), size(mh,3)) +lp_scores = fill(NaN, length(qq_good), size(lp,2), size(lp,3)) +for (idx,qg) in enumerate(qq_good) + mh_scores[idx,:,:] = [sum(skipmissing(mh[:,i,j]) .<= qg; init=0) for i in axes(mh,2), j in axes(mh,3)] + lp_scores[idx,:,:] = [sum(skipmissing(-2*lp[:,i,j]) .<= qg; init=0) for i in axes(lp,2), j in axes(lp,3)] + # divide by the # non-missing values. In case of all missing we divide by 1 + mh_scores[idx,:,:] ./= max.(mh1_nonmiss,1) + lp_scores[idx,:,:] ./= max.(lp1_nonmiss,1) +end + + From 22419cdc0a0ed5e3e124bb4de1478cafd8f53443 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 30 May 2026 23:14:25 -0700 Subject: [PATCH 40/86] copy postprocessing --- .../compute_leaderboard_metrics.jl | 5 +- .../EKIRace/hpc-variant/experiment_config.jl | 5 +- .../l63_exp_to_leaderboard_utilities.jl | 181 +++++++++++++++++ .../l96_exp_to_leaderboard_utilities.jl | 183 ++++++++++++++++++ 4 files changed, 371 insertions(+), 3 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl create mode 100644 examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index b17711eac..7ccb98191 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -2,7 +2,10 @@ using NCDatasets using Distributions using Statistics -filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" +#filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" +#filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" +#filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" +filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-29.nc" # load ncd = NCDataset(filename) diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl index 308d0e195..78bb43f2f 100644 --- a/examples/EKIRace/hpc-variant/experiment_config.jl +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -3,8 +3,9 @@ using Dates ######################################################################## ############### USER TOGGLE ######################################### ######################################################################## -# Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux -EXPERIMENT = :l96_const +# Set EXPERIMENT to one of: +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[4] # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). diff --git a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl new file mode 100644 index 000000000..d9a64582f --- /dev/null +++ b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl @@ -0,0 +1,181 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") + +# Step 1, Read and extract the available experiments + +#### CHOOSE YOUR CASE: +cfg = experiment_config(:l63) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +nc_save_filename = nc_filename(cfg, method) + +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l63.jl first.") +end + +### Load data + +# Load first valid file to determine parameter dimension and max k +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) + +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + # Load ekpobj for total calibration cost + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens + + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + end +end + + +### Saving the data in nc + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) + +# Coordinate variables +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + +# Data variables + +n_evals_v = defVar( + ds, "n_evals_to_target", Float64, + ("random_seed", "ensemble_size"); + fillvalue = NaN, +) +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." +n_evals_v[:, :] = n_evals_arr + +true_param_v = defVar( + ds, "true_param", Float64, + ("random_seed", "ensemble_size", "param_dim"); + fillvalue = NaN, +) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :] = true_param_arr + +post_mean_v = defVar( + ds, "post_mean", Float64, + ("random_seed", "ensemble_size", "k_iter", "param_dim"); + fillvalue = NaN, +) +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :] = post_mean_arr + +post_cov_v = defVar( + ds, "post_cov", Float64, + ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); + fillvalue = NaN, +) +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar( + ds, "mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" +mahalanobis_v[:, :, :] = mahal_arr + +posterior_logpdf_true_v_map_v = defVar( + ds, "posterior_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl new file mode 100644 index 000000000..a69532fe7 --- /dev/null +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -0,0 +1,183 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") + +# Step 1, Read and extract the available experiments + +#### CHOOSE YOUR CASE: +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For l96_exp_to_leaderboard_utilities.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +force_case = cfg.force_case +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +nc_save_filename = nc_filename(cfg, method) + +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") +end + +### Load data + +# Load first valid file to determine parameter dimension and max k +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) + +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + # Load ekpobj for total calibration cost + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens + + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + # (m - truth)' C^{-1} (m - truth) + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + # log p(truth | posterior) - log p(mode | posterior) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + end +end + + +### Saving the data in nc + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) + +# Coordinate variables +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + +# Data variables + +n_evals_v = defVar( + ds, "n_evals_to_target", Float64, + ("random_seed", "ensemble_size"); + fillvalue = NaN, +) +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." +n_evals_v[:, :] = n_evals_arr + +true_param_v = defVar( + ds, "true_param", Float64, + ("random_seed", "ensemble_size", "param_dim"); + fillvalue = NaN, +) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :] = true_param_arr + +post_mean_v = defVar( + ds, "post_mean", Float64, + ("random_seed", "ensemble_size", "k_iter", "param_dim"); + fillvalue = NaN, +) +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :] = post_mean_arr + +post_cov_v = defVar( + ds, "post_cov", Float64, + ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); + fillvalue = NaN, +) +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar( + ds, "mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +mahalanobis_v.attrib["description"] = "M = mahalanobis distance: for posterior ≈ N(m,C), it is (m-truth_param)'*C^{-1}*(m-truth_param). Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It is the multidimensional equivalent of `how many standard deviations is the truth away form the posterior mean`. For gaussian -2P = M, if M > -2P => heavy tails" +mahalanobis_v[:, :, :] = mahal_arr + +posterior_logpdf_true_v_map_v = defVar( + ds, "posterior_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" From c09f5d6cbd6d982342684ac6d52bd16d0eba0513 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 00:20:49 -0700 Subject: [PATCH 41/86] update metric display --- examples/EKIRace/Project.toml | 1 + .../EKIRace/compute_leaderboard_metrics.jl | 112 ++++++++++++++---- examples/EKIRace/hpc-variant/Project.toml | 1 + .../compute_leaderboard_metrics.jl | 107 ++++++++++++++--- 4 files changed, 182 insertions(+), 39 deletions(-) diff --git a/examples/EKIRace/Project.toml b/examples/EKIRace/Project.toml index 39c1c6f84..ae607c4cf 100644 --- a/examples/EKIRace/Project.toml +++ b/examples/EKIRace/Project.toml @@ -1,6 +1,7 @@ [deps] BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl index b17711eac..61de75089 100644 --- a/examples/EKIRace/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -1,32 +1,104 @@ using NCDatasets using Distributions using Statistics +using Printf +using DataFrames -filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" +indir = joinpath("hpc-variant","output","from-hpc_2026-05-29") +#filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" +#filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" +#filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" +filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-29.nc" + +filename = joinpath(indir, filename) +@info "computing leaderboard metrics from $(filename)" -# load ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] -# Gaussian: mh=-2lp +# Gaussian: mh = -2lp (n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) -# assume dimnames(mh) = "random_seed", "ensemble_size", "k_iter" - -# count missings -mh1_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] -lp1_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] - -# criteria of "good" posterior -qq_good = quantile(Chisq(n_par), [0.1, 0.5, 0.9]) -# now we compare over the rng dimension -mh_scores = fill(NaN, length(qq_good), size(mh,2), size(mh,3)) -lp_scores = fill(NaN, length(qq_good), size(lp,2), size(lp,3)) -for (idx,qg) in enumerate(qq_good) - mh_scores[idx,:,:] = [sum(skipmissing(mh[:,i,j]) .<= qg; init=0) for i in axes(mh,2), j in axes(mh,3)] - lp_scores[idx,:,:] = [sum(skipmissing(-2*lp[:,i,j]) .<= qg; init=0) for i in axes(lp,2), j in axes(lp,3)] - # divide by the # non-missing values. In case of all missing we divide by 1 - mh_scores[idx,:,:] ./= max.(mh1_nonmiss,1) - lp_scores[idx,:,:] ./= max.(lp1_nonmiss,1) + +mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] + +qq_probs = [0.1, 0.5, 0.9] +qq_good = quantile(Chisq(n_par), qq_probs) + +# === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === +# If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. +mh_pool = collect(skipmissing(vec(Array(mh)))) +lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) + +mh_emp_q = quantile(mh_pool, qq_probs) +lp_emp_q = quantile(lp_pool, qq_probs) + +println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") +println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) +for (i, p) in enumerate(qq_probs) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) end +# === 2. Per-experiment scores === +# missing where no valid seeds; * in printout where coverage < n_rng (less trusted) +mh_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) +lp_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) + +for (idx, qg) in enumerate(qq_good) + for j in axes(mh, 3), i in axes(mh, 2) + nm = mh_nonmiss[i, j] + nm == 0 && continue + mh_scores[idx, i, j] = sum(skipmissing(mh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(lp, 3), i in axes(lp, 2) + nm = lp_nonmiss[i, j] + nm == 0 && continue + lp_scores[idx, i, j] = sum(skipmissing(-2 .* lp[:, i, j]) .<= qg; init=0) / nm + end +end + +mh_coverage = mh_nonmiss ./ n_rng # 1.0 = all n_rng seeds valid; < 1 = partial +lp_coverage = lp_nonmiss ./ n_rng + +# === 3. Store scores as DataFrame === +ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end +kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end + +df_scores = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mh_score = Union{Float64,Missing}[], + lp_score = Union{Float64,Missing}[], + mh_coverage = Float64[], + lp_coverage = Float64[], +) + +for (qi, p) in enumerate(qq_probs) + for (ei, ev) in enumerate(ens_vals) + for (ki, kv) in enumerate(kiter_vals) + push!(df_scores, ( + q_prob = p, + ens_size = ev, + k_iter = kv, + mh_score = mh_scores[qi, ei, ki], + lp_score = lp_scores[qi, ei, ki], + mh_coverage = mh_coverage[ei, ki], + lp_coverage = lp_coverage[ei, ki], + )) + end + end +end + +# rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) +println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") +for ens in sort(unique(df_scores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(qq_probs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2))]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end +end diff --git a/examples/EKIRace/hpc-variant/Project.toml b/examples/EKIRace/hpc-variant/Project.toml index 39c1c6f84..ae607c4cf 100644 --- a/examples/EKIRace/hpc-variant/Project.toml +++ b/examples/EKIRace/hpc-variant/Project.toml @@ -1,6 +1,7 @@ [deps] BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0" CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index 7ccb98191..61de75089 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -1,35 +1,104 @@ using NCDatasets using Distributions using Statistics +using Printf +using DataFrames +indir = joinpath("hpc-variant","output","from-hpc_2026-05-29") #filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" #filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" #filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-29.nc" -# load +filename = joinpath(indir, filename) +@info "computing leaderboard metrics from $(filename)" + ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] -# Gaussian: mh=-2lp +# Gaussian: mh = -2lp (n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) -# assume dimnames(mh) = "random_seed", "ensemble_size", "k_iter" - -# count missings -mh1_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] -lp1_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] - -# criteria of "good" posterior -qq_good = quantile(Chisq(n_par), [0.1, 0.5, 0.9]) -# now we compare over the rng dimension -mh_scores = fill(NaN, length(qq_good), size(mh,2), size(mh,3)) -lp_scores = fill(NaN, length(qq_good), size(lp,2), size(lp,3)) -for (idx,qg) in enumerate(qq_good) - mh_scores[idx,:,:] = [sum(skipmissing(mh[:,i,j]) .<= qg; init=0) for i in axes(mh,2), j in axes(mh,3)] - lp_scores[idx,:,:] = [sum(skipmissing(-2*lp[:,i,j]) .<= qg; init=0) for i in axes(lp,2), j in axes(lp,3)] - # divide by the # non-missing values. In case of all missing we divide by 1 - mh_scores[idx,:,:] ./= max.(mh1_nonmiss,1) - lp_scores[idx,:,:] ./= max.(lp1_nonmiss,1) + +mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] + +qq_probs = [0.1, 0.5, 0.9] +qq_good = quantile(Chisq(n_par), qq_probs) + +# === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === +# If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. +mh_pool = collect(skipmissing(vec(Array(mh)))) +lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) + +mh_emp_q = quantile(mh_pool, qq_probs) +lp_emp_q = quantile(lp_pool, qq_probs) + +println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") +println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) +for (i, p) in enumerate(qq_probs) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) end +# === 2. Per-experiment scores === +# missing where no valid seeds; * in printout where coverage < n_rng (less trusted) +mh_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) +lp_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) + +for (idx, qg) in enumerate(qq_good) + for j in axes(mh, 3), i in axes(mh, 2) + nm = mh_nonmiss[i, j] + nm == 0 && continue + mh_scores[idx, i, j] = sum(skipmissing(mh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(lp, 3), i in axes(lp, 2) + nm = lp_nonmiss[i, j] + nm == 0 && continue + lp_scores[idx, i, j] = sum(skipmissing(-2 .* lp[:, i, j]) .<= qg; init=0) / nm + end +end + +mh_coverage = mh_nonmiss ./ n_rng # 1.0 = all n_rng seeds valid; < 1 = partial +lp_coverage = lp_nonmiss ./ n_rng + +# === 3. Store scores as DataFrame === +ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end +kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end + +df_scores = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mh_score = Union{Float64,Missing}[], + lp_score = Union{Float64,Missing}[], + mh_coverage = Float64[], + lp_coverage = Float64[], +) + +for (qi, p) in enumerate(qq_probs) + for (ei, ev) in enumerate(ens_vals) + for (ki, kv) in enumerate(kiter_vals) + push!(df_scores, ( + q_prob = p, + ens_size = ev, + k_iter = kv, + mh_score = mh_scores[qi, ei, ki], + lp_score = lp_scores[qi, ei, ki], + mh_coverage = mh_coverage[ei, ki], + lp_coverage = lp_coverage[ei, ki], + )) + end + end +end + +# rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) +println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") +for ens in sort(unique(df_scores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(qq_probs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2))]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end +end From f27624a4f29c56ba7c740869a4ca9ad86a046983 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 01:00:45 -0700 Subject: [PATCH 42/86] consistent plot ribbons --- .../EKIRace/hpc-variant/plot_l96_forcing.jl | 235 +++++++++ examples/EKIRace/plot_l96_forcing.jl | 483 +++++++----------- 2 files changed, 430 insertions(+), 288 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/plot_l96_forcing.jl diff --git a/examples/EKIRace/hpc-variant/plot_l96_forcing.jl b/examples/EKIRace/hpc-variant/plot_l96_forcing.jl new file mode 100644 index 000000000..fb763383f --- /dev/null +++ b/examples/EKIRace/hpc-variant/plot_l96_forcing.jl @@ -0,0 +1,235 @@ +# load packages +# CES + +using Random +using JLD2 +using Statistics +using LinearAlgebra +using Dates + +using Plots +using Plots.Measures + +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions +const EKP = EnsembleKalmanProcesses + +include("experiment_config.jl") + +#### CHOOSE YOUR CASE: +exp = l96_experiment() +@assert exp in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(exp) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +force_cases = [cfg.force_case] +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +homedir = pwd() + +valid_file_items = [] +valid_files = [] +for force_case in force_cases + for N_ens in N_enss + for rng_idx in rng_idxs + data_save_directory = joinpath(homedir, "output", calib_dir) + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (force_case, N_ens, rng_idx)) + end + end + end +end +@info "Plotting valid files:" +display(valid_files) +println(" ") + +for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) + + @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + # load + data_save_directory = joinpath(homedir, "output", calib_dir) + figure_save_directory = data_save_directory + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + truth_params = load(data_file)["truth_params_constrained"] + final_params = load(data_file)["final_params_constrained"] + posteriors_by_k = load(data_file)["posteriors_by_k"] + k_values = load(data_file)["k_values"] + prior = load(data_file)["priors"] + + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + ekpobj = load(ekp_file)["ekpobj"] + N_ens = get_N_ens(ekpobj) + nx = length(truth_params) + y = get_obs(ekpobj) + ny = length(y) + + # plot - EKI results + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + truth_params, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + zeros(nx), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p3 = deepcopy(p1) + p4 = deepcopy(p2) + + plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, + label = "ensemble inputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p4, + 1:length(y), + get_g_final(ekpobj)[:, 1], + color = :lightgreen, + label = "ensemble outputs", + linewidth = 4, + linealpha = 0.1, + ) + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + plt = plot(p3, p4, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + + # plot - UQ results, one figure per k + for k in k_values + posterior = posteriors_by_k[k] + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' + + gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + [truth_params final_params] .- truth_params, + label = ["solution" "EKI-opt"], + color = [:black :lightgreen], + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plot!( + p1, + range(0, nx - 1, step = 1), + quantiles[:, 2] .- truth_params, + color = :blue, + label = "posterior", + linewidth = 4, + ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], + fillalpha = 0.1, + ) + + figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + end +end diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl index 93c245e95..a4f743d34 100644 --- a/examples/EKIRace/plot_l96_forcing.jl +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -1,5 +1,5 @@ # load packages -# CES +# CES using Random using JLD2 @@ -14,20 +14,16 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions const EKP = EnsembleKalmanProcesses - -method_cases= ["Inversion", "TransformInversion", "Unscented", "GaussNewtonInversion"] -forcing_cases = ["const-force", "vec-force", "flux-force"] # problem types +include("experiment_config.jl") #### CHOOSE YOUR CASE: -# calib_data_dir -method=method_cases[1] -calibrate_date=Date("2026-05-18", "yyyy-mm-dd") -calib_directory="$(method)_$(calibrate_date)" - -# calib_filename_suffix items to loop over -force_cases=[forcing_cases[3]] -N_enss=[50,75,100] -rng_idxs=[1,2,3,4] +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +force_cases = [cfg.force_case] +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) homedir = pwd() @@ -36,11 +32,10 @@ valid_files = [] for force_case in force_cases for N_ens in N_enss for rng_idx in rng_idxs - data_save_directory = joinpath(homedir, "output", calib_directory) - calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") + data_save_directory = joinpath(homedir, "output", calib_dir) + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) if isfile(data_file) - push!(valid_files, calib_filename_suffix) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) push!(valid_file_items, (force_case, N_ens, rng_idx)) end end @@ -50,278 +45,190 @@ end display(valid_files) println(" ") -for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items,valid_files) - calib_filename_suffix = "$(force_case)_$(N_ens)_$(rng_idx)" - - @info("Plotting for L96: \n method: $(calib_directory) \n experiment: $(calib_filename_suffix)") - - # load - data_save_directory = joinpath(homedir, "output", calib_directory) - figure_save_directory = data_save_directory - data_file = joinpath(data_save_directory, "posterior_$(calib_filename_suffix).jld2") - truth_params = load(data_file)["truth_params_constrained"] - final_params = load(data_file)["final_params_constrained"] - posterior = load(data_file)["posterior"] - - prior_file = joinpath(data_save_directory, "l96_priors_$(force_case).jld2") - prior = load(prior_file)["prior"] - - ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") - ekpobj = load(ekp_file)["ekpobj"] - N_ens = get_N_ens(ekpobj) - nx = length(truth_params) - y = get_obs(ekpobj) - ny = length(y) - # get samples and quantiles from posterior - param_names = get_name(posterior) - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - - quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' # rows are quantiles for row of posterior samples - - - # plot - EKI results - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - truth_params, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - # xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) - # xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), - ) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) - - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - zeros(nx), - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - # xticks = (Int.(0:10:ny), Int.(nx+0:10:nx)) - # xticks = (Int.(0:10:ny), [0:10:nx], 10, 20, 30, (40, 0), 10, 20, 30, 40]), - ) +for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) + + @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + # load + data_save_directory = joinpath(homedir, "output", calib_dir) + figure_save_directory = data_save_directory + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + truth_params = load(data_file)["truth_params_constrained"] + final_params = load(data_file)["final_params_constrained"] + posteriors_by_k = load(data_file)["posteriors_by_k"] + k_values = load(data_file)["k_values"] + prior = load(data_file)["priors"] + + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + ekpobj = load(ekp_file)["ekpobj"] + N_ens = get_N_ens(ekpobj) + nx = length(truth_params) + y = get_obs(ekpobj) + ny = length(y) + + # plot - EKI results + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + truth_params, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + zeros(nx), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) p3 = deepcopy(p1) p4 = deepcopy(p2) - plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) - plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) - - # - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, - label = "ensemble inputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p4, - 1:length(y), - get_g_final(ekpobj)[:, 1], - color = :lightgreen, - label = "ensemble outputs", - linewidth = 4, - linealpha = 0.1, - ) - plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - - - plt = plot(p3, p4, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - - # plot - UQ results - - gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - [truth_params final_params] .- truth_params, - label = ["solution" "EKI-opt"], - color = [:black :lightgreen], - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - plot!( - p1, - range(0, nx - 1, step = 1), - quantiles[:, 2] .- truth_params, # median of all vals - color = :blue, - label = "posterior", - linewidth = 4, - ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], - fillalpha = 0.1, - ) - - figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)") - savefig(figpath * ".png") - savefig(figpath * ".pdf") - end -########################## -#cases = ["GP", "RF-scalar", "RF-nonsep"] -#= -# load -homedir = pwd() -println(homedir) -data_file_GP = joinpath(data_save_directory, "posterior_$(cases[1]).jld2") -data_file_RF = joinpath(data_save_directory, "posterior_$(case).jld2") -truth_params_GP = load(data_file_GP)["truth_params_constrained"] -final_params_GP = load(data_file_GP)["final_params_constrained"] -posterior_GP = load(data_file_GP)["posterior"] -truth_params_RF = load(data_file_RF)["truth_params_constrained"] -final_params_RF = load(data_file_RF)["final_params_constrained"] -posterior_RF = load(data_file_RF)["posterior"] - -ekp_file = joinpath(data_save_directory, "l96_ekp_$(calib_filename_suffix).jld2") -ekpobj = load(ekp_file)["ekpobj"] -N_ens = get_N_ens(ekpobj) -nx = length(truth_params) -y = get_obs(ekpobj) -ny = length(y) -# get samples and quantiles from posterior -param_names_GP = get_name(posterior_GP) -posterior_samples_GP = vcat([get_distribution(posterior_GP)[name] for name in get_name(posterior_GP)]...) #samples are columns -constrained_posterior_samples_GP = - mapslices(x -> transform_unconstrained_to_constrained(posterior_GP, x), posterior_samples_GP, dims = 1) - -quantiles_GP = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_GP)])' # rows are quantiles for row of posterior samples - -param_names_RF = get_name(posterior_RF) -posterior_samples_RF = vcat([get_distribution(posterior_RF)[name] for name in get_name(posterior_RF)]...) #samples are columns -constrained_posterior_samples_RF = - mapslices(x -> transform_unconstrained_to_constrained(posterior_RF, x), posterior_samples_RF, dims = 1) - -quantiles_RF = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples_RF)])' # rows are quantiles for row of posterior samples - -# plot - UQ results - both - -gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) -p1 = plot( - range(0, nx - 1, step = 1), - [truth_params final_params], - label = ["solution" "EKI-opt"], - color = [:black :lightgreen], - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing (input)", - left_margin = 15mm, - bottom_margin = 15mm, -) - -plot!( - p1, - range(0, nx - 1, step = 1), - quantiles_GP[:, 2], # median of all vals - color = :blue, - label = "GP posterior", - linewidth = 4, - ribbon = [quantiles_GP[:, 2] - quantiles_GP[:, 1] quantiles_GP[:, 3] - quantiles_GP[:, 2]], - linealpha = 0.5, - fillalpha = 0.1, -) - -plot!( - p1, - range(0, nx - 1, step = 1), - quantiles_RF[:, 2], # median of all vals - color = :red, - label = "posterior_RF", - linewidth = 4, - ribbon = [quantiles_RF[:, 2] - quantiles_RF[:, 1] quantiles_RF[:, 3] - quantiles_RF[:, 2]], - linealpha = 0.5, - fillalpha = 0.1, -) - - -figpath = joinpath(figure_save_directory, "posterior_ribbons_$(cases[1])_$(case)") -savefig(figpath * ".png") -savefig(figpath * ".pdf") -=# + plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + + l = @layout [a b] + plt = plot(p1, p2, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, + label = "ensemble inputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p3, + range(0, nx - 1, step = 1), + get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p4, + 1:length(y), + get_g_final(ekpobj)[:, 1], + color = :lightgreen, + label = "ensemble outputs", + linewidth = 4, + linealpha = 0.1, + ) + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + plt = plot(p3, p4, layout = l) + + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + + # plot - UQ results, one figure per k + for k in k_values + posterior = posteriors_by_k[k] + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' + + gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p1 = plot( + range(0, nx - 1, step = 1), + [truth_params final_params] .- truth_params, + label = ["solution" "EKI-opt"], + color = [:black :lightgreen], + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plot!( + p1, + range(0, nx - 1, step = 1), + quantiles[:, 2] .- truth_params, + color = :blue, + label = "posterior", + linewidth = 4, + ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], + fillalpha = 0.1, + ) + + figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)_k$(k)") + savefig(figpath * ".png") + savefig(figpath * ".pdf") + end +end From 78d120c29d4ba67594c3fa08d4b36425bbe8d27d Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 11:25:27 -0700 Subject: [PATCH 43/86] detach true_params from prior_mean in flux case --- examples/EKIRace/calibrate_l96.jl | 5 +++-- examples/EKIRace/hpc-variant/calibrate_l96.jl | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index fa957df62..c2dc8b12a 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -83,15 +83,16 @@ elseif case == "flux-force" nu = 61 true_sinusoid(x) = 8 .+ 6 * sin.((4 * pi * x) / 10) x_train = collect(-5.0:0.01:5.0) + Random.seed!(20260529) y_train = true_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) phi_structure = Chain(Dense(1 => 20, tanh), Dense(20 => 1)) - true_model, _ = train_network(phi_structure, x_train, y_train) + true_model, _ = train_network(deepcopy(phi_structure), x_train, y_train) sample_range = Float32.(collect(-5.0:0.1:4.9)) phi = FluxEMC(true_model, sample_range) prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) - prior_model, prior_mean = train_network(phi_structure, x_train, prior_train) + prior_model, prior_mean = train_network(deepcopy(phi_structure), x_train, prior_train) y_pred = true_model(reshape(sample_range, 1, :)) prior_pred = prior_model(reshape(sample_range, 1, :)) diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index fef4e5a05..537ede51b 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -68,14 +68,15 @@ function build_setup(cfg, output_dir) nu = 61 true_sinusoid(x) = 8 .+ 6 * sin.((4 * pi * x) / 10) x_train = collect(-5.0:0.01:5.0) + Random.seed!(20260529) y_train = true_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) phi_structure = Chain(Dense(1 => 20, tanh), Dense(20 => 1)) - true_model, _ = train_network(phi_structure, x_train, y_train) + true_model, _ = train_network(deepcopy(phi_structure), x_train, y_train) sample_range = Float32.(collect(-5.0:0.1:4.9)) phi = FluxEMC(true_model, sample_range) prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) - prior_model, prior_mean = train_network(phi_structure, x_train, prior_train) + prior_model, prior_mean = train_network(deepcopy(phi_structure), x_train, prior_train) prior_cov = (0.1^2) * I(length(prior_mean)) distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) From 749980117698f7fcaab07ad15ad1e32069932aa5 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 12:36:35 -0700 Subject: [PATCH 44/86] inflate prior to 0.5^2 in flux, and only run inversion cases --- examples/EKIRace/calibrate_l63.jl | 6 +++--- examples/EKIRace/calibrate_l96.jl | 8 ++++---- examples/EKIRace/hpc-variant/calibrate_l63.jl | 6 +++--- examples/EKIRace/hpc-variant/calibrate_l96.jl | 8 ++++---- 4 files changed, 14 insertions(+), 14 deletions(-) diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index 93acf6de6..fe2e7bd1c 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -133,9 +133,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) initial_params = construct_initial_ensemble(rng, prior, N_ens) methods = [ Inversion(), - TransformInversion(), - GaussNewtonInversion(prior), - Unscented(prior), +# TransformInversion(), +# GaussNewtonInversion(prior), +# Unscented(prior), ] @info "Ensemble size: $(N_ens)" diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index c2dc8b12a..8d3d75f0d 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -119,7 +119,7 @@ elseif case == "flux-force" # ylabel!("y") # title!("Offline 1D DNN") # display(p) - prior_cov = (0.1^2) * I(length(prior_mean)) + prior_cov = (0.5^2) * I(length(prior_mean)) distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) name = "l96_nn_prior" @@ -239,9 +239,9 @@ for (rr, rng_seed) in enumerate(rng_seeds) initial_params = construct_initial_ensemble(rng, prior, N_ens) methods = [ Inversion(), - TransformInversion(), - GaussNewtonInversion(prior), - Unscented(prior), +# TransformInversion(), +# GaussNewtonInversion(prior), +# Unscented(prior), ] @info "Ensemble size: $(N_ens)" diff --git a/examples/EKIRace/hpc-variant/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl index f44a0b706..5c8c2fe9f 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l63.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -123,9 +123,9 @@ function calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) initial_params = construct_initial_ensemble(rng, setup.prior, N_ens) methods = [ Inversion(), - TransformInversion(), - GaussNewtonInversion(setup.prior), - Unscented(setup.prior), +# TransformInversion(), +# GaussNewtonInversion(setup.prior), +# Unscented(setup.prior), ] conv_cell = fill(NaN, 4) diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index 537ede51b..720653229 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -77,7 +77,7 @@ function build_setup(cfg, output_dir) prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) prior_model, prior_mean = train_network(deepcopy(phi_structure), x_train, prior_train) - prior_cov = (0.1^2) * I(length(prior_mean)) + prior_cov = (0.5^2) * I(length(prior_mean)) distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) prior = ParameterDistribution(distribution, constraint, "l96_nn_prior") @@ -201,9 +201,9 @@ function calibrate_one(cfg, setup, N_ens, rng_idx, output_dir) initial_params = construct_initial_ensemble(rng, setup.prior, N_ens) methods = [ Inversion(), - TransformInversion(), - GaussNewtonInversion(setup.prior), - Unscented(setup.prior), +# TransformInversion(), +# GaussNewtonInversion(setup.prior), +# Unscented(setup.prior), ] conv_cell = fill(NaN, 4) From 9434d414e544ca760a5661954dd870776964a71d Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 12:40:34 -0700 Subject: [PATCH 45/86] allow plotting calibration and emulate-sample results separately --- .../EKIRace/hpc-variant/plot_l96_forcing.jl | 72 ++++++++++++------- examples/EKIRace/plot_l96_forcing.jl | 72 ++++++++++++------- 2 files changed, 96 insertions(+), 48 deletions(-) diff --git a/examples/EKIRace/hpc-variant/plot_l96_forcing.jl b/examples/EKIRace/hpc-variant/plot_l96_forcing.jl index fb763383f..81fcd58f0 100644 --- a/examples/EKIRace/hpc-variant/plot_l96_forcing.jl +++ b/examples/EKIRace/hpc-variant/plot_l96_forcing.jl @@ -15,10 +15,11 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions const EKP = EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction #### CHOOSE YOUR CASE: exp = l96_experiment() -@assert exp in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +@assert exp in (:l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_vec, or :l96_flux in experiment_config.jl" cfg = experiment_config(exp) method = method_cases[1] calib_dir = calib_directory(method, cfg) @@ -28,14 +29,16 @@ rng_idxs = collect(1:cfg.n_repeats) homedir = pwd() +# Discover runs that have calibrate output (ekp + results files). valid_file_items = [] valid_files = [] for force_case in force_cases for N_ens in N_enss for rng_idx in rng_idxs data_save_directory = joinpath(homedir, "output", calib_dir) - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if isfile(data_file) + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + if isfile(ekp_file) && isfile(res_file) push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) push!(valid_file_items, (force_case, N_ens, rng_idx)) end @@ -50,28 +53,39 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - # load - data_save_directory = joinpath(homedir, "output", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) figure_save_directory = data_save_directory - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - truth_params = load(data_file)["truth_params_constrained"] - final_params = load(data_file)["final_params_constrained"] - posteriors_by_k = load(data_file)["posteriors_by_k"] - k_values = load(data_file)["k_values"] - prior = load(data_file)["priors"] + # --- load calibrate outputs --- ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + pri_file = joinpath(data_save_directory, prior_filename(cfg)) + ekpobj = load(ekp_file)["ekpobj"] + prior = load(pri_file)["prior"] + + (truth_params_obj, _) = load(res_file)["truth_params_structure"] + truth_params_constrained = if force_case == "const-force" + [truth_params_obj.val] + elseif force_case == "vec-force" + truth_params_obj.val + elseif force_case == "flux-force" + tp, _ = Flux.destructure(truth_params_obj.model) + tp + end + + final_params_constrained = get_ϕ_mean_final(prior, ekpobj) + N_ens = get_N_ens(ekpobj) - nx = length(truth_params) - y = get_obs(ekpobj) - ny = length(y) + nx = length(truth_params_constrained) + y = get_obs(ekpobj) + ny = length(y) - # plot - EKI results + # --- data plot --- gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), - truth_params, + truth_params_constrained, label = "solution", color = :black, linewidth = 4, @@ -100,6 +114,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + # --- solution plots (mean ensemble) --- gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), @@ -129,7 +144,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item p3 = deepcopy(p1) p4 = deepcopy(p2) - plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p1, range(0, nx - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) l = @layout [a b] @@ -141,7 +156,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, + get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params_constrained, label = "ensemble inputs", color = :lightgreen, linewidth = 4, @@ -151,7 +166,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, + get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -160,7 +175,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, + get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -170,7 +185,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, + get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -195,7 +210,16 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - # plot - UQ results, one figure per k + # --- posterior ribbon plots (only if emulate_sample has been run) --- + post_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if !isfile(post_file) + @info "No posterior file found for $(calib_filename_suffix); skipping posterior_ribbons plots." + continue + end + + posteriors_by_k = load(post_file)["posteriors_by_k"] + k_values = load(post_file)["k_values"] + for k in k_values posterior = posteriors_by_k[k] @@ -207,7 +231,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), - [truth_params final_params] .- truth_params, + [truth_params_constrained final_params_constrained] .- truth_params_constrained, label = ["solution" "EKI-opt"], color = [:black :lightgreen], linewidth = 4, @@ -220,7 +244,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p1, range(0, nx - 1, step = 1), - quantiles[:, 2] .- truth_params, + quantiles[:, 2] .- truth_params_constrained, color = :blue, label = "posterior", linewidth = 4, diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl index a4f743d34..f5c005fbc 100644 --- a/examples/EKIRace/plot_l96_forcing.jl +++ b/examples/EKIRace/plot_l96_forcing.jl @@ -15,9 +15,10 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions const EKP = EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction #### CHOOSE YOUR CASE: -@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +@assert EXPERIMENT in (:l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_vec, or :l96_flux in experiment_config.jl" cfg = experiment_config(EXPERIMENT) method = method_cases[1] calib_dir = calib_directory(method, cfg) @@ -27,14 +28,16 @@ rng_idxs = collect(1:cfg.n_repeats) homedir = pwd() +# Discover runs that have calibrate output (ekp + results files). valid_file_items = [] valid_files = [] for force_case in force_cases for N_ens in N_enss for rng_idx in rng_idxs data_save_directory = joinpath(homedir, "output", calib_dir) - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if isfile(data_file) + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + if isfile(ekp_file) && isfile(res_file) push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) push!(valid_file_items, (force_case, N_ens, rng_idx)) end @@ -49,28 +52,39 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - # load - data_save_directory = joinpath(homedir, "output", calib_dir) + data_save_directory = joinpath(homedir, "output", calib_dir) figure_save_directory = data_save_directory - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - truth_params = load(data_file)["truth_params_constrained"] - final_params = load(data_file)["final_params_constrained"] - posteriors_by_k = load(data_file)["posteriors_by_k"] - k_values = load(data_file)["k_values"] - prior = load(data_file)["priors"] + # --- load calibrate outputs --- ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + pri_file = joinpath(data_save_directory, prior_filename(cfg)) + ekpobj = load(ekp_file)["ekpobj"] + prior = load(pri_file)["prior"] + + (truth_params_obj, _) = load(res_file)["truth_params_structure"] + truth_params_constrained = if force_case == "const-force" + [truth_params_obj.val] + elseif force_case == "vec-force" + truth_params_obj.val + elseif force_case == "flux-force" + tp, _ = Flux.destructure(truth_params_obj.model) + tp + end + + final_params_constrained = get_ϕ_mean_final(prior, ekpobj) + N_ens = get_N_ens(ekpobj) - nx = length(truth_params) - y = get_obs(ekpobj) - ny = length(y) + nx = length(truth_params_constrained) + y = get_obs(ekpobj) + ny = length(y) - # plot - EKI results + # --- data plot --- gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), - truth_params, + truth_params_constrained, label = "solution", color = :black, linewidth = 4, @@ -99,6 +113,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + # --- solution plots (mean ensemble) --- gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), @@ -128,7 +143,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item p3 = deepcopy(p1) p4 = deepcopy(p2) - plot!(p1, range(0, nx - 1, step = 1), final_params .- truth_params, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p1, range(0, nx - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) l = @layout [a b] @@ -140,7 +155,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params, + get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params_constrained, label = "ensemble inputs", color = :lightgreen, linewidth = 4, @@ -150,7 +165,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params, + get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -159,7 +174,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params, + get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -169,7 +184,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p3, range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params, + get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, @@ -194,7 +209,16 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - # plot - UQ results, one figure per k + # --- posterior ribbon plots (only if emulate_sample has been run) --- + post_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if !isfile(post_file) + @info "No posterior file found for $(calib_filename_suffix); skipping posterior_ribbons plots." + continue + end + + posteriors_by_k = load(post_file)["posteriors_by_k"] + k_values = load(post_file)["k_values"] + for k in k_values posterior = posteriors_by_k[k] @@ -206,7 +230,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p1 = plot( range(0, nx - 1, step = 1), - [truth_params final_params] .- truth_params, + [truth_params_constrained final_params_constrained] .- truth_params_constrained, label = ["solution" "EKI-opt"], color = [:black :lightgreen], linewidth = 4, @@ -219,7 +243,7 @@ for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_item plot!( p1, range(0, nx - 1, step = 1), - quantiles[:, 2] .- truth_params, + quantiles[:, 2] .- truth_params_constrained, color = :blue, label = "posterior", linewidth = 4, From 0be3305aabe9046b6bbbcb5ae4b1b6de876f00a0 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 11:45:08 -0700 Subject: [PATCH 46/86] deepcopy --- examples/EKIRace/hpc-variant/emulate_sample_l63.jl | 2 +- examples/EKIRace/hpc-variant/emulate_sample_l96.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l63.jl b/examples/EKIRace/hpc-variant/emulate_sample_l63.jl index 7e62d869b..0d5938968 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_l63.jl +++ b/examples/EKIRace/hpc-variant/emulate_sample_l63.jl @@ -93,7 +93,7 @@ function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) mlt, input_output_pairs; encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, + encoder_kwargs = deepcopy(encoder_kwargs), ) optimize_hyperparameters!(emulator) diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl index 4c3fcd5f9..6abb89167 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl +++ b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl @@ -104,7 +104,7 @@ function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) mlt, input_output_pairs; encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, + encoder_kwargs = deepcopy(encoder_kwargs), ) optimize_hyperparameters!(emulator) From 8b0572da857014bf38fd7f838c45acede1427510 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sun, 31 May 2026 14:08:49 -0700 Subject: [PATCH 47/86] scale prior and inflate --- examples/EKIRace/calibrate_l96.jl | 4 +++- examples/EKIRace/hpc-variant/calibrate_l96.jl | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index 8d3d75f0d..948dbcd83 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -119,7 +119,9 @@ elseif case == "flux-force" # ylabel!("y") # title!("Offline 1D DNN") # display(p) - prior_cov = (0.5^2) * I(length(prior_mean)) + + #prior_cov = (0.1^2) * I(length(prior_mean)) # original cov + prior_cov = (1.0^2) * Diagonal(prior_mean.^2) # scaled to mean distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) name = "l96_nn_prior" diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index 720653229..5083f6452 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -77,7 +77,8 @@ function build_setup(cfg, output_dir) prior_sinusoid(x) = 8.02 .+ 6.5 * sin.(1.02 * (4 * pi * x) / 10 + 0.2) prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) prior_model, prior_mean = train_network(deepcopy(phi_structure), x_train, prior_train) - prior_cov = (0.5^2) * I(length(prior_mean)) + #prior_cov = (0.1^2) * I(length(prior_mean)) # original cov + prior_cov = (1.0^2) * Diagonal(prior_mean.^2) # scaled to mean distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) prior = ParameterDistribution(distribution, constraint, "l96_nn_prior") From 70e1fb5f810e2b422b609c28dd53ce78a3214b4e Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 1 Jun 2026 00:49:30 -0700 Subject: [PATCH 48/86] first go at pushing forward the posteriors --- examples/EKIRace/ensemble_from_posterior.jl | 214 ++++++++++++++++++ .../hpc-variant/ensemble_from_posterior.jl | 214 ++++++++++++++++++ 2 files changed, 428 insertions(+) create mode 100644 examples/EKIRace/ensemble_from_posterior.jl create mode 100644 examples/EKIRace/hpc-variant/ensemble_from_posterior.jl diff --git a/examples/EKIRace/ensemble_from_posterior.jl b/examples/EKIRace/ensemble_from_posterior.jl new file mode 100644 index 000000000..402a8ace4 --- /dev/null +++ b/examples/EKIRace/ensemble_from_posterior.jl @@ -0,0 +1,214 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates +using Plots +using Plots.Measures + + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz96.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true + +n_samples_pushforward = 100 # num. samples to pushforward through lorenz + + +######################################################################## +################## file-certainty for loading ########################## +######################################################################## + +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +force_case = cfg.force_case +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +# get some preliminaries +prelim_dir = joinpath(@__DIR__, "output") +if !isdir(prelim_dir) + mkdir(prelim_dir) +end +prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") +if isfile(prelim_file) + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + R_inv_var = loaded_data["R_inv_var"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + @info "loaded precomputed preliminary quantities from $(prelim_file)" +else + throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) +end +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Pushing forward posteriors through the forward map from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") +end + +### Then load data + +# Load first valid file to determine parameter dimension and max k +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +for (N_ens, rng_idx) in valid_file_items + calib_fn = results_filename(cfg, N_ens, rng_idx) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] + + (truth_params_obj, _) = loaded_c["truth_params_structure"] + + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + truth_params_constrained, _ = destructure(truth_params_obj.model) + + ( + truth_params_constrained, + truth_params_obj.model, + truth_params_obj.sample_range + ) + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + + for k in k_values + post_dist = posteriors_by_k[k] + push_ensemble = sample(post_dist, n_samples_pushforward) + # note: push_ensemble is unconstrained + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + G_ens = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + n_p = length(truth_params_constrained) + ny = length(y) + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p3 = plot( + range(0, n_p - 1, step = 1), + zeros(n_p), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plot!( + p3, + range(0, n_p - 1, step = 1), + constrained_push_ensemble[:, 1] .- truth_params_constrained, + label = "posterior samples", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!( + p3, + range(0, n_p - 1, step = 1), + constrained_push_ensemble[:, 2:end] .- truth_params_constrained, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p4, + 1:ny, + G_ens[:, 1], + label = "pushforward outputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + l = @layout [a b] + plt = plot(p3, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + end +end + + diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl new file mode 100644 index 000000000..402a8ace4 --- /dev/null +++ b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl @@ -0,0 +1,214 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates +using Plots +using Plots.Measures + + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz96.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true + +n_samples_pushforward = 100 # num. samples to pushforward through lorenz + + +######################################################################## +################## file-certainty for loading ########################## +######################################################################## + +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +force_case = cfg.force_case +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +# get some preliminaries +prelim_dir = joinpath(@__DIR__, "output") +if !isdir(prelim_dir) + mkdir(prelim_dir) +end +prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") +if isfile(prelim_file) + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + R_inv_var = loaded_data["R_inv_var"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + @info "loaded precomputed preliminary quantities from $(prelim_file)" +else + throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) +end +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Pushing forward posteriors through the forward map from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") +end + +### Then load data + +# Load first valid file to determine parameter dimension and max k +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +for (N_ens, rng_idx) in valid_file_items + calib_fn = results_filename(cfg, N_ens, rng_idx) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] + + (truth_params_obj, _) = loaded_c["truth_params_structure"] + + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + truth_params_constrained, _ = destructure(truth_params_obj.model) + + ( + truth_params_constrained, + truth_params_obj.model, + truth_params_obj.sample_range + ) + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + + for k in k_values + post_dist = posteriors_by_k[k] + push_ensemble = sample(post_dist, n_samples_pushforward) + # note: push_ensemble is unconstrained + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + G_ens = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + n_p = length(truth_params_constrained) + ny = length(y) + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p3 = plot( + range(0, n_p - 1, step = 1), + zeros(n_p), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "Forcing difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plot!( + p3, + range(0, n_p - 1, step = 1), + constrained_push_ensemble[:, 1] .- truth_params_constrained, + label = "posterior samples", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!( + p3, + range(0, n_p - 1, step = 1), + constrained_push_ensemble[:, 2:end] .- truth_params_constrained, + label = "", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + + plot!( + p4, + 1:ny, + G_ens[:, 1], + label = "pushforward outputs", + color = :lightgreen, + linewidth = 4, + linealpha = 0.1, + ) + plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + l = @layout [a b] + plt = plot(p3, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + end +end + + From f7aa5d0f9ce400766b56d8f465a6d120cf47ea99 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 1 Jun 2026 01:31:49 -0700 Subject: [PATCH 49/86] add new job scripts for the pushforward of the posterior --- examples/EKIRace/ensemble_from_posterior.jl | 34 +-- .../hpc-variant/ensemble_from_posterior.jl | 203 +++++++----------- .../ensemble_from_posterior.sbatch | 31 +++ .../EKIRace/hpc-variant/submit_l96_flux.sh | 10 + .../EKIRace/hpc-variant/submit_l96_vec.sh | 10 + 5 files changed, 141 insertions(+), 147 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch diff --git a/examples/EKIRace/ensemble_from_posterior.jl b/examples/EKIRace/ensemble_from_posterior.jl index 402a8ace4..63a2bebe4 100644 --- a/examples/EKIRace/ensemble_from_posterior.jl +++ b/examples/EKIRace/ensemble_from_posterior.jl @@ -142,18 +142,22 @@ for (N_ens, rng_idx) in valid_file_items ]..., ) - n_p = length(truth_params_constrained) ny = length(y) + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p3 = plot( - range(0, n_p - 1, step = 1), - zeros(n_p), + xaxis_forcing, + truth_forcing, label = "solution", color = :black, linewidth = 4, - xlabel = "Parameter index", - ylabel = "Forcing difference (input)", + xlabel = "Spatial index", + ylabel = "Forcing (input)", left_margin = 15mm, bottom_margin = 15mm, ) @@ -170,24 +174,8 @@ for (N_ens, rng_idx) in valid_file_items bottom_margin = 15mm, ) - plot!( - p3, - range(0, n_p - 1, step = 1), - constrained_push_ensemble[:, 1] .- truth_params_constrained, - label = "posterior samples", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!( - p3, - range(0, n_p - 1, step = 1), - constrained_push_ensemble[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) plot!( p4, diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl index 402a8ace4..daf14a722 100644 --- a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl +++ b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl @@ -1,5 +1,5 @@ # Import modules -using Distributions # probability distributions and associated functions +using Distributions using LinearAlgebra using Random using JLD2 @@ -7,6 +7,7 @@ using Statistics using Flux using BSON using Dates +ENV["GKSwstype"] = "100" using Plots using Plots.Measures @@ -16,97 +17,51 @@ using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses -include("Lorenz96.jl") # Contains Lorenz 96 source code +include("Lorenz96.jl") include("experiment_config.jl") -verbose_flag = false -save_all_ekp = true - -n_samples_pushforward = 100 # num. samples to pushforward through lorenz - - ######################################################################## -################## file-certainty for loading ########################## +############### Per-cell pushforward ################################### ######################################################################## -@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" -cfg = experiment_config(EXPERIMENT) -method = method_cases[1] # method_cases defined in experiment_config.jl -calib_dir = calib_directory(method, cfg) -force_case = cfg.force_case -N_enss = cfg.N_ens_sizes -rng_idxs = collect(1:cfg.n_repeats) - -# get some preliminaries -prelim_dir = joinpath(@__DIR__, "output") -if !isdir(prelim_dir) - mkdir(prelim_dir) -end -prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") -if isfile(prelim_file) - loaded_data = JLD2.load(prelim_file) - x0 = loaded_data["x0"] - nx = length(x0) - y = loaded_data["y"] - ic_cov_sqrt = loaded_data["ic_cov_sqrt"] - R = loaded_data["R"] - R_inv_var = loaded_data["R_inv_var"] - lorenz_config_settings = loaded_data["lorenz_config_settings"] - observation_config = loaded_data["observation_config"] - - @info "loaded precomputed preliminary quantities from $(prelim_file)" -else - throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) -end -### determine valid files before we try to load - -homedir = joinpath(pwd()) -data_save_directory = joinpath(homedir, "output", calib_dir) - -valid_file_items = [] -valid_files = [] -for N_ens in N_enss - for rng_idx in rng_idxs - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if isfile(data_file) - push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) - push!(valid_file_items, (N_ens, rng_idx)) - end - end -end - -@info "Pushing forward posteriors through the forward map from valid files:" -display(valid_files) +function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) + n_samples_pushforward = 100 -if isempty(valid_file_items) - error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") -end + force_case = cfg.force_case + calib_dir = calib_directory(method, cfg) + calib_fn = results_filename(cfg, N_ens, rng_idx) + post_fn = posterior_filename(cfg, N_ens, rng_idx) -### Then load data - -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + # load preliminaries + prelim_dir = joinpath(@__DIR__, "output") + prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") + if !isfile(prelim_file) + throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] -n_k = maximum( - maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) - for (N_ens, rng_idx) in valid_file_items -) + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) -n_rng = length(rng_idxs) -n_ens = length(N_enss) + if !isfile(joinpath(data_save_directory, post_fn)) + @warn "No posterior file found for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end -for (N_ens, rng_idx) in valid_file_items - calib_fn = results_filename(cfg, N_ens, rng_idx) - post_fn = posterior_filename(cfg, N_ens, rng_idx) @info "loading case $(post_fn)" loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) - + posteriors_by_k = loaded_p["posteriors_by_k"] priors = loaded_p["priors"] k_values = loaded_p["k_values"] - truth_params = loaded_p["truth_params"] (truth_params_obj, _) = loaded_c["truth_params_structure"] @@ -115,16 +70,10 @@ for (N_ens, rng_idx) in valid_file_items elseif force_case == "vec-force" (truth_params_obj.val, nothing, nothing) elseif force_case == "flux-force" - truth_params_constrained, _ = destructure(truth_params_obj.model) - - ( - truth_params_constrained, - truth_params_obj.model, - truth_params_obj.sample_range - ) + tp, _ = destructure(truth_params_obj.model) + (tp, truth_params_obj.model, truth_params_obj.sample_range) end - truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) - + for k in k_values post_dist = posteriors_by_k[k] push_ensemble = sample(post_dist, n_samples_pushforward) @@ -134,26 +83,30 @@ for (N_ens, rng_idx) in valid_file_items G_ens = hcat( [ lorenz_forward( - build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), + build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), lorenz_config_settings, observation_config, ) for j in 1:n_samples_pushforward - ]..., + ]..., ) - n_p = length(truth_params_constrained) - ny = length(y) + ny = length(y) + + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) p3 = plot( - range(0, n_p - 1, step = 1), - zeros(n_p), + xaxis_forcing, + truth_forcing, label = "solution", color = :black, linewidth = 4, - xlabel = "Parameter index", - ylabel = "Forcing difference (input)", + xlabel = "Spatial index", + ylabel = "Forcing (input)", left_margin = 15mm, bottom_margin = 15mm, ) @@ -170,45 +123,47 @@ for (N_ens, rng_idx) in valid_file_items bottom_margin = 15mm, ) - plot!( - p3, - range(0, n_p - 1, step = 1), - constrained_push_ensemble[:, 1] .- truth_params_constrained, - label = "posterior samples", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!( - p3, - range(0, n_p - 1, step = 1), - constrained_push_ensemble[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!( - p4, - 1:ny, - G_ens[:, 1], - label = "pushforward outputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) l = @layout [a b] plt = plot(p3, p4, layout = l) - suffix = case_suffix(cfg, N_ens, rng_idx) + suffix = case_suffix(cfg, N_ens, rng_idx) figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" savefig(plt, joinpath(data_save_directory, figure_fn)) savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) - end -end +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## +function main() + exp = l96_experiment() + @assert( + exp in (:l96_const, :l96_vec, :l96_flux), + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)", + ) + cfg = experiment_config(exp) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end + else + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end +end +main() diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch b/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch new file mode 100644 index 000000000..69d451033 --- /dev/null +++ b/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +# Single job that pushes all posteriors forward through the Lorenz96 forward map +# and saves pushforward plots. Runs all (N_ens, rng_idx) cases serially in one job +# (no array needed — the script loops internally). +# +# L96: +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_vec ensemble_from_posterior.sbatch +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_flux ensemble_from_posterior.sbatch + +#SBATCH --job-name=postfwd +#SBATCH --output=output/slurm/postfwd_%j.out +#SBATCH --error=output/slurm/postfwd_%j.err +#SBATCH --time=08:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +julia --project=. ensemble_from_posterior.jl diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 867d9ae72..1fe7524df 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -36,4 +36,14 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" +echo "=== Submitting ensemble_from_posterior (L96 flux-force, after ${EMU_JID}) ===" +POST_JID=$(sbatch --parsable \ + -A esm \ + --job-name="post_${LABEL}" \ + --dependency=afterok:${EMU_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_flux \ + ensemble_from_posterior.sbatch) +echo " ensemble_from_posterior job ID: ${POST_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 289154ac0..e07c01fd3 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -36,4 +36,14 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" +echo "=== Submitting ensemble_from_posterior (L96 vec-force, after ${EMU_JID}) ===" +POST_JID=$(sbatch --parsable \ + -A esm \ + --job-name="post_${LABEL}" \ + --dependency=afterok:${EMU_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_vec \ + ensemble_from_posterior.sbatch) +echo " ensemble_from_posterior job ID: ${POST_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" From 32a0cfee0942a132de591576edd608277ba5a535 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 1 Jun 2026 13:42:00 -0700 Subject: [PATCH 50/86] update readmes --- examples/EKIRace/README.md | 115 +++++++++++++++++++++++++ examples/EKIRace/hpc-variant/README.md | 36 ++++++++ 2 files changed, 151 insertions(+) create mode 100644 examples/EKIRace/README.md diff --git a/examples/EKIRace/README.md b/examples/EKIRace/README.md new file mode 100644 index 000000000..3de352e49 --- /dev/null +++ b/examples/EKIRace/README.md @@ -0,0 +1,115 @@ +# EKIRace + +Serial (local) versions of the calibrate, emulate-sample, and posterior-pushforward +scripts for the Lorenz 63 and Lorenz 96 benchmark experiments. + +## Experiments + +| Symbol | Model | Forcing type | +|--------|-------|-------------| +| `:l63` | Lorenz 63 | — | +| `:l96_const` | Lorenz 96 | scalar constant | +| `:l96_vec` | Lorenz 96 | spatial vector | +| `:l96_flux` | Lorenz 96 | neural-network flux | + +## One-time setup + +In `experiment_config.jl`, set `EXPERIMENT` to the case you want to run and pin +`calibrate_date` to today's date before starting a run, e.g. + +```julia +EXPERIMENT = :l96_vec +calibrate_date = Date("2026-05-31", "yyyy-mm-dd") +``` + +All subsequent stages (emulate_sample, ensemble_from_posterior, leaderboard +utilities) include `experiment_config.jl` and read these values, so keeping them +fixed ensures every stage finds the right output directory. + +## Running scripts + +Run all scripts from the `examples/EKIRace/` directory using: + +```bash +julia --project -e 'include("scriptname.jl")' +``` + +Each script loops over all `(N_ens, rng_idx)` cells defined in `experiment_config.jl` +and writes output to `output/_/`. + +## Stages + +### 1. Calibrate + +```bash +# L63 (EXPERIMENT setting is ignored) +julia --project -e 'include("calibrate_l63.jl")' + +# L96 — reads EXPERIMENT from experiment_config.jl +julia --project -e 'include("calibrate_l96.jl")' +``` + +The L96 calibration also writes a `l96_computed_preliminaries_.jld2` +file to `output/` that the later stages require. + +### 2. Emulate and sample + +```bash +# L63 +julia --project -e 'include("emulate_sample_l63.jl")' + +# L96 +julia --project -e 'include("emulate_sample_l96.jl")' +``` + +Reads the calibration output for each `(N_ens, rng_idx)` cell, trains an +emulator, runs MCMC, and writes per-cell posterior `.jld2` files plus +summary netCDF files to the calibration output directory. + +### 3. Posterior pushforward (L96 only) + +```bash +julia --project -e 'include("ensemble_from_posterior.jl")' +``` + +Requires `emulate_sample_l96.jl` to have completed. Loops over all valid +`(N_ens, rng_idx)` cells, samples 100 points from each posterior, runs them +through the Lorenz 96 forward map, and writes +`pushforward_from_posterior_*_k_full_ens.{png,pdf}` plots into the +calibration output directory. + +### 4. Leaderboard + +The leaderboard utility scripts (`l63_exp_to_leaderboard_utilities.jl`, +`l96_exp_to_leaderboard_utilities.jl`) are not standalone — they are included +by downstream analysis. + +To compute summary metrics from a netCDF result file, edit the `filename` +variable at the top of `compute_leaderboard_metrics.jl` to point to the +target file, then run: + +```bash +julia --project -e 'include("compute_leaderboard_metrics.jl")' +``` + +This prints per-ensemble-size Mahalanobis and log-posterior scores against +chi-squared reference quantiles. + +## Full run example (L96 vector forcing) + +```julia +# In experiment_config.jl: +EXPERIMENT = :l96_vec +calibrate_date = Date("2026-05-31", "yyyy-mm-dd") +``` + +```bash +julia --project -e 'include("calibrate_l96.jl")' +julia --project -e 'include("emulate_sample_l96.jl")' +julia --project -e 'include("ensemble_from_posterior.jl")' +``` + +## HPC variant + +For parallelised SLURM job-array runs on the Caltech Resnick cluster, see +[`hpc-variant/README.md`](hpc-variant/README.md). diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index 05ba76478..b19821cfe 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -45,6 +45,19 @@ julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only julia --project=. emulate_sample_l63.jl 5 # fifth cell only ``` +After `emulate_sample` completes, push the posteriors forward through the L96 +forward map and save plots: + +```bash +EXPERIMENT=l96_const julia --project=. ensemble_from_posterior.jl +EXPERIMENT=l96_vec julia --project=. ensemble_from_posterior.jl +EXPERIMENT=l96_flux julia --project=. ensemble_from_posterior.jl +``` + +Each call loops over all `(N_ens, rng_idx)` cells, samples 100 points from each +posterior, runs them through the forward model, and writes `pushforward_from_posterior_*.{png,pdf}` +plots into the per-case calibration output directory. + ## HPC (Caltech Resnick cluster, SLURM) ### Submission scripts (recommended) @@ -100,6 +113,29 @@ sbatch -A esm \ Each task writes its per-method `ekp` and `results` files, which are consumed directly by the emulate_sample stage. +### Posterior pushforward (`ensemble_from_posterior`) + +After the `emulate_sample` jobs finish, submit a single (non-array) job that +loops over all `(N_ens, rng_idx)` cells internally: + +```bash +eid=$(sbatch --parsable \ + --dependency=afterok: --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_const \ + ensemble_from_posterior.sbatch) + +sbatch --dependency=afterok: --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_vec ensemble_from_posterior.sbatch + +sbatch --dependency=afterok: --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_flux ensemble_from_posterior.sbatch +``` + +Replace `` with the job ID returned by the corresponding +`emulate_sample_array.sbatch` submission. Each run takes up to 8 h and writes +`pushforward_from_posterior_*.{png,pdf}` into the same output directory as the +calibration results. + ### Adjusting array size The sbatch files default to `--array=1-60` (3 ensemble sizes × 20 repeats). From 33318a5b98bcedcafb427be3ecbef19c26679147 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 2 Jun 2026 00:04:16 -0700 Subject: [PATCH 51/86] add panel to pushforward plots --- examples/EKIRace/ensemble_from_posterior.jl | 31 ++++++++++++++----- .../hpc-variant/ensemble_from_posterior.jl | 24 +++++++++++--- 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/examples/EKIRace/ensemble_from_posterior.jl b/examples/EKIRace/ensemble_from_posterior.jl index 63a2bebe4..5361787cf 100644 --- a/examples/EKIRace/ensemble_from_posterior.jl +++ b/examples/EKIRace/ensemble_from_posterior.jl @@ -143,13 +143,27 @@ for (N_ens, rng_idx) in valid_file_items ) ny = length(y) - + n_par = length(truth_params_constrained) truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) truth_forcing = forcing(truth_emc, x0) xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range - push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + push_forcings = reduce(hcat,[forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]) + + param_diffs = reduce(hcat,[constrained_push_ensemble[:, j]- truth_params_constrained for j in 1:n_samples_pushforward]) - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p2 = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters(input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) p3 = plot( xaxis_forcing, truth_forcing, @@ -157,7 +171,7 @@ for (N_ens, rng_idx) in valid_file_items color = :black, linewidth = 4, xlabel = "Spatial index", - ylabel = "Forcing (input)", + ylabel = "Forcing (transformed-input)", left_margin = 15mm, bottom_margin = 15mm, ) @@ -169,11 +183,14 @@ for (N_ens, rng_idx) in valid_file_items color = :black, linewidth = 4, xlabel = "Spatial index", - ylabel = "State mean/std output", + ylabel = "State mean/std (output)", left_margin = 15mm, bottom_margin = 15mm, ) + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) @@ -188,8 +205,8 @@ for (N_ens, rng_idx) in valid_file_items ) plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - l = @layout [a b] - plt = plot(p3, p4, layout = l) + l = @layout [a b c] + plt = plot(p2, p3, p4, layout = l) suffix = case_suffix(cfg, N_ens, rng_idx) figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl index daf14a722..44972e5cd 100644 --- a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl +++ b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl @@ -91,14 +91,27 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ ]..., ) - ny = length(y) + ny = length(y) + n_par = length(truth_params_constrained) truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) truth_forcing = forcing(truth_emc, x0) xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + p2 = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters(input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) p3 = plot( xaxis_forcing, truth_forcing, @@ -123,14 +136,17 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ bottom_margin = 15mm, ) + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - l = @layout [a b] - plt = plot(p3, p4, layout = l) + l = @layout [a b c] + plt = plot(p2, p3, p4, layout = l) suffix = case_suffix(cfg, N_ens, rng_idx) figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" From 5cedcf6b5d73de1afce6130cece1e6449896011a Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 2 Jun 2026 16:20:09 -0700 Subject: [PATCH 52/86] improved plotting consistecy and script organization --- .../calibration_diagnostic_plots_l96.jl | 290 ++++++++++++++ examples/EKIRace/ensemble_from_posterior.jl | 219 ----------- .../calibration_diagnostic_plots_l96.jl | 291 ++++++++++++++ .../calibration_diagnostic_plots_l96.sbatch | 30 ++ .../hpc-variant/ensemble_from_posterior.jl | 185 --------- .../ensemble_from_posterior.sbatch | 31 -- .../EKIRace/hpc-variant/plot_l96_forcing.jl | 259 ------------ .../posterior_diagnostic_plots_l96.jl | 343 ++++++++++++++++ .../posterior_diagnostic_plots_l96.sbatch | 30 ++ .../EKIRace/hpc-variant/submit_l96_const.sh | 22 +- .../EKIRace/hpc-variant/submit_l96_flux.sh | 22 +- .../EKIRace/hpc-variant/submit_l96_vec.sh | 22 +- examples/EKIRace/plot_l96_forcing.jl | 258 ------------ .../EKIRace/posterior_diagnostic_plots_l96.jl | 368 ++++++++++++++++++ 14 files changed, 1405 insertions(+), 965 deletions(-) create mode 100644 examples/EKIRace/calibration_diagnostic_plots_l96.jl delete mode 100644 examples/EKIRace/ensemble_from_posterior.jl create mode 100644 examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.jl create mode 100644 examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch delete mode 100644 examples/EKIRace/hpc-variant/ensemble_from_posterior.jl delete mode 100644 examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch delete mode 100644 examples/EKIRace/hpc-variant/plot_l96_forcing.jl create mode 100644 examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl create mode 100644 examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch delete mode 100644 examples/EKIRace/plot_l96_forcing.jl create mode 100644 examples/EKIRace/posterior_diagnostic_plots_l96.jl diff --git a/examples/EKIRace/calibration_diagnostic_plots_l96.jl b/examples/EKIRace/calibration_diagnostic_plots_l96.jl new file mode 100644 index 000000000..baa66cf1f --- /dev/null +++ b/examples/EKIRace/calibration_diagnostic_plots_l96.jl @@ -0,0 +1,290 @@ +# load packages +# CES + +using Random +using JLD2 +using Statistics +using LinearAlgebra +using Dates + +using Plots +using Plots.Measures + +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions +const EKP = EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction + +#### CHOOSE YOUR CASE: +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +force_cases = [cfg.force_case] +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +homedir = pwd() + +# Discover runs that have calibrate output (ekp + results files). +valid_file_items = [] +valid_files = [] +for force_case in force_cases + for N_ens in N_enss + for rng_idx in rng_idxs + data_save_directory = joinpath(homedir, "output", calib_dir) + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + if isfile(ekp_file) && isfile(res_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (force_case, N_ens, rng_idx)) + end + end + end +end +@info "Plotting valid files:" +display(valid_files) +println(" ") + +# Load x0 from preliminaries (needed for computing panel-ii forcing). +prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(cfg.force_case).jld2") +if !isfile(prelim_file) + error("Preliminaries file not found: $(prelim_file). Run calibrate_l96.jl first.") +end +x0 = JLD2.load(prelim_file)["x0"] + +for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) + + @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + data_save_directory = joinpath(homedir, "output", calib_dir) + figure_save_directory = data_save_directory + + # --- load calibrate outputs --- + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + pri_file = joinpath(data_save_directory, prior_filename(cfg)) + + ekpobj = load(ekp_file)["ekpobj"] + prior = load(pri_file)["prior"] + + (truth_params_obj, _) = load(res_file)["truth_params_structure"] + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + tp, _ = Flux.destructure(truth_params_obj.model) + (tp, truth_params_obj.model, truth_params_obj.sample_range) + end + + final_params_constrained = get_ϕ_mean_final(prior, ekpobj) + + N_ens = get_N_ens(ekpobj) + n_par = length(truth_params_constrained) + y = get_obs(ekpobj) + ny = length(y) + + # --- compute forcing quantities --- + is_const = force_case == "const-force" + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + final_emc = build_forcing(truth_params_obj, final_params_constrained, structure, sample_range) + final_forcing = forcing(final_emc, x0) + + # --- data plot --- + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + p1 = if is_const + hline( + [truth_params_constrained[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + range(0, n_par - 1, step = 1), + truth_params_constrained, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p_mid = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plt = plot(p1, p_mid, p2, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + # --- solution plots (mean ensemble) --- + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + p1 = if is_const + hline( + [0.0], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + range(0, n_par - 1, step = 1), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p_mid = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p3 = deepcopy(p1) + p_mid3 = deepcopy(p_mid) + p4 = deepcopy(p2) + + if is_const + hline!(p1, [final_params_constrained[1] - truth_params_constrained[1]], label = "mean ensemble input", color = :lightgreen, linewidth = 4) + hline!(p_mid, [final_forcing[1]], label = "mean ensemble forcing", color = :lightgreen, linewidth = 4) + else + plot!(p1, range(0, n_par - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p_mid, xaxis_forcing, final_forcing, label = "mean ensemble forcing", color = :lightgreen, linewidth = 4) + end + plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + + plt = plot(p1, p_mid, p2, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + + # --- full ensemble plot --- + ens_final_params = get_ϕ_final(prior, ekpobj) + ens_init_params = get_ϕ(prior, ekpobj, 1) + ens_final_forcings = hcat([forcing(build_forcing(truth_params_obj, ens_final_params[:, j], structure, sample_range), x0) for j in axes(ens_final_params, 2)]...) + ens_init_forcings = hcat([forcing(build_forcing(truth_params_obj, ens_init_params[:, j], structure, sample_range), x0) for j in axes(ens_init_params, 2)]...) + + if is_const + hline!(p3, [ens_final_params[1, 1] - truth_params_constrained[1]], label = "ensemble inputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, ens_final_params[1, 2:end] .- truth_params_constrained[1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, ens_init_params[1, :] .- truth_params_constrained[1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + hline!(p_mid3, [ens_final_forcings[1, 1]], label = "ensemble forcings", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p_mid3, ens_final_forcings[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p_mid3, ens_init_forcings[1, :], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + else + plot!(p3, range(0, n_par - 1, step = 1), ens_final_params[:, 1] .- truth_params_constrained, label = "ensemble inputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_final_params[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_init_params[:, 1] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_init_params[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + plot!(p_mid3, xaxis_forcing, ens_final_forcings[:, 1], label = "ensemble forcings", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_final_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_init_forcings[:, 1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_init_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + end + + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 1], color = :lightgreen, label = "ensemble outputs", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + plt = plot(p3, p_mid3, p4, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + +end diff --git a/examples/EKIRace/ensemble_from_posterior.jl b/examples/EKIRace/ensemble_from_posterior.jl deleted file mode 100644 index 5361787cf..000000000 --- a/examples/EKIRace/ensemble_from_posterior.jl +++ /dev/null @@ -1,219 +0,0 @@ -# Import modules -using Distributions # probability distributions and associated functions -using LinearAlgebra -using Random -using JLD2 -using Statistics -using Flux -using BSON -using Dates -using Plots -using Plots.Measures - - -# CES -using CalibrateEmulateSample.ParameterDistributions -using CalibrateEmulateSample.DataContainers -using CalibrateEmulateSample.EnsembleKalmanProcesses - -include("Lorenz96.jl") # Contains Lorenz 96 source code -include("experiment_config.jl") - -verbose_flag = false -save_all_ekp = true - -n_samples_pushforward = 100 # num. samples to pushforward through lorenz - - -######################################################################## -################## file-certainty for loading ########################## -######################################################################## - -@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" -cfg = experiment_config(EXPERIMENT) -method = method_cases[1] # method_cases defined in experiment_config.jl -calib_dir = calib_directory(method, cfg) -force_case = cfg.force_case -N_enss = cfg.N_ens_sizes -rng_idxs = collect(1:cfg.n_repeats) - -# get some preliminaries -prelim_dir = joinpath(@__DIR__, "output") -if !isdir(prelim_dir) - mkdir(prelim_dir) -end -prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") -if isfile(prelim_file) - loaded_data = JLD2.load(prelim_file) - x0 = loaded_data["x0"] - nx = length(x0) - y = loaded_data["y"] - ic_cov_sqrt = loaded_data["ic_cov_sqrt"] - R = loaded_data["R"] - R_inv_var = loaded_data["R_inv_var"] - lorenz_config_settings = loaded_data["lorenz_config_settings"] - observation_config = loaded_data["observation_config"] - - @info "loaded precomputed preliminary quantities from $(prelim_file)" -else - throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) -end -### determine valid files before we try to load - -homedir = joinpath(pwd()) -data_save_directory = joinpath(homedir, "output", calib_dir) - -valid_file_items = [] -valid_files = [] -for N_ens in N_enss - for rng_idx in rng_idxs - data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if isfile(data_file) - push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) - push!(valid_file_items, (N_ens, rng_idx)) - end - end -end - -@info "Pushing forward posteriors through the forward map from valid files:" -display(valid_files) - -if isempty(valid_file_items) - error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") -end - -### Then load data - -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) - -n_k = maximum( - maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) - for (N_ens, rng_idx) in valid_file_items -) - -n_rng = length(rng_idxs) -n_ens = length(N_enss) - -for (N_ens, rng_idx) in valid_file_items - calib_fn = results_filename(cfg, N_ens, rng_idx) - post_fn = posterior_filename(cfg, N_ens, rng_idx) - @info "loading case $(post_fn)" - loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) - loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) - - posteriors_by_k = loaded_p["posteriors_by_k"] - priors = loaded_p["priors"] - k_values = loaded_p["k_values"] - truth_params = loaded_p["truth_params"] - - (truth_params_obj, _) = loaded_c["truth_params_structure"] - - (truth_params_constrained, structure, sample_range) = if force_case == "const-force" - ([truth_params_obj.val], nothing, nothing) - elseif force_case == "vec-force" - (truth_params_obj.val, nothing, nothing) - elseif force_case == "flux-force" - truth_params_constrained, _ = destructure(truth_params_obj.model) - - ( - truth_params_constrained, - truth_params_obj.model, - truth_params_obj.sample_range - ) - end - truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) - - for k in k_values - post_dist = posteriors_by_k[k] - push_ensemble = sample(post_dist, n_samples_pushforward) - # note: push_ensemble is unconstrained - constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) - - G_ens = hcat( - [ - lorenz_forward( - build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) for j in 1:n_samples_pushforward - ]..., - ) - - ny = length(y) - n_par = length(truth_params_constrained) - truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) - truth_forcing = forcing(truth_emc, x0) - xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range - push_forcings = reduce(hcat,[forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]) - - param_diffs = reduce(hcat,[constrained_push_ensemble[:, j]- truth_params_constrained for j in 1:n_samples_pushforward]) - - - gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p2 = plot( - collect(1:n_par), - zeros(n_par), - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "parameters(input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - p3 = plot( - xaxis_forcing, - truth_forcing, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing (transformed-input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - p4 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(R)), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std (output)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - - plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - - plot!( - p4, - 1:ny, - G_ens[:, 1], - label = "pushforward outputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - - l = @layout [a b c] - plt = plot(p2, p3, p4, layout = l) - - suffix = case_suffix(cfg, N_ens, rng_idx) - figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" - savefig(plt, joinpath(data_save_directory, figure_fn)) - savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) - - end -end - - diff --git a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.jl new file mode 100644 index 000000000..e27392123 --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.jl @@ -0,0 +1,291 @@ +# load packages +# CES + +using Random +using JLD2 +using Statistics +using LinearAlgebra +using Dates + +using Plots +using Plots.Measures + +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions +const EKP = EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction + +#### CHOOSE YOUR CASE: +exp = l96_experiment() +@assert exp in (:l96_const, :l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(exp) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +force_cases = [cfg.force_case] +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +homedir = pwd() + +# Discover runs that have calibrate output (ekp + results files). +valid_file_items = [] +valid_files = [] +for force_case in force_cases + for N_ens in N_enss + for rng_idx in rng_idxs + data_save_directory = joinpath(homedir, "output", calib_dir) + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + if isfile(ekp_file) && isfile(res_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (force_case, N_ens, rng_idx)) + end + end + end +end +@info "Plotting valid files:" +display(valid_files) +println(" ") + +# Load x0 from preliminaries (needed for computing panel-ii forcing). +prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(cfg.force_case).jld2") +if !isfile(prelim_file) + error("Preliminaries file not found: $(prelim_file). Run calibrate_l96.jl first.") +end +x0 = JLD2.load(prelim_file)["x0"] + +for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) + + @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + data_save_directory = joinpath(homedir, "output", calib_dir) + figure_save_directory = data_save_directory + + # --- load calibrate outputs --- + ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) + res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + pri_file = joinpath(data_save_directory, prior_filename(cfg)) + + ekpobj = load(ekp_file)["ekpobj"] + prior = load(pri_file)["prior"] + + (truth_params_obj, _) = load(res_file)["truth_params_structure"] + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + tp, _ = Flux.destructure(truth_params_obj.model) + (tp, truth_params_obj.model, truth_params_obj.sample_range) + end + + final_params_constrained = get_ϕ_mean_final(prior, ekpobj) + + N_ens = get_N_ens(ekpobj) + n_par = length(truth_params_constrained) + y = get_obs(ekpobj) + ny = length(y) + + # --- compute forcing quantities --- + is_const = force_case == "const-force" + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + final_emc = build_forcing(truth_params_obj, final_params_constrained, structure, sample_range) + final_forcing = forcing(final_emc, x0) + + # --- data plot --- + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + p1 = if is_const + hline( + [truth_params_constrained[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + range(0, n_par - 1, step = 1), + truth_params_constrained, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p_mid = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + plt = plot(p1, p_mid, p2, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) + + # --- solution plots (mean ensemble) --- + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + p1 = if is_const + hline( + [0.0], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + range(0, n_par - 1, step = 1), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Parameter difference (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p_mid = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + ylims = (0, 16), + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + p2 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + p3 = deepcopy(p1) + p_mid3 = deepcopy(p_mid) + p4 = deepcopy(p2) + + if is_const + hline!(p1, [final_params_constrained[1] - truth_params_constrained[1]], label = "mean ensemble input", color = :lightgreen, linewidth = 4) + hline!(p_mid, [final_forcing[1]], label = "mean ensemble forcing", color = :lightgreen, linewidth = 4) + else + plot!(p1, range(0, n_par - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) + plot!(p_mid, xaxis_forcing, final_forcing, label = "mean ensemble forcing", color = :lightgreen, linewidth = 4) + end + plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) + + plt = plot(p1, p_mid, p2, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) + + # --- full ensemble plot --- + ens_final_params = get_ϕ_final(prior, ekpobj) + ens_init_params = get_ϕ(prior, ekpobj, 1) + ens_final_forcings = hcat([forcing(build_forcing(truth_params_obj, ens_final_params[:, j], structure, sample_range), x0) for j in axes(ens_final_params, 2)]...) + ens_init_forcings = hcat([forcing(build_forcing(truth_params_obj, ens_init_params[:, j], structure, sample_range), x0) for j in axes(ens_init_params, 2)]...) + + if is_const + hline!(p3, [ens_final_params[1, 1] - truth_params_constrained[1]], label = "ensemble inputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, ens_final_params[1, 2:end] .- truth_params_constrained[1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, ens_init_params[1, :] .- truth_params_constrained[1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + hline!(p_mid3, [ens_final_forcings[1, 1]], label = "ensemble forcings", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p_mid3, ens_final_forcings[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p_mid3, ens_init_forcings[1, :], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + else + plot!(p3, range(0, n_par - 1, step = 1), ens_final_params[:, 1] .- truth_params_constrained, label = "ensemble inputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_final_params[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_init_params[:, 1] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, range(0, n_par - 1, step = 1), ens_init_params[:, 2:end] .- truth_params_constrained, label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + plot!(p_mid3, xaxis_forcing, ens_final_forcings[:, 1], label = "ensemble forcings", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_final_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_init_forcings[:, 1], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p_mid3, xaxis_forcing, ens_init_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + end + + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 1], color = :lightgreen, label = "ensemble outputs", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + plt = plot(p3, p_mid3, p4, layout = @layout [a b c]) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) + savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) + +end diff --git a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch new file mode 100644 index 000000000..d31fc0d3d --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash +# Single job that loads calibration output and produces diagnostic plots: +# data, solution (mean ensemble), and full-ensemble spread figures. +# Runs all (N_ens, rng_idx) cases serially in one job (no array needed). +# +# Invoked automatically by submit_l96_*.sh after calibrate completes. +# Manual use: +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_vec calibration_diagnostic_plots_l96.sbatch + +#SBATCH --job-name=calib_diag +#SBATCH --output=output/slurm/calib_diag_%j.out +#SBATCH --error=output/slurm/calib_diag_%j.err +#SBATCH --time=01:00:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=2 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +julia --project=. calibration_diagnostic_plots_l96.jl diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl b/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl deleted file mode 100644 index 44972e5cd..000000000 --- a/examples/EKIRace/hpc-variant/ensemble_from_posterior.jl +++ /dev/null @@ -1,185 +0,0 @@ -# Import modules -using Distributions -using LinearAlgebra -using Random -using JLD2 -using Statistics -using Flux -using BSON -using Dates -ENV["GKSwstype"] = "100" -using Plots -using Plots.Measures - - -# CES -using CalibrateEmulateSample.ParameterDistributions -using CalibrateEmulateSample.DataContainers -using CalibrateEmulateSample.EnsembleKalmanProcesses - -include("Lorenz96.jl") -include("experiment_config.jl") - -######################################################################## -############### Per-cell pushforward ################################### -######################################################################## - -function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) - n_samples_pushforward = 100 - - force_case = cfg.force_case - calib_dir = calib_directory(method, cfg) - calib_fn = results_filename(cfg, N_ens, rng_idx) - post_fn = posterior_filename(cfg, N_ens, rng_idx) - - # load preliminaries - prelim_dir = joinpath(@__DIR__, "output") - prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") - if !isfile(prelim_file) - throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) - end - loaded_data = JLD2.load(prelim_file) - x0 = loaded_data["x0"] - nx = length(x0) - y = loaded_data["y"] - ic_cov_sqrt = loaded_data["ic_cov_sqrt"] - R = loaded_data["R"] - lorenz_config_settings = loaded_data["lorenz_config_settings"] - observation_config = loaded_data["observation_config"] - - homedir = joinpath(pwd()) - data_save_directory = joinpath(homedir, "output", calib_dir) - - if !isfile(joinpath(data_save_directory, post_fn)) - @warn "No posterior file found for $(case_suffix(cfg, N_ens, rng_idx)); skipping." - return - end - - @info "loading case $(post_fn)" - loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) - loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) - - posteriors_by_k = loaded_p["posteriors_by_k"] - priors = loaded_p["priors"] - k_values = loaded_p["k_values"] - - (truth_params_obj, _) = loaded_c["truth_params_structure"] - - (truth_params_constrained, structure, sample_range) = if force_case == "const-force" - ([truth_params_obj.val], nothing, nothing) - elseif force_case == "vec-force" - (truth_params_obj.val, nothing, nothing) - elseif force_case == "flux-force" - tp, _ = destructure(truth_params_obj.model) - (tp, truth_params_obj.model, truth_params_obj.sample_range) - end - - for k in k_values - post_dist = posteriors_by_k[k] - push_ensemble = sample(post_dist, n_samples_pushforward) - # note: push_ensemble is unconstrained - constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) - - G_ens = hcat( - [ - lorenz_forward( - build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) for j in 1:n_samples_pushforward - ]..., - ) - - ny = length(y) - n_par = length(truth_params_constrained) - - truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) - truth_forcing = forcing(truth_emc, x0) - xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range - push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) - param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) - - gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p2 = plot( - collect(1:n_par), - zeros(n_par), - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "parameters(input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - p3 = plot( - xaxis_forcing, - truth_forcing, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - p4 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(R)), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - - plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - - plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) - plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) - - l = @layout [a b c] - plt = plot(p2, p3, p4, layout = l) - - suffix = case_suffix(cfg, N_ens, rng_idx) - figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" - savefig(plt, joinpath(data_save_directory, figure_fn)) - savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) - end -end - -######################################################################## -############### Main dispatcher ######################################## -######################################################################## - -function main() - exp = l96_experiment() - @assert( - exp in (:l96_const, :l96_vec, :l96_flux), - "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)", - ) - cfg = experiment_config(exp) - tasks = flat_tasks(cfg) - idx = task_index_from_args() - - if isnothing(idx) - for (N_ens, rng_idx) in tasks - ensemble_from_posterior_one(cfg, N_ens, rng_idx) - end - else - if idx < 1 || idx > length(tasks) - error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") - end - N_ens, rng_idx = tasks[idx] - ensemble_from_posterior_one(cfg, N_ens, rng_idx) - end -end - -main() diff --git a/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch b/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch deleted file mode 100644 index 69d451033..000000000 --- a/examples/EKIRace/hpc-variant/ensemble_from_posterior.sbatch +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/bash -# Single job that pushes all posteriors forward through the Lorenz96 forward map -# and saves pushforward plots. Runs all (N_ens, rng_idx) cases serially in one job -# (no array needed — the script loops internally). -# -# L96: -# sbatch --dependency=afterok: \ -# --export=ALL,EXPERIMENT=l96_vec ensemble_from_posterior.sbatch -# sbatch --dependency=afterok: \ -# --export=ALL,EXPERIMENT=l96_flux ensemble_from_posterior.sbatch - -#SBATCH --job-name=postfwd -#SBATCH --output=output/slurm/postfwd_%j.out -#SBATCH --error=output/slurm/postfwd_%j.err -#SBATCH --time=08:00:00 -#SBATCH --mem=16G -#SBATCH --cpus-per-task=4 -#SBATCH --ntasks=1 - -set -euo pipefail - -module load julia/1.12.2 - -export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} -export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} - -cd "${SLURM_SUBMIT_DIR}" - -export JULIA_PKG_PRECOMPILE_AUTO=0 - -julia --project=. ensemble_from_posterior.jl diff --git a/examples/EKIRace/hpc-variant/plot_l96_forcing.jl b/examples/EKIRace/hpc-variant/plot_l96_forcing.jl deleted file mode 100644 index 81fcd58f0..000000000 --- a/examples/EKIRace/hpc-variant/plot_l96_forcing.jl +++ /dev/null @@ -1,259 +0,0 @@ -# load packages -# CES - -using Random -using JLD2 -using Statistics -using LinearAlgebra -using Dates - -using Plots -using Plots.Measures - -using CalibrateEmulateSample.EnsembleKalmanProcesses -using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions -const EKP = EnsembleKalmanProcesses - -include("experiment_config.jl") -include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction - -#### CHOOSE YOUR CASE: -exp = l96_experiment() -@assert exp in (:l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_vec, or :l96_flux in experiment_config.jl" -cfg = experiment_config(exp) -method = method_cases[1] -calib_dir = calib_directory(method, cfg) -force_cases = [cfg.force_case] -N_enss = cfg.N_ens_sizes -rng_idxs = collect(1:cfg.n_repeats) - -homedir = pwd() - -# Discover runs that have calibrate output (ekp + results files). -valid_file_items = [] -valid_files = [] -for force_case in force_cases - for N_ens in N_enss - for rng_idx in rng_idxs - data_save_directory = joinpath(homedir, "output", calib_dir) - ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) - res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) - if isfile(ekp_file) && isfile(res_file) - push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) - push!(valid_file_items, (force_case, N_ens, rng_idx)) - end - end - end -end -@info "Plotting valid files:" -display(valid_files) -println(" ") - -for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) - - @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - - data_save_directory = joinpath(homedir, "output", calib_dir) - figure_save_directory = data_save_directory - - # --- load calibrate outputs --- - ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) - res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) - pri_file = joinpath(data_save_directory, prior_filename(cfg)) - - ekpobj = load(ekp_file)["ekpobj"] - prior = load(pri_file)["prior"] - - (truth_params_obj, _) = load(res_file)["truth_params_structure"] - truth_params_constrained = if force_case == "const-force" - [truth_params_obj.val] - elseif force_case == "vec-force" - truth_params_obj.val - elseif force_case == "flux-force" - tp, _ = Flux.destructure(truth_params_obj.model) - tp - end - - final_params_constrained = get_ϕ_mean_final(prior, ekpobj) - - N_ens = get_N_ens(ekpobj) - nx = length(truth_params_constrained) - y = get_obs(ekpobj) - ny = length(y) - - # --- data plot --- - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - truth_params_constrained, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) - - # --- solution plots (mean ensemble) --- - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - zeros(nx), - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p3 = deepcopy(p1) - p4 = deepcopy(p2) - - plot!(p1, range(0, nx - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) - plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params_constrained, - label = "ensemble inputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p4, - 1:length(y), - get_g_final(ekpobj)[:, 1], - color = :lightgreen, - label = "ensemble outputs", - linewidth = 4, - linealpha = 0.1, - ) - plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - - plt = plot(p3, p4, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - - # --- posterior ribbon plots (only if emulate_sample has been run) --- - post_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if !isfile(post_file) - @info "No posterior file found for $(calib_filename_suffix); skipping posterior_ribbons plots." - continue - end - - posteriors_by_k = load(post_file)["posteriors_by_k"] - k_values = load(post_file)["k_values"] - - for k in k_values - posterior = posteriors_by_k[k] - - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' - - gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - [truth_params_constrained final_params_constrained] .- truth_params_constrained, - label = ["solution" "EKI-opt"], - color = [:black :lightgreen], - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - plot!( - p1, - range(0, nx - 1, step = 1), - quantiles[:, 2] .- truth_params_constrained, - color = :blue, - label = "posterior", - linewidth = 4, - ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], - fillalpha = 0.1, - ) - - figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)_k$(k)") - savefig(figpath * ".png") - savefig(figpath * ".pdf") - end -end diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl new file mode 100644 index 000000000..42148cf2b --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -0,0 +1,343 @@ +# Import modules +using Distributions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates +ENV["GKSwstype"] = "100" +using Plots +using Plots.Measures + + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz96.jl") +include("experiment_config.jl") + +######################################################################## +############### Per-cell pushforward ################################### +######################################################################## + +function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) + n_samples_pushforward = 100 + + force_case = cfg.force_case + calib_dir = calib_directory(method, cfg) + calib_fn = results_filename(cfg, N_ens, rng_idx) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + + # load preliminaries + prelim_dir = joinpath(@__DIR__, "output") + prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") + if !isfile(prelim_file) + throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) + end + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + if !isfile(joinpath(data_save_directory, post_fn)) + @warn "No posterior file found for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + + (truth_params_obj, _) = loaded_c["truth_params_structure"] + + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + tp, _ = destructure(truth_params_obj.model) + (tp, truth_params_obj.model, truth_params_obj.sample_range) + end + + is_const = force_case == "const-force" + + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + + # --- sample from prior for grey background --- + prior_samples_unconstrained = sample(priors, n_samples_pushforward) + prior_samples_constrained = transform_unconstrained_to_constrained(priors, prior_samples_unconstrained) + prior_param_diffs = reduce(hcat, [prior_samples_constrained[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + prior_forcings = hcat([forcing(build_forcing(truth_params_obj, prior_samples_constrained[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + G_prior = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, prior_samples_constrained[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + for k in k_values + post_dist = posteriors_by_k[k] + push_ensemble = sample(post_dist, n_samples_pushforward) + # note: push_ensemble is unconstrained + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + G_ens = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + ny = length(y) + n_par = length(truth_params_constrained) + push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameters + p2 = if is_const + hline( + [0.0], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + # Panel (ii): forcing + p3 = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + # Panel (iii): output + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std output", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Add prior samples (grey) + if is_const + hline!(p2, [prior_param_diffs[1, 1]], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p2, prior_param_diffs[1, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p3, [prior_forcings[1, 1]], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p3, prior_forcings[1, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + else + plot!(p2, collect(1:n_par), prior_param_diffs[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, prior_forcings[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, prior_forcings[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + end + plot!(p4, 1:ny, G_prior[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + + # Add posterior samples (green) + if is_const + hline!(p2, [param_diffs[1, 1]], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p2, param_diffs[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, [push_forcings[1, 1]], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, push_forcings[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + else + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + end + plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + l = @layout [a b c] + plt = plot(p2, p3, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + # --- posterior ribbons plot --- + # Ribbons computed from the pushforward ensembles already in hand. + post_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(param_diffs)])' + post_forcing_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(push_forcings)])' + post_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_ens)])' + prior_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_param_diffs)])' + prior_forcing_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_forcings)])' + prior_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_prior)])' + + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameter ribbons + pa = if is_const + hline([0.0], label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "parameters (input)", + left_margin = 15mm, bottom_margin = 15mm) + else + plot(collect(1:n_par), zeros(n_par), label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "parameters (input)", + left_margin = 15mm, bottom_margin = 15mm) + end + + if is_const + hline!(pa, [prior_param_quantiles[1, 2]], color = :grey, label = "prior median", linewidth = 4) + hline!(pa, [prior_param_quantiles[1, 1]], color = :grey, label = "prior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pa, [prior_param_quantiles[1, 3]], color = :grey, label = "", linewidth = 2, linestyle = :dash) + hline!(pa, [post_param_quantiles[1, 2]], color = :blue, label = "posterior median", linewidth = 4) + hline!(pa, [post_param_quantiles[1, 1]], color = :blue, label = "posterior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pa, [post_param_quantiles[1, 3]], color = :blue, label = "", linewidth = 2, linestyle = :dash) + else + plot!(pa, collect(1:n_par), prior_param_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_param_quantiles[:, 2] - prior_param_quantiles[:, 1] prior_param_quantiles[:, 3] - prior_param_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pa, collect(1:n_par), post_param_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_param_quantiles[:, 2] - post_param_quantiles[:, 1] post_param_quantiles[:, 3] - post_param_quantiles[:, 2]], + fillalpha = 0.2) + end + + # Panel (ii): forcing ribbons + pb = if is_const + hline([truth_forcing[1]], label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "Forcing (transformed-input)", + ylims = (0, 16), left_margin = 15mm, bottom_margin = 15mm) + else + plot(xaxis_forcing, truth_forcing, label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "Forcing (transformed-input)", + ylims = (0, 16), left_margin = 15mm, bottom_margin = 15mm) + end + + if is_const + hline!(pb, [prior_forcing_quantiles[1, 2]], color = :grey, label = "prior median", linewidth = 4) + hline!(pb, [prior_forcing_quantiles[1, 1]], color = :grey, label = "prior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pb, [prior_forcing_quantiles[1, 3]], color = :grey, label = "", linewidth = 2, linestyle = :dash) + hline!(pb, [post_forcing_quantiles[1, 2]], color = :blue, label = "posterior median", linewidth = 4) + hline!(pb, [post_forcing_quantiles[1, 1]], color = :blue, label = "posterior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pb, [post_forcing_quantiles[1, 3]], color = :blue, label = "", linewidth = 2, linestyle = :dash) + else + plot!(pb, xaxis_forcing, prior_forcing_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_forcing_quantiles[:, 2] - prior_forcing_quantiles[:, 1] prior_forcing_quantiles[:, 3] - prior_forcing_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pb, xaxis_forcing, post_forcing_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_forcing_quantiles[:, 2] - post_forcing_quantiles[:, 1] post_forcing_quantiles[:, 3] - post_forcing_quantiles[:, 2]], + fillalpha = 0.2) + end + + # Panel (iii): output ribbons + pc = plot(1:ny, y, ribbon = sqrt.(diag(R)), label = "data", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "State mean/std (output)", + left_margin = 15mm, bottom_margin = 15mm) + plot!(pc, 1:ny, prior_G_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_G_quantiles[:, 2] - prior_G_quantiles[:, 1] prior_G_quantiles[:, 3] - prior_G_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pc, 1:ny, post_G_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_G_quantiles[:, 2] - post_G_quantiles[:, 1] post_G_quantiles[:, 3] - post_G_quantiles[:, 2]], + fillalpha = 0.2) + + ribbons_plt = plot(pa, pb, pc, layout = @layout [a b c]) + ribbons_fn = "posterior_ribbons_$(suffix)_k$(k)" + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".png")) + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".pdf")) + end +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + exp = l96_experiment() + @assert( + exp in (:l96_const, :l96_vec, :l96_flux), + "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)", + ) + cfg = experiment_config(exp) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end + else + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch new file mode 100644 index 000000000..e5ab9e885 --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -0,0 +1,30 @@ +#!/bin/bash +# Single job that pushes all posteriors forward through the Lorenz96 forward map, +# saves pushforward plots, and produces posterior-ribbon diagnostic figures. +# Runs all (N_ens, rng_idx) cases serially in one job (no array needed). +# +# Invoked automatically by submit_l96_*.sh after emulate_sample completes. +# Manual use: +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch + +#SBATCH --job-name=post_diag +#SBATCH --output=output/slurm/post_diag_%j.out +#SBATCH --error=output/slurm/post_diag_%j.err +#SBATCH --time=08:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +julia --project=. posterior_diagnostic_plots_l96.jl diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 30fc9e80f..62907e330 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Submit calibrate + emulate_sample for the L96 const-force case. +# Submit calibrate + emulate_sample + diagnostic plots for the L96 const-force case. # # Usage: bash submit_l96_const.sh [EXP_ID] # EXP_ID (optional): label appended to SLURM job names so the queue stays @@ -26,6 +26,16 @@ CALIB_JID=$(sbatch --parsable \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" +echo "=== Submitting calibration_diagnostic_plots (L96 const-force, after ${CALIB_JID}) ===" +CALIB_DIAG_JID=$(sbatch --parsable \ + -A esm \ + --job-name="calib_diag_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_const \ + calibration_diagnostic_plots_l96.sbatch) +echo " calibration_diagnostic_plots job ID: ${CALIB_DIAG_JID}" + echo "=== Submitting emulate_sample (L96 const-force, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ -A esm \ @@ -36,4 +46,14 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" +echo "=== Submitting posterior_diagnostic_plots (L96 const-force, after ${EMU_JID}) ===" +POST_DIAG_JID=$(sbatch --parsable \ + -A esm \ + --job-name="post_diag_${LABEL}" \ + --dependency=afterok:${EMU_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_const \ + posterior_diagnostic_plots_l96.sbatch) +echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 1fe7524df..634ac5b9f 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Submit calibrate + emulate_sample for the L96 flux-force case. +# Submit calibrate + emulate_sample + diagnostic plots for the L96 flux-force case. # # Usage: bash submit_l96_flux.sh [EXP_ID] # EXP_ID (optional): label appended to SLURM job names so the queue stays @@ -26,6 +26,16 @@ CALIB_JID=$(sbatch --parsable \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" +echo "=== Submitting calibration_diagnostic_plots (L96 flux-force, after ${CALIB_JID}) ===" +CALIB_DIAG_JID=$(sbatch --parsable \ + -A esm \ + --job-name="calib_diag_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_flux \ + calibration_diagnostic_plots_l96.sbatch) +echo " calibration_diagnostic_plots job ID: ${CALIB_DIAG_JID}" + echo "=== Submitting emulate_sample (L96 flux-force, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ -A esm \ @@ -36,14 +46,14 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting ensemble_from_posterior (L96 flux-force, after ${EMU_JID}) ===" -POST_JID=$(sbatch --parsable \ +echo "=== Submitting posterior_diagnostic_plots (L96 flux-force, after ${EMU_JID}) ===" +POST_DIAG_JID=$(sbatch --parsable \ -A esm \ - --job-name="post_${LABEL}" \ + --job-name="post_diag_${LABEL}" \ --dependency=afterok:${EMU_JID} \ --kill-on-invalid-dep=yes \ --export=ALL,EXPERIMENT=l96_flux \ - ensemble_from_posterior.sbatch) -echo " ensemble_from_posterior job ID: ${POST_JID}" + posterior_diagnostic_plots_l96.sbatch) +echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index e07c01fd3..e5b840e32 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Submit calibrate + emulate_sample for the L96 vec-force case. +# Submit calibrate + emulate_sample + diagnostic plots for the L96 vec-force case. # # Usage: bash submit_l96_vec.sh [EXP_ID] # EXP_ID (optional): label appended to SLURM job names so the queue stays @@ -26,6 +26,16 @@ CALIB_JID=$(sbatch --parsable \ calibrate_array.sbatch) echo " calibrate job ID: ${CALIB_JID}" +echo "=== Submitting calibration_diagnostic_plots (L96 vec-force, after ${CALIB_JID}) ===" +CALIB_DIAG_JID=$(sbatch --parsable \ + -A esm \ + --job-name="calib_diag_${LABEL}" \ + --dependency=afterok:${CALIB_JID} \ + --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_vec \ + calibration_diagnostic_plots_l96.sbatch) +echo " calibration_diagnostic_plots job ID: ${CALIB_DIAG_JID}" + echo "=== Submitting emulate_sample (L96 vec-force, after ${CALIB_JID}) ===" EMU_JID=$(sbatch --parsable \ -A esm \ @@ -36,14 +46,14 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting ensemble_from_posterior (L96 vec-force, after ${EMU_JID}) ===" -POST_JID=$(sbatch --parsable \ +echo "=== Submitting posterior_diagnostic_plots (L96 vec-force, after ${EMU_JID}) ===" +POST_DIAG_JID=$(sbatch --parsable \ -A esm \ - --job-name="post_${LABEL}" \ + --job-name="post_diag_${LABEL}" \ --dependency=afterok:${EMU_JID} \ --kill-on-invalid-dep=yes \ --export=ALL,EXPERIMENT=l96_vec \ - ensemble_from_posterior.sbatch) -echo " ensemble_from_posterior job ID: ${POST_JID}" + posterior_diagnostic_plots_l96.sbatch) +echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/plot_l96_forcing.jl b/examples/EKIRace/plot_l96_forcing.jl deleted file mode 100644 index f5c005fbc..000000000 --- a/examples/EKIRace/plot_l96_forcing.jl +++ /dev/null @@ -1,258 +0,0 @@ -# load packages -# CES - -using Random -using JLD2 -using Statistics -using LinearAlgebra -using Dates - -using Plots -using Plots.Measures - -using CalibrateEmulateSample.EnsembleKalmanProcesses -using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions -const EKP = EnsembleKalmanProcesses - -include("experiment_config.jl") -include("Lorenz96.jl") # provides Flux.destructure for flux-force truth_params extraction - -#### CHOOSE YOUR CASE: -@assert EXPERIMENT in (:l96_vec, :l96_flux) "For plot_l96_forcing.jl, set EXPERIMENT to :l96_vec, or :l96_flux in experiment_config.jl" -cfg = experiment_config(EXPERIMENT) -method = method_cases[1] -calib_dir = calib_directory(method, cfg) -force_cases = [cfg.force_case] -N_enss = cfg.N_ens_sizes -rng_idxs = collect(1:cfg.n_repeats) - -homedir = pwd() - -# Discover runs that have calibrate output (ekp + results files). -valid_file_items = [] -valid_files = [] -for force_case in force_cases - for N_ens in N_enss - for rng_idx in rng_idxs - data_save_directory = joinpath(homedir, "output", calib_dir) - ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) - res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) - if isfile(ekp_file) && isfile(res_file) - push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) - push!(valid_file_items, (force_case, N_ens, rng_idx)) - end - end - end -end -@info "Plotting valid files:" -display(valid_files) -println(" ") - -for ((force_case, N_ens, rng_idx), calib_filename_suffix) in zip(valid_file_items, valid_files) - - @info("Plotting for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") - - data_save_directory = joinpath(homedir, "output", calib_dir) - figure_save_directory = data_save_directory - - # --- load calibrate outputs --- - ekp_file = joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx)) - res_file = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) - pri_file = joinpath(data_save_directory, prior_filename(cfg)) - - ekpobj = load(ekp_file)["ekpobj"] - prior = load(pri_file)["prior"] - - (truth_params_obj, _) = load(res_file)["truth_params_structure"] - truth_params_constrained = if force_case == "const-force" - [truth_params_obj.val] - elseif force_case == "vec-force" - truth_params_obj.val - elseif force_case == "flux-force" - tp, _ = Flux.destructure(truth_params_obj.model) - tp - end - - final_params_constrained = get_ϕ_mean_final(prior, ekpobj) - - N_ens = get_N_ens(ekpobj) - nx = length(truth_params_constrained) - y = get_obs(ekpobj) - ny = length(y) - - # --- data plot --- - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - truth_params_constrained, - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "data_$(calib_filename_suffix).pdf")) - - # --- solution plots (mean ensemble) --- - gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - zeros(nx), - label = "solution", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p2 = plot( - 1:ny, - y, - ribbon = sqrt.(diag(get_obs_noise_cov(ekpobj))), - label = "data", - color = :black, - linewidth = 4, - xlabel = "Spatial index", - ylabel = "State mean/std output)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - p3 = deepcopy(p1) - p4 = deepcopy(p2) - - plot!(p1, range(0, nx - 1, step = 1), final_params_constrained .- truth_params_constrained, label = "mean ensemble input", color = :lightgreen, linewidth = 4) - plot!(p2, 1:length(y), get_g_mean_final(ekpobj), label = "mean ensemble output", color = :lightgreen, linewidth = 4) - - l = @layout [a b] - plt = plot(p1, p2, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix).pdf")) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 1] .- truth_params_constrained, - label = "ensemble inputs", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ_final(prior, ekpobj)[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 1] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p3, - range(0, nx - 1, step = 1), - get_ϕ(prior, ekpobj, 1)[:, 2:end] .- truth_params_constrained, - label = "", - color = :lightgreen, - linewidth = 4, - linealpha = 0.1, - ) - - plot!( - p4, - 1:length(y), - get_g_final(ekpobj)[:, 1], - color = :lightgreen, - label = "ensemble outputs", - linewidth = 4, - linealpha = 0.1, - ) - plot!(p4, 1:length(y), get_g_final(ekpobj)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 1], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - plot!(p4, 1:length(y), get_g(ekpobj, 1)[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) - - plt = plot(p3, p4, layout = l) - - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.png")) - savefig(plt, joinpath(figure_save_directory, "solution_$(calib_filename_suffix)_full_ens.pdf")) - - # --- posterior ribbon plots (only if emulate_sample has been run) --- - post_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) - if !isfile(post_file) - @info "No posterior file found for $(calib_filename_suffix); skipping posterior_ribbons plots." - continue - end - - posteriors_by_k = load(post_file)["posteriors_by_k"] - k_values = load(post_file)["k_values"] - - for k in k_values - posterior = posteriors_by_k[k] - - posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) - constrained_posterior_samples = - mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) - quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(constrained_posterior_samples)])' - - gr(size = (1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) - p1 = plot( - range(0, nx - 1, step = 1), - [truth_params_constrained final_params_constrained] .- truth_params_constrained, - label = ["solution" "EKI-opt"], - color = [:black :lightgreen], - linewidth = 4, - xlabel = "Spatial index", - ylabel = "Forcing difference (input)", - left_margin = 15mm, - bottom_margin = 15mm, - ) - - plot!( - p1, - range(0, nx - 1, step = 1), - quantiles[:, 2] .- truth_params_constrained, - color = :blue, - label = "posterior", - linewidth = 4, - ribbon = [quantiles[:, 2] - quantiles[:, 1] quantiles[:, 3] - quantiles[:, 2]], - fillalpha = 0.1, - ) - - figpath = joinpath(figure_save_directory, "posterior_ribbons_$(calib_filename_suffix)_k$(k)") - savefig(figpath * ".png") - savefig(figpath * ".pdf") - end -end diff --git a/examples/EKIRace/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl new file mode 100644 index 000000000..7c9947165 --- /dev/null +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -0,0 +1,368 @@ +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Flux +using BSON +using Dates +using Plots +using Plots.Measures + + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz96.jl") # Contains Lorenz 96 source code +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true + +n_samples_pushforward = 100 # num. samples to pushforward through lorenz + + +######################################################################## +################## file-certainty for loading ########################## +######################################################################## + +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "For calibrate_l96.jl, set EXPERIMENT to :l96_const, :l96_vec, or :l96_flux in experiment_config.jl" +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] # method_cases defined in experiment_config.jl +calib_dir = calib_directory(method, cfg) +force_case = cfg.force_case +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +# get some preliminaries +prelim_dir = joinpath(@__DIR__, "output") +if !isdir(prelim_dir) + mkdir(prelim_dir) +end +prelim_file = joinpath(prelim_dir, "l96_computed_preliminaries_$(force_case).jld2") +if isfile(prelim_file) + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + R_inv_var = loaded_data["R_inv_var"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + @info "loaded precomputed preliminary quantities from $(prelim_file)" +else + throw(ErrorException("preliminaries files not found. \n First run: \n > julia --project calibrate_l96.jl")) +end +### determine valid files before we try to load + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Pushing forward posteriors through the forward map from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l96.jl first.") +end + +### Then load data + +# Load first valid file to determine parameter dimension and max k +first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +for (N_ens, rng_idx) in valid_file_items + calib_fn = results_filename(cfg, N_ens, rng_idx) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] + + (truth_params_obj, _) = loaded_c["truth_params_structure"] + + (truth_params_constrained, structure, sample_range) = if force_case == "const-force" + ([truth_params_obj.val], nothing, nothing) + elseif force_case == "vec-force" + (truth_params_obj.val, nothing, nothing) + elseif force_case == "flux-force" + truth_params_constrained, _ = destructure(truth_params_obj.model) + + ( + truth_params_constrained, + truth_params_obj.model, + truth_params_obj.sample_range + ) + end + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + + is_const = force_case == "const-force" + + truth_emc = build_forcing(truth_params_obj, truth_params_constrained, structure, sample_range) + truth_forcing = forcing(truth_emc, x0) + xaxis_forcing = isnothing(sample_range) ? range(0, length(truth_forcing) - 1, step = 1) : sample_range + + # --- sample from prior for grey background --- + prior_samples_unconstrained = sample(priors, n_samples_pushforward) + prior_samples_constrained = transform_unconstrained_to_constrained(priors, prior_samples_unconstrained) + prior_param_diffs = reduce(hcat, [prior_samples_constrained[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + prior_forcings = hcat([forcing(build_forcing(truth_params_obj, prior_samples_constrained[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + G_prior = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, prior_samples_constrained[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + for k in k_values + post_dist = posteriors_by_k[k] + push_ensemble = sample(post_dist, n_samples_pushforward) + # note: push_ensemble is unconstrained + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + G_ens = hcat( + [ + lorenz_forward( + build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + ny = length(y) + n_par = length(truth_params_constrained) + push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameters + p2 = if is_const + hline( + [0.0], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + # Panel (ii): forcing + p3 = if is_const + hline( + [truth_forcing[1]], + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + else + plot( + xaxis_forcing, + truth_forcing, + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "Forcing (transformed-input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + end + + # Panel (iii): output + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Spatial index", + ylabel = "State mean/std (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Add prior samples (grey) + if is_const + hline!(p2, [prior_param_diffs[1, 1]], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p2, prior_param_diffs[1, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p3, [prior_forcings[1, 1]], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + hline!(p3, prior_forcings[1, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + else + plot!(p2, collect(1:n_par), prior_param_diffs[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, prior_forcings[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, prior_forcings[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + end + plot!(p4, 1:ny, G_prior[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + + # Add posterior samples (green) + if is_const + hline!(p2, [param_diffs[1, 1]], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p2, param_diffs[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, [push_forcings[1, 1]], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + hline!(p3, push_forcings[1, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + else + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p3, xaxis_forcing, push_forcings[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + end + plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 2:end], color = :lightgreen, label = "", linewidth = 4, linealpha = 0.1) + + l = @layout [a b c] + plt = plot(p2, p3, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + # --- posterior ribbons plot --- + # Ribbons computed from the pushforward ensembles already in hand. + post_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(param_diffs)])' + post_forcing_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(push_forcings)])' + post_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_ens)])' + prior_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_param_diffs)])' + prior_forcing_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_forcings)])' + prior_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_prior)])' + + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameter ribbons + pa = if is_const + hline([0.0], label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "parameters (input)", + left_margin = 15mm, bottom_margin = 15mm) + else + plot(collect(1:n_par), zeros(n_par), label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "parameters (input)", + left_margin = 15mm, bottom_margin = 15mm) + end + + if is_const + hline!(pa, [prior_param_quantiles[1, 2]], color = :grey, label = "prior median", linewidth = 4) + hline!(pa, [prior_param_quantiles[1, 1]], color = :grey, label = "prior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pa, [prior_param_quantiles[1, 3]], color = :grey, label = "", linewidth = 2, linestyle = :dash) + hline!(pa, [post_param_quantiles[1, 2]], color = :blue, label = "posterior median", linewidth = 4) + hline!(pa, [post_param_quantiles[1, 1]], color = :blue, label = "posterior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pa, [post_param_quantiles[1, 3]], color = :blue, label = "", linewidth = 2, linestyle = :dash) + else + plot!(pa, collect(1:n_par), prior_param_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_param_quantiles[:, 2] - prior_param_quantiles[:, 1] prior_param_quantiles[:, 3] - prior_param_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pa, collect(1:n_par), post_param_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_param_quantiles[:, 2] - post_param_quantiles[:, 1] post_param_quantiles[:, 3] - post_param_quantiles[:, 2]], + fillalpha = 0.2) + end + + # Panel (ii): forcing ribbons + pb = if is_const + hline([truth_forcing[1]], label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "Forcing (transformed-input)", + ylims = (0, 16), left_margin = 15mm, bottom_margin = 15mm) + else + plot(xaxis_forcing, truth_forcing, label = "solution", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "Forcing (transformed-input)", + ylims = (0, 16), left_margin = 15mm, bottom_margin = 15mm) + end + + if is_const + hline!(pb, [prior_forcing_quantiles[1, 2]], color = :grey, label = "prior median", linewidth = 4) + hline!(pb, [prior_forcing_quantiles[1, 1]], color = :grey, label = "prior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pb, [prior_forcing_quantiles[1, 3]], color = :grey, label = "", linewidth = 2, linestyle = :dash) + hline!(pb, [post_forcing_quantiles[1, 2]], color = :blue, label = "posterior median", linewidth = 4) + hline!(pb, [post_forcing_quantiles[1, 1]], color = :blue, label = "posterior 5/95%", linewidth = 2, linestyle = :dash) + hline!(pb, [post_forcing_quantiles[1, 3]], color = :blue, label = "", linewidth = 2, linestyle = :dash) + else + plot!(pb, xaxis_forcing, prior_forcing_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_forcing_quantiles[:, 2] - prior_forcing_quantiles[:, 1] prior_forcing_quantiles[:, 3] - prior_forcing_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pb, xaxis_forcing, post_forcing_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_forcing_quantiles[:, 2] - post_forcing_quantiles[:, 1] post_forcing_quantiles[:, 3] - post_forcing_quantiles[:, 2]], + fillalpha = 0.2) + end + + # Panel (iii): output ribbons + pc = plot(1:ny, y, ribbon = sqrt.(diag(R)), label = "data", color = :black, linewidth = 4, + xlabel = "Spatial index", ylabel = "State mean/std (output)", + left_margin = 15mm, bottom_margin = 15mm) + plot!(pc, 1:ny, prior_G_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_G_quantiles[:, 2] - prior_G_quantiles[:, 1] prior_G_quantiles[:, 3] - prior_G_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pc, 1:ny, post_G_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_G_quantiles[:, 2] - post_G_quantiles[:, 1] post_G_quantiles[:, 3] - post_G_quantiles[:, 2]], + fillalpha = 0.2) + + ribbons_plt = plot(pa, pb, pc, layout = @layout [a b c]) + ribbons_fn = "posterior_ribbons_$(suffix)_k$(k)" + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".png")) + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".pdf")) + + end +end + + From bfdb4c8acd7c01c6ec91a9aa41b5b36ae1ad1f17 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 2 Jun 2026 17:31:12 -0700 Subject: [PATCH 53/86] updated config and emulation --- examples/EKIRace/calibrate_l96.jl | 2 +- examples/EKIRace/emulate_sample_l96.jl | 13 ++++++++++--- examples/EKIRace/experiment_config.jl | 10 +++++----- examples/EKIRace/hpc-variant/calibrate_l96.jl | 2 +- examples/EKIRace/hpc-variant/emulate_sample_l96.jl | 10 ++++++++-- examples/EKIRace/hpc-variant/experiment_config.jl | 6 +++--- 6 files changed, 28 insertions(+), 15 deletions(-) diff --git a/examples/EKIRace/calibrate_l96.jl b/examples/EKIRace/calibrate_l96.jl index 948dbcd83..79d2ae795 100644 --- a/examples/EKIRace/calibrate_l96.jl +++ b/examples/EKIRace/calibrate_l96.jl @@ -121,7 +121,7 @@ elseif case == "flux-force" # display(p) #prior_cov = (0.1^2) * I(length(prior_mean)) # original cov - prior_cov = (1.0^2) * Diagonal(prior_mean.^2) # scaled to mean + prior_cov = (0.1^2) * Diagonal(prior_mean.^2) # scaled to mean distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) name = "l96_nn_prior" diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl index 30011025d..45b48d545 100644 --- a/examples/EKIRace/emulate_sample_l96.jl +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -110,12 +110,19 @@ function main() kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), optimizer_options = overrides, ) - encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] +# encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + + encoder_schedule = [ + (decorrelate_structure_mat(), "in_and_out"), + (likelihood_informed(retain_info=retain_var, iters=1:k),"in"), + (likelihood_informed(retain_info=retain_var),"out"), # iters=1:1 for output space + ] + emulator = Emulator( mlt, input_output_pairs; encoder_schedule = deepcopy(encoder_schedule), - encoder_kwargs = encoder_kwargs, + encoder_kwargs = deepcopy(encoder_kwargs), ) optimize_hyperparameters!(emulator) @@ -127,7 +134,7 @@ function main() ### Sample: MCMC u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) - mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + mcmc = MCMCWrapper(pCNMHSampling(), truth_sample, priors, emulator; init_params = u0) new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) println("k=$(k) | Begin MCMC - step size ", new_step) chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) diff --git a/examples/EKIRace/experiment_config.jl b/examples/EKIRace/experiment_config.jl index e6edba6e2..cee4c77c4 100644 --- a/examples/EKIRace/experiment_config.jl +++ b/examples/EKIRace/experiment_config.jl @@ -4,11 +4,11 @@ using Dates ############### USER TOGGLE ######################################### ######################################################################## # Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux -EXPERIMENT = :l96_const +EXPERIMENT = :l96_flux # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). -#calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +#calibrate_date = Date("2026-05-31", "yyyy-mm-dd") calibrate_date = today() ######################################################################## @@ -69,7 +69,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, @@ -83,7 +83,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, @@ -97,7 +97,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, diff --git a/examples/EKIRace/hpc-variant/calibrate_l96.jl b/examples/EKIRace/hpc-variant/calibrate_l96.jl index 5083f6452..b10cb724b 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l96.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -78,7 +78,7 @@ function build_setup(cfg, output_dir) prior_train = prior_sinusoid.(x_train) .+ 0.2 .* randn(length(x_train)) prior_model, prior_mean = train_network(deepcopy(phi_structure), x_train, prior_train) #prior_cov = (0.1^2) * I(length(prior_mean)) # original cov - prior_cov = (1.0^2) * Diagonal(prior_mean.^2) # scaled to mean + prior_cov = (0.1^2) * Diagonal(prior_mean.^2) # scaled to mean distribution = Parameterized(MvNormal(prior_mean, prior_cov)) constraint = repeat([no_constraint()], 61) prior = ParameterDistribution(distribution, constraint, "l96_nn_prior") diff --git a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl index 6abb89167..6759f9e4a 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_l96.jl +++ b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl @@ -99,7 +99,13 @@ function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) kernel_structure = SeparableKernel(DiagonalFactor(nugget), OneDimFactor()), optimizer_options = overrides, ) - encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + # encoder_schedule = [(decorrelate_structure_mat(retain_var = retain_var), "in_and_out")] + + encoder_schedule = [ + (decorrelate_structure_mat(), "in_and_out"), + (likelihood_informed(retain_info=retain_var, iters=1:k),"in"), + (likelihood_informed(retain_info=retain_var),"out"), # iters=1:1 for output space + ] emulator = Emulator( mlt, input_output_pairs; @@ -116,7 +122,7 @@ function emulate_sample_one(cfg, N_ens, rng_idx; method = method_cases[1]) ### Sample: MCMC u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) - mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + mcmc = MCMCWrapper(pCNMHSampling(), truth_sample, priors, emulator; init_params = u0) new_step = optimize_stepsize(mcmc; init_stepsize = 0.2, N = 2000, discard_initial = 0) println("k=$(k) | Begin MCMC - step size ", new_step) chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl index 78bb43f2f..52bcdcd35 100644 --- a/examples/EKIRace/hpc-variant/experiment_config.jl +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -71,7 +71,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, @@ -85,7 +85,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, @@ -99,7 +99,7 @@ function experiment_config(case::Symbol) terminate_at = 2.0, n_repeats = 20, max_iter = 15, - retain_var = 0.95, + retain_var = 0.99, n_features = 200, n_features_opt = 160, calibrate_date = calibrate_date, From 08f622cf0b70fc0d1c4498833fef507d35f844c8 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 2 Jun 2026 17:33:55 -0700 Subject: [PATCH 54/86] today() --- examples/EKIRace/hpc-variant/experiment_config.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl index 52bcdcd35..d1f9d354f 100644 --- a/examples/EKIRace/hpc-variant/experiment_config.jl +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -10,8 +10,8 @@ EXPERIMENT = experiments[4] # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). # PIN this before submitting an array job so all tasks use the same directory. -calibrate_date = Date("2026-05-29", "yyyy-mm-dd") -#calibrate_date = today() +#calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +calibrate_date = today() ######################################################################## ############### SHARED CONSTANTS #################################### From 65d2f3915a8974a1ce7467d595af0666dea66093 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 11:01:47 -0700 Subject: [PATCH 55/86] reduced requirements --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- .../hpc-variant/calibration_diagnostic_plots_l96.sbatch | 4 ++-- .../EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch | 4 ++-- examples/EKIRace/hpc-variant/submit_l96_const.sh | 2 ++ examples/EKIRace/hpc-variant/submit_l96_flux.sh | 2 ++ examples/EKIRace/hpc-variant/submit_l96_vec.sh | 2 ++ 6 files changed, 11 insertions(+), 5 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 36973b7ce..7ca5ae8c9 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -21,7 +21,7 @@ #SBATCH --output=output/slurm/calib_%A_%a.out #SBATCH --error=output/slurm/calib_%A_%a.err #SBATCH --array=1-60%100 -#SBATCH --time=04:00:00 +#SBATCH --time=02:00:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 diff --git a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch index d31fc0d3d..b6176cdd3 100644 --- a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch @@ -11,9 +11,9 @@ #SBATCH --job-name=calib_diag #SBATCH --output=output/slurm/calib_diag_%j.out #SBATCH --error=output/slurm/calib_diag_%j.err -#SBATCH --time=01:00:00 +#SBATCH --time=00:30:00 #SBATCH --mem=8G -#SBATCH --cpus-per-task=2 +#SBATCH --cpus-per-task=1 #SBATCH --ntasks=1 set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index e5ab9e885..e152db7e3 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -11,9 +11,9 @@ #SBATCH --job-name=post_diag #SBATCH --output=output/slurm/post_diag_%j.out #SBATCH --error=output/slurm/post_diag_%j.err -#SBATCH --time=08:00:00 +#SBATCH --time=00:30:00 #SBATCH --mem=16G -#SBATCH --cpus-per-task=4 +#SBATCH --cpus-per-task=1 #SBATCH --ntasks=1 set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 62907e330..a5d929fd7 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -57,3 +57,5 @@ POST_DIAG_JID=$(sbatch --parsable \ echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" + +# sbatch -A esm --job-name="post_diag_l96_const" --export=ALL,EXPERIMENT=l96_const posterior_diagnostic_plots_l96.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 634ac5b9f..998fdd1fc 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -57,3 +57,5 @@ POST_DIAG_JID=$(sbatch --parsable \ echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" + +#sbatch --parsable -A esm --job-name="post_diag_l96_flux" --export=ALL,EXPERIMENT=l96_flux posterior_diagnostic_plots_l96.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index e5b840e32..8cd80dcff 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -57,3 +57,5 @@ POST_DIAG_JID=$(sbatch --parsable \ echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" + +# sbatch -A esm --job-name="post_diag_l96_vec" --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch From 983dcd8a9922892380c2dfd698f6fddaf085aba5 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 13:56:59 -0700 Subject: [PATCH 56/86] updated diagnostics to include the physically meaningful variables, saved in NetCDF --- .../l96_exp_to_leaderboard_utilities.jl | 73 +++++++++++++++++-- .../posterior_diagnostic_plots_l96.jl | 26 ++++++- .../l96_exp_to_leaderboard_utilities.jl | 73 +++++++++++++++++-- .../EKIRace/posterior_diagnostic_plots_l96.jl | 24 ++++++ 4 files changed, 185 insertions(+), 11 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index a69532fe7..0ac83daf6 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -3,11 +3,16 @@ using Dates using JLD2 using Distributions using LinearAlgebra +using Flux +using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz96.jl") + +n_pushforward_samples = 1000 # Step 1, Read and extract the available experiments @@ -60,6 +65,28 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) +# Load Lorenz preliminaries and truth structure for pushforward +prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") +prelim_data = JLD2.load(prelim_file) +x0 = prelim_data["x0"] +nx = length(x0) +ic_cov_sqrt = prelim_data["ic_cov_sqrt"] +lorenz_config_settings = prelim_data["lorenz_config_settings"] +observation_config = prelim_data["observation_config"] +n_output = 2 * nx + +first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) +(truth_phi, _) = first_calib["truth_params_structure"] +(phi_structure, sample_range) = if force_case == "const-force" + (nothing, nothing) +elseif force_case == "vec-force" + (nothing, nothing) +elseif force_case == "flux-force" + (truth_phi.model, truth_phi.sample_range) +end +n_forcing = force_case == "flux-force" ? length(sample_range) : nx + # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) n_evals_arr = fill(NaN, n_rng, n_ens) @@ -67,6 +94,8 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -103,6 +132,21 @@ for (N_ens, rng_idx) in valid_file_items mahal_arr[i, j, k] = diff' * (C_reg \ diff) # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + # pushforward ensemble for forcing and output space + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" + for s in 1:n_pushforward_samples + emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) + forcing_samples_arr[i, j, k, s, :] = forcing(emc, x0) + output_samples_arr[i, j, k, s, :] = lorenz_forward( + emc, + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end end end @@ -111,11 +155,14 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "forcing_dim", n_forcing) +defDim(ds, "output_dim", n_output) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -179,5 +226,21 @@ posterior_logpdf_true_v_map_v = defVar( posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr +forcing_samples_v = defVar( + ds, "forcing_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "forcing_dim"); + fillvalue = NaN, +) +forcing_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to forcing space" +forcing_samples_v[:, :, :, :, :] = forcing_samples_arr + +output_samples_v = defVar( + ds, "output_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); + fillvalue = NaN, +) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz96 output space (mean/std per spatial point)" +output_samples_v[:, :, :, :, :] = output_samples_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl index 42148cf2b..1c11d690a 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -24,6 +24,19 @@ include("experiment_config.jl") ############### Per-cell pushforward ################################### ######################################################################## +function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) + m = vec(mean(samples, dims=2)) + C_raw = cov(samples') + λ = max(1e-10, 1e-4 * mean(diag(C_raw))) + C = Symmetric(C_raw + λ * I) + dist = MvNormal(m, C) + pmode = samples[:, argmax(logpdf(dist, samples))] + diff = m - truth + mah = diff' * (C \ diff) + lp = logpdf(dist, truth) - logpdf(dist, pmode) + return mah, lp +end + function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) n_samples_pushforward = 100 @@ -118,6 +131,17 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces + frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings + truth_frc_m = is_const ? truth_forcing[1:1] : truth_forcing + param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) + forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) + output_mah, output_lp = pushforward_metrics(G_ens, y) + @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" + @info " param [d=$(n_par)]: mahal=$(round(param_mah, digits=2)) logpdf_ratio=$(round(param_lp, digits=2))" + @info " forcing [d=$(length(truth_frc_m))]: mahal=$(round(forcing_mah, digits=2)) logpdf_ratio=$(round(forcing_lp, digits=2))" + @info " output [d=$(ny)]: mahal=$(round(output_mah, digits=2)) logpdf_ratio=$(round(output_lp, digits=2))" + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) # Panel (i): parameters @@ -181,7 +205,7 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ color = :black, linewidth = 4, xlabel = "Spatial index", - ylabel = "State mean/std output", + ylabel = "State mean/std (output)", left_margin = 15mm, bottom_margin = 15mm, ) diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index a69532fe7..0ac83daf6 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -3,11 +3,16 @@ using Dates using JLD2 using Distributions using LinearAlgebra +using Flux +using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz96.jl") + +n_pushforward_samples = 1000 # Step 1, Read and extract the available experiments @@ -60,6 +65,28 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) +# Load Lorenz preliminaries and truth structure for pushforward +prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") +prelim_data = JLD2.load(prelim_file) +x0 = prelim_data["x0"] +nx = length(x0) +ic_cov_sqrt = prelim_data["ic_cov_sqrt"] +lorenz_config_settings = prelim_data["lorenz_config_settings"] +observation_config = prelim_data["observation_config"] +n_output = 2 * nx + +first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) +(truth_phi, _) = first_calib["truth_params_structure"] +(phi_structure, sample_range) = if force_case == "const-force" + (nothing, nothing) +elseif force_case == "vec-force" + (nothing, nothing) +elseif force_case == "flux-force" + (truth_phi.model, truth_phi.sample_range) +end +n_forcing = force_case == "flux-force" ? length(sample_range) : nx + # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) n_evals_arr = fill(NaN, n_rng, n_ens) @@ -67,6 +94,8 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -103,6 +132,21 @@ for (N_ens, rng_idx) in valid_file_items mahal_arr[i, j, k] = diff' * (C_reg \ diff) # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + # pushforward ensemble for forcing and output space + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" + for s in 1:n_pushforward_samples + emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) + forcing_samples_arr[i, j, k, s, :] = forcing(emc, x0) + output_samples_arr[i, j, k, s, :] = lorenz_forward( + emc, + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end end end @@ -111,11 +155,14 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "forcing_dim", n_forcing) +defDim(ds, "output_dim", n_output) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -179,5 +226,21 @@ posterior_logpdf_true_v_map_v = defVar( posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr +forcing_samples_v = defVar( + ds, "forcing_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "forcing_dim"); + fillvalue = NaN, +) +forcing_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to forcing space" +forcing_samples_v[:, :, :, :, :] = forcing_samples_arr + +output_samples_v = defVar( + ds, "output_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); + fillvalue = NaN, +) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz96 output space (mean/std per spatial point)" +output_samples_v[:, :, :, :, :] = output_samples_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl index 7c9947165..6d236112c 100644 --- a/examples/EKIRace/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -96,6 +96,19 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) +function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) + m = vec(mean(samples, dims=2)) + C_raw = cov(samples') + λ = max(1e-10, 1e-4 * mean(diag(C_raw))) + C = Symmetric(C_raw + λ * I) + dist = MvNormal(m, C) + pmode = samples[:, argmax(logpdf(dist, samples))] + diff = m - truth + mah = diff' * (C \ diff) + lp = logpdf(dist, truth) - logpdf(dist, pmode) + return mah, lp +end + for (N_ens, rng_idx) in valid_file_items calib_fn = results_filename(cfg, N_ens, rng_idx) post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -169,6 +182,17 @@ for (N_ens, rng_idx) in valid_file_items push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces + frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings + truth_frc_m = is_const ? truth_forcing[1:1] : truth_forcing + param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) + forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) + output_mah, output_lp = pushforward_metrics(G_ens, y) + @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" + @info " param [d=$(n_par)]: mahal=$(round(param_mah, digits=2)) logpdf_ratio=$(round(param_lp, digits=2))" + @info " forcing [d=$(length(truth_frc_m))]: mahal=$(round(forcing_mah, digits=2)) logpdf_ratio=$(round(forcing_lp, digits=2))" + @info " output [d=$(ny)]: mahal=$(round(output_mah, digits=2)) logpdf_ratio=$(round(output_lp, digits=2))" + gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) # Panel (i): parameters From 3956660fabd09f145a7f2c7186eb022a55d5f7d7 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 15:38:54 -0700 Subject: [PATCH 57/86] add leaderboard conversions to the pipelines --- .../hpc-variant/l96_exp_to_leaderboard_utilities.jl | 2 ++ examples/EKIRace/hpc-variant/submit_l63.sh | 8 ++++++++ examples/EKIRace/hpc-variant/submit_l96_const.sh | 12 ++++++++++-- examples/EKIRace/hpc-variant/submit_l96_flux.sh | 12 ++++++++++-- examples/EKIRace/hpc-variant/submit_l96_vec.sh | 12 ++++++++++-- 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index 0ac83daf6..b9ff6655c 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -12,6 +12,8 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") include("Lorenz96.jl") +EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg + n_pushforward_samples = 1000 # Step 1, Read and extract the available experiments diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index ae0bc1365..1e9c9ac17 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -36,4 +36,12 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" +echo "=== Submitting exp_to_leaderboard (L63, after ${EMU_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + exp_to_leaderboard.sbatch) +echo " exp_to_leaderboard job ID: ${LB_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index a5d929fd7..65121dce8 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -50,12 +50,20 @@ echo "=== Submitting posterior_diagnostic_plots (L96 const-force, after ${EMU_JI POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterok:${EMU_JID} \ - --kill-on-invalid-dep=yes \ + --dependency=afterany:${EMU_JID} \ --export=ALL,EXPERIMENT=l96_const \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" +echo "=== Submitting exp_to_leaderboard (L96 const-force, after ${EMU_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + exp_to_leaderboard.sbatch) +echo " exp_to_leaderboard job ID: ${LB_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch -A esm --job-name="post_diag_l96_const" --export=ALL,EXPERIMENT=l96_const posterior_diagnostic_plots_l96.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index 998fdd1fc..da7caacfd 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -50,12 +50,20 @@ echo "=== Submitting posterior_diagnostic_plots (L96 flux-force, after ${EMU_JID POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterok:${EMU_JID} \ - --kill-on-invalid-dep=yes \ + --dependency=afterany:${EMU_JID} \ --export=ALL,EXPERIMENT=l96_flux \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" +echo "=== Submitting exp_to_leaderboard (L96 flux-force, after ${EMU_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_flux \ + exp_to_leaderboard.sbatch) +echo " exp_to_leaderboard job ID: ${LB_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" #sbatch --parsable -A esm --job-name="post_diag_l96_flux" --export=ALL,EXPERIMENT=l96_flux posterior_diagnostic_plots_l96.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 8cd80dcff..5d5cc9390 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -50,12 +50,20 @@ echo "=== Submitting posterior_diagnostic_plots (L96 vec-force, after ${EMU_JID} POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterok:${EMU_JID} \ - --kill-on-invalid-dep=yes \ + --dependency=afterany:${EMU_JID} \ --export=ALL,EXPERIMENT=l96_vec \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" +echo "=== Submitting exp_to_leaderboard (L96 vec-force, after ${EMU_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_vec \ + exp_to_leaderboard.sbatch) +echo " exp_to_leaderboard job ID: ${LB_JID}" + echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch -A esm --job-name="post_diag_l96_vec" --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch From ec1d96bec07f4b975f7982b0d98726f787523f71 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 15:41:21 -0700 Subject: [PATCH 58/86] missed file --- .../hpc-variant/exp_to_leaderboard.sbatch | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch new file mode 100644 index 000000000..cbf5cc297 --- /dev/null +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch @@ -0,0 +1,34 @@ +#!/bin/bash +# Single job: converts emulate_sample posterior files to a leaderboard NetCDF. +# Runs all (N_ens, rng_idx) cases serially in one job (no array needed). +# +# Invoked automatically by submit_l*.sh after emulate_sample completes. +# Manual use: +# L96: sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_flux exp_to_leaderboard.sbatch +# L63: sbatch --dependency=afterok: exp_to_leaderboard.sbatch + +#SBATCH --job-name=leaderboard +#SBATCH --output=output/slurm/leaderboard_%j.out +#SBATCH --error=output/slurm/leaderboard_%j.err +#SBATCH --time=04:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +if [[ -z "${EXPERIMENT:-}" || "${EXPERIMENT}" == "l63" ]]; then + julia --project=. l63_exp_to_leaderboard_utilities.jl +else + julia --project=. l96_exp_to_leaderboard_utilities.jl +fi From dc540495785902e7cafe5ef455d9b64f6de300af Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 19:54:07 -0700 Subject: [PATCH 59/86] increase diagnostic samples --- .../l96_exp_to_leaderboard_utilities.jl | 6 ++++++ .../posterior_diagnostic_plots_l96.jl | 20 ++++++++++++++----- .../l96_exp_to_leaderboard_utilities.jl | 6 ++++++ .../EKIRace/posterior_diagnostic_plots_l96.jl | 18 +++++++++++++---- 4 files changed, 41 insertions(+), 9 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index b9ff6655c..c7f6fb054 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -128,6 +128,12 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(n_params)." + end + post_mean_arr[i, j, k, :] = pm post_cov_arr[i, j, k, :, :] = pc # (m - truth)' C^{-1} (m - truth) diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl index 1c11d690a..715ac7e85 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -26,7 +26,13 @@ include("experiment_config.jl") function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) m = vec(mean(samples, dims=2)) - C_raw = cov(samples') + C_raw = cov(samples, dims=2) + num_samples = size(samples, 2) + dim_samples = size(samples, 1) + r = rank(C_raw) + if r == num_samples - 1 && r < dim_samples - 1 + @warn "Covariance rank $(r) = num_samples-1 = $(num_samples-1) < dim-1 = $(dim_samples-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(dim_samples)." + end λ = max(1e-10, 1e-4 * mean(diag(C_raw))) C = Symmetric(C_raw + λ * I) dist = MvNormal(m, C) @@ -38,7 +44,7 @@ function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) end function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) - n_samples_pushforward = 100 + n_samples_pushforward = 1000 force_case = cfg.force_case calib_dir = calib_directory(method, cfg) @@ -137,10 +143,14 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) output_mah, output_lp = pushforward_metrics(G_ens, y) + n_frc = length(truth_frc_m) + lowbd_par, upbd_par = round.(quantile(Chisq(n_par), [0.01, 0.99]), digits=2) + lowbd_frc, upbd_frc = round.(quantile(Chisq(n_frc), [0.01, 0.99]), digits=2) + lowbd_out, upbd_out = round.(quantile(Chisq(ny), [0.01, 0.99]), digits=2) @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" - @info " param [d=$(n_par)]: mahal=$(round(param_mah, digits=2)) logpdf_ratio=$(round(param_lp, digits=2))" - @info " forcing [d=$(length(truth_frc_m))]: mahal=$(round(forcing_mah, digits=2)) logpdf_ratio=$(round(forcing_lp, digits=2))" - @info " output [d=$(ny)]: mahal=$(round(output_mah, digits=2)) logpdf_ratio=$(round(output_lp, digits=2))" + @info " param [d=$(n_par)]: target: [$(lowbd_par), $(upbd_par)] mahal=$(round(param_mah, digits=2)) -2*logpdf_ratio=$(round(-2*param_lp, digits=2))" + @info " forcing [d=$(n_frc)]: target: [$(lowbd_frc), $(upbd_frc)] mahal=$(round(forcing_mah, digits=2)) -2*logpdf_ratio=$(round(-2*forcing_lp, digits=2))" + @info " output [d=$(ny)]: target: [$(lowbd_out), $(upbd_out)] mahal=$(round(output_mah, digits=2)) -2*logpdf_ratio=$(round(-2*output_lp, digits=2))" gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 0ac83daf6..7d1ea78a7 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -126,6 +126,12 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(n_params)." + end + post_mean_arr[i, j, k, :] = pm post_cov_arr[i, j, k, :, :] = pc # (m - truth)' C^{-1} (m - truth) diff --git a/examples/EKIRace/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl index 6d236112c..6e8686165 100644 --- a/examples/EKIRace/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -22,7 +22,7 @@ include("experiment_config.jl") verbose_flag = false save_all_ekp = true -n_samples_pushforward = 100 # num. samples to pushforward through lorenz +n_samples_pushforward = 400 # num. samples to pushforward through lorenz ######################################################################## @@ -99,6 +99,12 @@ n_ens = length(N_enss) function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) m = vec(mean(samples, dims=2)) C_raw = cov(samples') + num_samples = size(samples, 2) + dim_samples = size(samples, 1) + r = rank(C_raw) + if r == num_samples - 1 && r < dim_samples - 1 + @warn "Covariance rank $(r) = num_samples-1 = $(num_samples-1) < dim-1 = $(dim_samples-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(dim_samples)." + end λ = max(1e-10, 1e-4 * mean(diag(C_raw))) C = Symmetric(C_raw + λ * I) dist = MvNormal(m, C) @@ -188,10 +194,14 @@ for (N_ens, rng_idx) in valid_file_items param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) output_mah, output_lp = pushforward_metrics(G_ens, y) + n_frc = length(truth_frc_m) + lowbd_par, upbd_par = round.(quantile(Chisq(n_par), [0.01, 0.99]), digits=2) + lowbd_frc, upbd_frc = round.(quantile(Chisq(n_frc), [0.01, 0.99]), digits=2) + lowbd_out, upbd_out = round.(quantile(Chisq(ny), [0.01, 0.99]), digits=2) @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" - @info " param [d=$(n_par)]: mahal=$(round(param_mah, digits=2)) logpdf_ratio=$(round(param_lp, digits=2))" - @info " forcing [d=$(length(truth_frc_m))]: mahal=$(round(forcing_mah, digits=2)) logpdf_ratio=$(round(forcing_lp, digits=2))" - @info " output [d=$(ny)]: mahal=$(round(output_mah, digits=2)) logpdf_ratio=$(round(output_lp, digits=2))" + @info " param [d=$(n_par)]: target: [$(lowbd_par), $(upbd_par)] mahal=$(round(param_mah, digits=2)) -2*logpdf_ratio=$(round(-2*param_lp, digits=2))" + @info " forcing [d=$(n_frc)]: target: [$(lowbd_frc), $(upbd_frc)] mahal=$(round(forcing_mah, digits=2)) -2*logpdf_ratio=$(round(-2*forcing_lp, digits=2))" + @info " output [d=$(ny)]: target: [$(lowbd_out), $(upbd_out)] mahal=$(round(output_mah, digits=2)) -2*logpdf_ratio=$(round(-2*output_lp, digits=2))" gr(size = (3 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) From 7b9e98cebc1475a11da90fa1721bad8989facddf Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 14:02:20 -0700 Subject: [PATCH 60/86] cpu=4 for plots --- .../EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index e152db7e3..6e0840f36 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -13,7 +13,7 @@ #SBATCH --error=output/slurm/post_diag_%j.err #SBATCH --time=00:30:00 #SBATCH --mem=16G -#SBATCH --cpus-per-task=1 +#SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 set -euo pipefail From f205007c80f0a21b4402cc046cefe5b91b8a6bb4 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 3 Jun 2026 19:55:06 -0700 Subject: [PATCH 61/86] add sbatch comment --- examples/EKIRace/hpc-variant/submit_l96_const.sh | 2 ++ examples/EKIRace/hpc-variant/submit_l96_flux.sh | 2 ++ examples/EKIRace/hpc-variant/submit_l96_vec.sh | 2 ++ 3 files changed, 6 insertions(+) diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 65121dce8..36f8fa455 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -67,3 +67,5 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch -A esm --job-name="post_diag_l96_const" --export=ALL,EXPERIMENT=l96_const posterior_diagnostic_plots_l96.sbatch + +# sbatch -A esm --job-name="leaderboard_l96_const" --export=ALL,EXPERIMENT=l96_const exp_to_leaderboard.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index da7caacfd..ed260875d 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -67,3 +67,5 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" #sbatch --parsable -A esm --job-name="post_diag_l96_flux" --export=ALL,EXPERIMENT=l96_flux posterior_diagnostic_plots_l96.sbatch + +# sbatch -A esm --job-name="leaderboard_l96_flux" --export=ALL,EXPERIMENT=l96_flux exp_to_leaderboard.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 5d5cc9390..65bc1ae95 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -67,3 +67,5 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch -A esm --job-name="post_diag_l96_vec" --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch + +# sbatch -A esm --job-name="leaderboard_l96_vec" --export=ALL,EXPERIMENT=l96_vec exp_to_leaderboard.sbatch From ccb2a2e0a070633a5984fe14049fb93041d69d42 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 4 Jun 2026 11:57:34 -0700 Subject: [PATCH 62/86] add array to post_diag sbatch --- .../posterior_diagnostic_plots_l96.sbatch | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index 6e0840f36..964696335 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -1,7 +1,7 @@ #!/bin/bash -# Single job that pushes all posteriors forward through the Lorenz96 forward map, -# saves pushforward plots, and produces posterior-ribbon diagnostic figures. -# Runs all (N_ens, rng_idx) cases serially in one job (no array needed). +# Array job: one task per (N_ens, rng_idx) pair (60 tasks for each L96 case). +# Each task pushes its posterior forward through the Lorenz96 forward map and +# saves pushforward + posterior-ribbon diagnostic figures. # # Invoked automatically by submit_l96_*.sh after emulate_sample completes. # Manual use: @@ -9,8 +9,9 @@ # --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch #SBATCH --job-name=post_diag -#SBATCH --output=output/slurm/post_diag_%j.out -#SBATCH --error=output/slurm/post_diag_%j.err +#SBATCH --output=output/slurm/post_diag_%A_%a.out +#SBATCH --error=output/slurm/post_diag_%A_%a.err +#SBATCH --array=1-60%20 #SBATCH --time=00:30:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 @@ -27,4 +28,4 @@ cd "${SLURM_SUBMIT_DIR}" export JULIA_PKG_PRECOMPILE_AUTO=0 -julia --project=. posterior_diagnostic_plots_l96.jl +julia --project=. posterior_diagnostic_plots_l96.jl "${SLURM_ARRAY_TASK_ID}" From dd0cb1a18687a51a5f32b8a0b1296f51cf58c8b7 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 4 Jun 2026 12:36:17 -0700 Subject: [PATCH 63/86] update readme --- examples/EKIRace/README.md | 63 +++++++---- examples/EKIRace/hpc-variant/README.md | 146 ++++++++++++++----------- 2 files changed, 120 insertions(+), 89 deletions(-) diff --git a/examples/EKIRace/README.md b/examples/EKIRace/README.md index 3de352e49..d7532b33d 100644 --- a/examples/EKIRace/README.md +++ b/examples/EKIRace/README.md @@ -15,16 +15,15 @@ scripts for the Lorenz 63 and Lorenz 96 benchmark experiments. ## One-time setup In `experiment_config.jl`, set `EXPERIMENT` to the case you want to run and pin -`calibrate_date` to today's date before starting a run, e.g. +`calibrate_date` before starting a run, e.g. ```julia EXPERIMENT = :l96_vec -calibrate_date = Date("2026-05-31", "yyyy-mm-dd") +calibrate_date = Date("2026-06-04", "yyyy-mm-dd") ``` -All subsequent stages (emulate_sample, ensemble_from_posterior, leaderboard -utilities) include `experiment_config.jl` and read these values, so keeping them -fixed ensures every stage finds the right output directory. +All subsequent stages include `experiment_config.jl` and read these values, so +keeping them fixed ensures every stage finds the right output directory. ## Running scripts @@ -52,7 +51,17 @@ julia --project -e 'include("calibrate_l96.jl")' The L96 calibration also writes a `l96_computed_preliminaries_.jld2` file to `output/` that the later stages require. -### 2. Emulate and sample +### 2. Calibration diagnostic plots (L96 only) + +```bash +julia --project -e 'include("calibration_diagnostic_plots_l96.jl")' +``` + +Reads calibration output and produces data, solution (mean ensemble), and +full-ensemble spread figures for each `(N_ens, rng_idx)` cell. Can run +in parallel with stage 3. + +### 3. Emulate and sample ```bash # L63 @@ -63,30 +72,34 @@ julia --project -e 'include("emulate_sample_l96.jl")' ``` Reads the calibration output for each `(N_ens, rng_idx)` cell, trains an -emulator, runs MCMC, and writes per-cell posterior `.jld2` files plus -summary netCDF files to the calibration output directory. +emulator, runs MCMC, and writes per-cell posterior `.jld2` files to the +calibration output directory. -### 3. Posterior pushforward (L96 only) +### 4. Posterior diagnostic plots (L96 only) ```bash -julia --project -e 'include("ensemble_from_posterior.jl")' +julia --project -e 'include("posterior_diagnostic_plots_l96.jl")' ``` -Requires `emulate_sample_l96.jl` to have completed. Loops over all valid -`(N_ens, rng_idx)` cells, samples 100 points from each posterior, runs them -through the Lorenz 96 forward map, and writes -`pushforward_from_posterior_*_k_full_ens.{png,pdf}` plots into the -calibration output directory. +Reads each posterior `.jld2`, pushes 400 samples through the Lorenz 96 +forward map, and writes pushforward and posterior-ribbon diagnostic plots +into the calibration output directory. -### 4. Leaderboard +### 5. Leaderboard -The leaderboard utility scripts (`l63_exp_to_leaderboard_utilities.jl`, -`l96_exp_to_leaderboard_utilities.jl`) are not standalone — they are included -by downstream analysis. +Convert the per-cell posterior files into a leaderboard NetCDF file: + +```bash +# L63 +julia --project -e 'include("l63_exp_to_leaderboard_utilities.jl")' + +# L96 +julia --project -e 'include("l96_exp_to_leaderboard_utilities.jl")' +``` -To compute summary metrics from a netCDF result file, edit the `filename` -variable at the top of `compute_leaderboard_metrics.jl` to point to the -target file, then run: +To compute summary metrics from the resulting netCDF file, edit the `filename` +variable at the top of `compute_leaderboard_metrics.jl` to point to the target +file, then run: ```bash julia --project -e 'include("compute_leaderboard_metrics.jl")' @@ -100,13 +113,15 @@ chi-squared reference quantiles. ```julia # In experiment_config.jl: EXPERIMENT = :l96_vec -calibrate_date = Date("2026-05-31", "yyyy-mm-dd") +calibrate_date = Date("2026-06-04", "yyyy-mm-dd") ``` ```bash julia --project -e 'include("calibrate_l96.jl")' +julia --project -e 'include("calibration_diagnostic_plots_l96.jl")' # optional julia --project -e 'include("emulate_sample_l96.jl")' -julia --project -e 'include("ensemble_from_posterior.jl")' +julia --project -e 'include("posterior_diagnostic_plots_l96.jl")' # optional +julia --project -e 'include("l96_exp_to_leaderboard_utilities.jl")' ``` ## HPC variant diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index b19821cfe..531f83dc6 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -6,58 +6,74 @@ SLURM job array where every `(N_ens, rng_idx)` cell is an independent task. ## One-time setup -In `experiment_config.jl`, pin the calibrate date to today's date before -starting a run, e.g. +In `experiment_config.jl`, pin the calibrate date before starting a run, e.g. ```julia -calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +calibrate_date = Date("2026-06-04", "yyyy-mm-dd") ``` -All subsequent stages (emulate_sample, leaderboard) read this value to locate -the right output directory; keeping it fixed avoids mismatches when jobs run -past midnight or across days. +All subsequent stages read this value to locate the right output directory; +keeping it fixed avoids mismatches when jobs run past midnight or across days. Precompilation is handled by a dedicated `submit_precompile.sh` script that queues `precompile.sbatch` as a compute job. Run it once before your experiments and again whenever the environment changes (fresh checkout, package updates). -The `submit_l*.sh` scripts do not precompile — they will remind you at submission -time. +The `submit_l*.sh` scripts do not precompile — they will remind you at +submission time. + +## Pipeline + +### L63 + +``` +calibrate_array ──afterok──► emulate_sample_array ──afterany──► exp_to_leaderboard +``` + +### L96 (const / vec / flux) + +``` + ┌──afterok──► calibration_diagnostic_plots_l96 +calibrate_array ──afterok──► emulate_sample_array ──afterany──► posterior_diagnostic_plots_l96 + ──afterany──► exp_to_leaderboard +``` + +`calibration_diagnostic_plots` and `emulate_sample` both start once calibrate +succeeds (they run in parallel). `posterior_diagnostic_plots` and +`exp_to_leaderboard` both start once `emulate_sample` finishes (whether or not +it succeeded, so partial results are still processed). ## Standalone (serial) -Run with no arguments to sweep all `(N_ens, rng_idx)` cells sequentially, -exactly as the original scripts did. +Run with no arguments to sweep all `(N_ens, rng_idx)` cells sequentially. ```bash # L63 julia --project=. calibrate_l63.jl julia --project=. emulate_sample_l63.jl -# L96 — set EXPERIMENT or edit the toggle in experiment_config.jl +# L96 — set EXPERIMENT env var or edit the toggle in experiment_config.jl EXPERIMENT=l96_const julia --project=. calibrate_l96.jl +EXPERIMENT=l96_const julia --project=. calibration_diagnostic_plots_l96.jl EXPERIMENT=l96_const julia --project=. emulate_sample_l96.jl +EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl ``` -You can also run a single cell by passing its 1-based task index: +You can also run a single cell by passing its 1-based task index for the +array-capable scripts: ```bash -julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only -julia --project=. emulate_sample_l63.jl 5 # fifth cell only +julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only +julia --project=. emulate_sample_l63.jl 5 # fifth cell only +EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl 3 ``` -After `emulate_sample` completes, push the posteriors forward through the L96 -forward map and save plots: +Leaderboard conversion runs all cells serially in a single call: ```bash -EXPERIMENT=l96_const julia --project=. ensemble_from_posterior.jl -EXPERIMENT=l96_vec julia --project=. ensemble_from_posterior.jl -EXPERIMENT=l96_flux julia --project=. ensemble_from_posterior.jl +julia --project=. l63_exp_to_leaderboard_utilities.jl +EXPERIMENT=l96_const julia --project=. l96_exp_to_leaderboard_utilities.jl ``` -Each call loops over all `(N_ens, rng_idx)` cells, samples 100 points from each -posterior, runs them through the forward model, and writes `pushforward_from_posterior_*.{png,pdf}` -plots into the per-case calibration output directory. - ## HPC (Caltech Resnick cluster, SLURM) ### Submission scripts (recommended) @@ -73,10 +89,10 @@ bash submit_l96_vec.sh [EXP_ID] bash submit_l96_flux.sh [EXP_ID] ``` -Each `submit_l*.sh` script chains `calibrate_array.sbatch` → `emulate_sample_array.sbatch` -for its case. All four cases can be launched simultaneously — output files are -case-specific so there are no write conflicts. The optional `EXP_ID` argument -suffixes SLURM job names to keep the queue readable. +Each `submit_l*.sh` script chains the full pipeline for its case automatically. +All four cases can be launched simultaneously — output files are case-specific +so there are no write conflicts. The optional `EXP_ID` argument suffixes SLURM +job names to keep the queue readable. If you are launching all four cases together you only need one precompile run: @@ -90,59 +106,59 @@ wait ### Manual submission -Precompile via `submit_precompile.sh` (or directly), then submit calibrate and -emulate_sample: +Precompile via `submit_precompile.sh` (or directly), then submit each stage: ```bash # L63 -cid=$(sbatch --parsable -A esm \ - --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) +CALIB_JID=$(sbatch --parsable -A esm \ + --export=ALL,SCRIPT=calibrate_l63.jl calibrate_array.sbatch) +EMU_JID=$(sbatch --parsable -A esm \ + --dependency=afterok:${CALIB_JID} --kill-on-invalid-dep=yes \ + --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch) sbatch -A esm \ - --dependency=afterok:$cid --kill-on-invalid-dep=yes \ - --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch + --dependency=afterany:${EMU_JID} \ + exp_to_leaderboard.sbatch # L96 -cid=$(sbatch --parsable -A esm \ - --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const calibrate_array.sbatch) +CALIB_JID=$(sbatch --parsable -A esm \ + --export=ALL,SCRIPT=calibrate_l96.jl,EXPERIMENT=l96_const \ + calibrate_array.sbatch) +sbatch -A esm \ + --dependency=afterok:${CALIB_JID} --kill-on-invalid-dep=yes \ + --export=ALL,EXPERIMENT=l96_const \ + calibration_diagnostic_plots_l96.sbatch +EMU_JID=$(sbatch --parsable -A esm \ + --dependency=afterok:${CALIB_JID} --kill-on-invalid-dep=yes \ + --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ + emulate_sample_array.sbatch) +sbatch -A esm \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + posterior_diagnostic_plots_l96.sbatch sbatch -A esm \ - --dependency=afterok:$cid --kill-on-invalid-dep=yes \ - --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ - emulate_sample_array.sbatch + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + exp_to_leaderboard.sbatch ``` -Each task writes its per-method `ekp` and `results` files, which are consumed -directly by the emulate_sample stage. - -### Posterior pushforward (`ensemble_from_posterior`) - -After the `emulate_sample` jobs finish, submit a single (non-array) job that -loops over all `(N_ens, rng_idx)` cells internally: - -```bash -eid=$(sbatch --parsable \ - --dependency=afterok: --kill-on-invalid-dep=yes \ - --export=ALL,EXPERIMENT=l96_const \ - ensemble_from_posterior.sbatch) - -sbatch --dependency=afterok: --kill-on-invalid-dep=yes \ - --export=ALL,EXPERIMENT=l96_vec ensemble_from_posterior.sbatch - -sbatch --dependency=afterok: --kill-on-invalid-dep=yes \ - --export=ALL,EXPERIMENT=l96_flux ensemble_from_posterior.sbatch -``` +### Sbatch files reference -Replace `` with the job ID returned by the corresponding -`emulate_sample_array.sbatch` submission. Each run takes up to 8 h and writes -`pushforward_from_posterior_*.{png,pdf}` into the same output directory as the -calibration results. +| File | Type | Description | +|------|------|-------------| +| `calibrate_array.sbatch` | array (1–60) | One task per `(N_ens, rng_idx)` cell | +| `emulate_sample_array.sbatch` | array (1–60) | One task per `(N_ens, rng_idx)` cell | +| `calibration_diagnostic_plots_l96.sbatch` | single job | Calibration figures, all cells serially (L96) | +| `posterior_diagnostic_plots_l96.sbatch` | array (1–60) | Pushforward figures, one task per cell (L96) | +| `exp_to_leaderboard.sbatch` | single job | NetCDF leaderboard file, all cells serially | +| `precompile.sbatch` | single job | `Pkg.instantiate()` + `Pkg.precompile()` | ### Adjusting array size The sbatch files default to `--array=1-60` (3 ensemble sizes × 20 repeats). If you change `N_ens_sizes` or `n_repeats` in `experiment_config.jl`, update -the upper bound to `length(N_ens_sizes) * n_repeats`. The `%20` suffix caps -concurrent tasks to 20 at a time as a cluster-courtesy limit; raise or remove -it if you want faster turnaround. +the upper bound to `length(N_ens_sizes) * n_repeats`. The `%20` suffix on +`posterior_diagnostic_plots_l96.sbatch` caps concurrent tasks to 20 at a time +as a cluster-courtesy limit; raise or remove it if you want faster turnaround. ### Smoke test From 389f40307c2dcc95cb546624a02294eba0c2fad5 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 4 Jun 2026 18:25:52 -0700 Subject: [PATCH 64/86] add metric computation into the nc files --- .../l96_exp_to_leaderboard_utilities.jl | 69 ++++++++++++++++++- .../l96_exp_to_leaderboard_utilities.jl | 69 ++++++++++++++++++- 2 files changed, 132 insertions(+), 6 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index c7f6fb054..76cd0302d 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -96,8 +96,12 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) -forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) -output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -106,7 +110,8 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] - truth_params = loaded["truth_params"] + truth_params = loaded["truth_params"] + truth_params_constrained = loaded["truth_params_constrained"] # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -118,6 +123,10 @@ for (N_ens, rng_idx) in valid_file_items true_param_arr[i, j, :] = truth_params n_evals_arr[i, j] = conv_alg_iters * N_ens + emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) + truth_forcing_vec = forcing(emc_truth, x0) + truth_output_vec = lorenz_forward(emc_truth, x0, lorenz_config_settings, observation_config) + for k in k_values post_dist = posteriors_by_k[k] pm = vec(mean(post_dist)) @@ -155,6 +164,28 @@ for (N_ens, rng_idx) in valid_file_items observation_config, ) end + + # forcing-space metrics (empirical distribution of pushforward samples vs truth forcing) + fs = forcing_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_forcing) + fm = vec(mean(fs, dims=1)) + fc = Symmetric(cov(fs) + 1e-10 * I) + f_normal = MvNormal(fm, fc) + f_cols = Matrix(fs') + f_mode = f_cols[:, argmax(logpdf(f_normal, f_cols))] + f_diff = fm - truth_forcing_vec + forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) + forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + + # output-space metrics (empirical distribution of pushforward samples vs truth output) + os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - truth_output_vec + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, truth_output_vec) - logpdf(o_normal, o_mode) end end @@ -250,5 +281,37 @@ output_samples_v = defVar( output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz96 output space (mean/std per spatial point)" output_samples_v[:, :, :, :, :] = output_samples_arr +forcing_mahal_v = defVar( + ds, "forcing_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +forcing_mahal_v.attrib["description"] = "Mahalanobis distance in forcing space: (fm - truth_forcing)' Cf^{-1} (fm - truth_forcing), where fm and Cf are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward forcing samples" +forcing_mahal_v[:, :, :] = forcing_mahal_arr + +forcing_logpdf_true_v_map_v = defVar( + ds, "forcing_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +forcing_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in forcing space: logpdf(N(fm,Cf), truth_forcing) - logpdf(N(fm,Cf), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward forcing distribution" +forcing_logpdf_true_v_map_v[:, :, :] = forcing_logpdf_true_v_map_arr + +output_mahal_v = defVar( + ds, "output_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - truth_output)' Co^{-1} (om - truth_output), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples" +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar( + ds, "output_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), truth_output) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution" +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 7d1ea78a7..95a88ab46 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -94,8 +94,12 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) -forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) -output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -104,7 +108,8 @@ for (N_ens, rng_idx) in valid_file_items posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] - truth_params = loaded["truth_params"] + truth_params = loaded["truth_params"] + truth_params_constrained = loaded["truth_params_constrained"] # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -116,6 +121,10 @@ for (N_ens, rng_idx) in valid_file_items true_param_arr[i, j, :] = truth_params n_evals_arr[i, j] = conv_alg_iters * N_ens + emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) + truth_forcing_vec = forcing(emc_truth, x0) + truth_output_vec = lorenz_forward(emc_truth, x0, lorenz_config_settings, observation_config) + for k in k_values post_dist = posteriors_by_k[k] pm = vec(mean(post_dist)) @@ -153,6 +162,28 @@ for (N_ens, rng_idx) in valid_file_items observation_config, ) end + + # forcing-space metrics (empirical distribution of pushforward samples vs truth forcing) + fs = forcing_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_forcing) + fm = vec(mean(fs, dims=1)) + fc = Symmetric(cov(fs) + 1e-10 * I) + f_normal = MvNormal(fm, fc) + f_cols = Matrix(fs') + f_mode = f_cols[:, argmax(logpdf(f_normal, f_cols))] + f_diff = fm - truth_forcing_vec + forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) + forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + + # output-space metrics (empirical distribution of pushforward samples vs truth output) + os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - truth_output_vec + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, truth_output_vec) - logpdf(o_normal, o_mode) end end @@ -248,5 +279,37 @@ output_samples_v = defVar( output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz96 output space (mean/std per spatial point)" output_samples_v[:, :, :, :, :] = output_samples_arr +forcing_mahal_v = defVar( + ds, "forcing_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +forcing_mahal_v.attrib["description"] = "Mahalanobis distance in forcing space: (fm - truth_forcing)' Cf^{-1} (fm - truth_forcing), where fm and Cf are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward forcing samples" +forcing_mahal_v[:, :, :] = forcing_mahal_arr + +forcing_logpdf_true_v_map_v = defVar( + ds, "forcing_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +forcing_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in forcing space: logpdf(N(fm,Cf), truth_forcing) - logpdf(N(fm,Cf), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward forcing distribution" +forcing_logpdf_true_v_map_v[:, :, :] = forcing_logpdf_true_v_map_arr + +output_mahal_v = defVar( + ds, "output_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - truth_output)' Co^{-1} (om - truth_output), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples" +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar( + ds, "output_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), truth_output) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution" +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" From 2f76ddbed712483d5e61bf13469eba60cff45062 Mon Sep 17 00:00:00 2001 From: odunbar Date: Sat, 6 Jun 2026 13:55:27 -0700 Subject: [PATCH 65/86] added low-rank + residual mahalanobis, to get more robust metrics --- .../EKIRace/compute_leaderboard_metrics.jl | 272 ++++++++++++++++- .../compute_leaderboard_metrics.jl | 280 +++++++++++++++++- .../l96_exp_to_leaderboard_utilities.jl | 74 ++++- .../posterior_diagnostic_plots_l96.jl | 3 +- .../l96_exp_to_leaderboard_utilities.jl | 74 ++++- .../EKIRace/posterior_diagnostic_plots_l96.jl | 3 +- 6 files changed, 664 insertions(+), 42 deletions(-) diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl index 61de75089..18f66b190 100644 --- a/examples/EKIRace/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -4,15 +4,22 @@ using Statistics using Printf using DataFrames -indir = joinpath("hpc-variant","output","from-hpc_2026-05-29") +indir = "." #joinpath("output","from-hpc_2026-05-31") #filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" #filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" #filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" -filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-29.nc" +filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-31.nc" filename = joinpath(indir, filename) @info "computing leaderboard metrics from $(filename)" +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis + ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] @@ -22,27 +29,26 @@ lp = ncd[:posterior_logpdf_true_v_map] mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] -qq_probs = [0.1, 0.5, 0.9] -qq_good = quantile(Chisq(n_par), qq_probs) +qq_good = quantile(Chisq(n_par), calibration_check_quantiles) # === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === # If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. mh_pool = collect(skipmissing(vec(Array(mh)))) lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) -mh_emp_q = quantile(mh_pool, qq_probs) -lp_emp_q = quantile(lp_pool, qq_probs) +mh_emp_q = quantile(mh_pool, calibration_check_quantiles) +lp_emp_q = quantile(lp_pool, calibration_check_quantiles) println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) -for (i, p) in enumerate(qq_probs) +for (i, p) in enumerate(calibration_check_quantiles) println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) end # === 2. Per-experiment scores === # missing where no valid seeds; * in printout where coverage < n_rng (less trusted) -mh_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) -lp_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) +mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) +lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good) for j in axes(mh, 3), i in axes(mh, 2) @@ -74,7 +80,7 @@ df_scores = DataFrame( lp_coverage = Float64[], ) -for (qi, p) in enumerate(qq_probs) +for (qi, p) in enumerate(calibration_check_quantiles) for (ei, ev) in enumerate(ens_vals) for (ki, kv) in enumerate(kiter_vals) push!(df_scores, ( @@ -92,13 +98,255 @@ end # rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") +println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") for ens in sort(unique(df_scores.ens_size)) println("\n=== ens_size = $(ens) ===") - for (qi, p) in enumerate(qq_probs) + for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) - println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2))]") + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) println() end end +# === 4. Forcing-space M and L scores === +if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") + n_for = ncd.dim["forcing_dim"] + fmh = ncd[:forcing_mahalanobis] + flp = ncd[:forcing_logpdf_true_v_map] + qq_good_f = quantile(Chisq(n_for), calibration_check_quantiles) + + fmh_nonmiss = sum(.!ismissing.(fmh), dims=1)[1,:,:] + flp_nonmiss = sum(.!ismissing.(flp), dims=1)[1,:,:] + fmh_pool = collect(skipmissing(vec(Array(fmh)))) + flp_pool = collect(skipmissing(vec(-2 .* Array(flp)))) + fmh_emp_q = quantile(fmh_pool, calibration_check_quantiles) + flp_emp_q = quantile(flp_pool, calibration_check_quantiles) + + println("\n" * "="^72) + println("Forcing-space M and L (chi-sq ref dim = $(n_for))") + println("="^72) + println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_for)):") + println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_f", "-2*lp_f")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_f[i], fmh_emp_q[i], flp_emp_q[i])) + end + + fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good_f) + for j in axes(fmh, 3), i in axes(fmh, 2) + nm = fmh_nonmiss[i, j]; nm == 0 && continue + fmh_scores[idx, i, j] = sum(skipmissing(fmh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(flp, 3), i in axes(flp, 2) + nm = flp_nonmiss[i, j]; nm == 0 && continue + flp_scores[idx, i, j] = sum(skipmissing(-2 .* flp[:, i, j]) .<= qg; init=0) / nm + end + end + + fmh_coverage = fmh_nonmiss ./ n_rng + flp_coverage = flp_nonmiss ./ n_rng + df_fscores = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_coverage=Float64[], lp_coverage=Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_fscores, (q_prob=p, ens_size=ev, k_iter=kv, + mh_score=fmh_scores[qi,ei,ki], lp_score=flp_scores[qi,ei,ki], + mh_coverage=fmh_coverage[ei,ki], lp_coverage=flp_coverage[ei,ki])) + end + + println("\nForcing-space scores by ens_size and q_prob:") + println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") + for ens in sort(unique(df_fscores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_fscores), :k_iter) + println(" q = $(p) [chi-sq($(n_for)) bound = $(round(qq_good_f[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end + end +end + +# === 5. Output-space M and L scores === +if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") + n_out = ncd.dim["output_dim"] + omh = ncd[:output_mahalanobis] + olp = ncd[:output_logpdf_true_v_map] + qq_good_o = quantile(Chisq(n_out), calibration_check_quantiles) + + omh_nonmiss = sum(.!ismissing.(omh), dims=1)[1,:,:] + olp_nonmiss = sum(.!ismissing.(olp), dims=1)[1,:,:] + omh_pool = collect(skipmissing(vec(Array(omh)))) + olp_pool = collect(skipmissing(vec(-2 .* Array(olp)))) + omh_emp_q = quantile(omh_pool, calibration_check_quantiles) + olp_emp_q = quantile(olp_pool, calibration_check_quantiles) + + println("\n" * "="^72) + println("Output-space M and L (chi-sq ref dim = $(n_out))") + println("="^72) + println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_out)):") + println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_o", "-2*lp_o")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_o[i], omh_emp_q[i], olp_emp_q[i])) + end + + omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good_o) + for j in axes(omh, 3), i in axes(omh, 2) + nm = omh_nonmiss[i, j]; nm == 0 && continue + omh_scores[idx, i, j] = sum(skipmissing(omh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(olp, 3), i in axes(olp, 2) + nm = olp_nonmiss[i, j]; nm == 0 && continue + olp_scores[idx, i, j] = sum(skipmissing(-2 .* olp[:, i, j]) .<= qg; init=0) / nm + end + end + + omh_coverage = omh_nonmiss ./ n_rng + olp_coverage = olp_nonmiss ./ n_rng + df_oscores = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_coverage=Float64[], lp_coverage=Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_oscores, (q_prob=p, ens_size=ev, k_iter=kv, + mh_score=omh_scores[qi,ei,ki], lp_score=olp_scores[qi,ei,ki], + mh_coverage=omh_coverage[ei,ki], lp_coverage=olp_coverage[ei,ki])) + end + + println("\nOutput-space scores by ens_size and q_prob:") + println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") + for ens in sort(unique(df_oscores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_oscores), :k_iter) + println(" q = $(p) [chi-sq($(n_out)) bound = $(round(qq_good_o[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end + end +end + +# === 6. Perturbed-low-rank Mahalanobis (C ≈ aI + U*D*U', k = n_lowrank_modes top modes) === +# Each space yields three terms with separate reference distributions: +# top ~ Chisq(k) — miscalibration in the k principal directions +# residual ~ Chisq(n-k) — miscalibration in the n-k noise-floor directions (a = mean of bottom n-k eigenvalues) +# total ~ Chisq(n) — full-space calibration (top + residual, identical to Woodbury inverse applied to full diff) +for (space_label, top_sym, res_sym, n_dim_key) in [ + ("Forcing", :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), + ("Output", :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), +] + (haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue + n_dim = ncd.dim[n_dim_key] + top_raw = Array(ncd[top_sym]) # (n_rng, n_ens_size, n_k_iter), Union{Float64,Missing} + res_raw = Array(ncd[res_sym]) + tot_raw = top_raw .+ res_raw # missing propagates where either component is missing + + println("\n" * "="^72) + println("$(space_label)-space perturbed-low-rank M (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") + println(" Covariance model: C ≈ a*I + U*D*U' (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of empirical C)") + println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") + println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") + println(" total ~ Chisq($(n_dim)) — full-space calibration (top + residual)") + println("="^72) + + for (part_label, arr, ref_dof) in [ + ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), + ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), + ("total (n=$(n_dim))", tot_raw, n_dim), + ] + nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] + pool = collect(skipmissing(vec(arr))) + isempty(pool) && continue + qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) + emp_q = quantile(pool, calibration_check_quantiles) + coverage = nonmiss ./ n_rng + + println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") + println(" Calibration check — pooled empirical quantiles vs chi-sq($(ref_dof)):") + println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) + end + + scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_ref) + for j in axes(arr, 3), i in axes(arr, 2) + nm = nonmiss[i, j]; nm == 0 && continue + scores[idx, i, j] = sum(skipmissing(arr[:, i, j]) .<= qg; init=0) / nm + end + end + + df_plr = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + score = Union{Float64,Missing}[], + coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, + score=scores[qi,ei,ki], coverage=coverage[ei,ki])) + end + + println("\n Scores by ens_size and q_prob (expected score ≈ q_prob):") + for ens in sort(unique(df_plr.ens_size)) + println("\n === ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) + println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :score, :coverage), allrows=true) + println() + end + end + end +end + +# === 7. Marginal coverage fractions (forcing and output) === +for (space_label, cov_sym) in [("Forcing", :forcing_coverage), ("Output", :output_coverage)] + haskey(ncd, cov_sym) || continue + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs = ncd["coverage_quantile"][:] + + println("\n" * "="^72) + println("$(space_label)-space marginal coverage fractions") + println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("="^72) + println("Pooled mean coverage across all experiments:") + for (qi, qp) in enumerate(cov_qs) + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) + println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + df_cov = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mean_coverage = Union{Float64,Missing}[], + n_valid = Int[], + ) + for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = collect(skipmissing(cov[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + n_valid = length(vals))) + end + + println("\n$(space_label)-space coverage by ens_size and q_prob:") + println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + for ens in sort(unique(df_cov.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(cov_qs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :n_valid), allrows=true) + println() + end + end +end diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index 61de75089..61a6209da 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -3,16 +3,25 @@ using Distributions using Statistics using Printf using DataFrames +using Dates -indir = joinpath("hpc-variant","output","from-hpc_2026-05-29") -#filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" -#filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" -#filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" -filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-29.nc" +calib_date = Date("2026-06-03", "yyyy-mm-dd") +indir = joinpath("output","from-hpc_$(calib_date)") +#filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" +#filename = "ces-eki-dmc_l96_ensemble_results_$(calib_date).nc" +filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_$(calib_date).nc" +#filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_$(calib_date).nc" filename = joinpath(indir, filename) @info "computing leaderboard metrics from $(filename)" +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis + ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] @@ -22,27 +31,26 @@ lp = ncd[:posterior_logpdf_true_v_map] mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] -qq_probs = [0.1, 0.5, 0.9] -qq_good = quantile(Chisq(n_par), qq_probs) +qq_good = quantile(Chisq(n_par), calibration_check_quantiles) # === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === # If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. mh_pool = collect(skipmissing(vec(Array(mh)))) lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) -mh_emp_q = quantile(mh_pool, qq_probs) -lp_emp_q = quantile(lp_pool, qq_probs) +mh_emp_q = quantile(mh_pool, calibration_check_quantiles) +lp_emp_q = quantile(lp_pool, calibration_check_quantiles) println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) -for (i, p) in enumerate(qq_probs) +for (i, p) in enumerate(calibration_check_quantiles) println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) end # === 2. Per-experiment scores === # missing where no valid seeds; * in printout where coverage < n_rng (less trusted) -mh_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) -lp_scores = Array{Union{Float64,Missing}}(missing, length(qq_probs), n_ens_size, n_k_iter) +mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) +lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good) for j in axes(mh, 3), i in axes(mh, 2) @@ -74,7 +82,7 @@ df_scores = DataFrame( lp_coverage = Float64[], ) -for (qi, p) in enumerate(qq_probs) +for (qi, p) in enumerate(calibration_check_quantiles) for (ei, ev) in enumerate(ens_vals) for (ki, kv) in enumerate(kiter_vals) push!(df_scores, ( @@ -92,13 +100,255 @@ end # rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") +println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") for ens in sort(unique(df_scores.ens_size)) println("\n=== ens_size = $(ens) ===") - for (qi, p) in enumerate(qq_probs) + for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) - println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2))]") + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) println() end end +# === 4. Forcing-space M and L scores === +if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") + n_for = ncd.dim["forcing_dim"] + fmh = ncd[:forcing_mahalanobis] + flp = ncd[:forcing_logpdf_true_v_map] + qq_good_f = quantile(Chisq(n_for), calibration_check_quantiles) + + fmh_nonmiss = sum(.!ismissing.(fmh), dims=1)[1,:,:] + flp_nonmiss = sum(.!ismissing.(flp), dims=1)[1,:,:] + fmh_pool = collect(skipmissing(vec(Array(fmh)))) + flp_pool = collect(skipmissing(vec(-2 .* Array(flp)))) + fmh_emp_q = quantile(fmh_pool, calibration_check_quantiles) + flp_emp_q = quantile(flp_pool, calibration_check_quantiles) + + println("\n" * "="^72) + println("Forcing-space M and L (chi-sq ref dim = $(n_for))") + println("="^72) + println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_for)):") + println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_f", "-2*lp_f")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_f[i], fmh_emp_q[i], flp_emp_q[i])) + end + + fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good_f) + for j in axes(fmh, 3), i in axes(fmh, 2) + nm = fmh_nonmiss[i, j]; nm == 0 && continue + fmh_scores[idx, i, j] = sum(skipmissing(fmh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(flp, 3), i in axes(flp, 2) + nm = flp_nonmiss[i, j]; nm == 0 && continue + flp_scores[idx, i, j] = sum(skipmissing(-2 .* flp[:, i, j]) .<= qg; init=0) / nm + end + end + + fmh_coverage = fmh_nonmiss ./ n_rng + flp_coverage = flp_nonmiss ./ n_rng + df_fscores = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_coverage=Float64[], lp_coverage=Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_fscores, (q_prob=p, ens_size=ev, k_iter=kv, + mh_score=fmh_scores[qi,ei,ki], lp_score=flp_scores[qi,ei,ki], + mh_coverage=fmh_coverage[ei,ki], lp_coverage=flp_coverage[ei,ki])) + end + + println("\nForcing-space scores by ens_size and q_prob:") + println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") + for ens in sort(unique(df_fscores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_fscores), :k_iter) + println(" q = $(p) [chi-sq($(n_for)) bound = $(round(qq_good_f[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end + end +end + +# === 5. Output-space M and L scores === +if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") + n_out = ncd.dim["output_dim"] + omh = ncd[:output_mahalanobis] + olp = ncd[:output_logpdf_true_v_map] + qq_good_o = quantile(Chisq(n_out), calibration_check_quantiles) + + omh_nonmiss = sum(.!ismissing.(omh), dims=1)[1,:,:] + olp_nonmiss = sum(.!ismissing.(olp), dims=1)[1,:,:] + omh_pool = collect(skipmissing(vec(Array(omh)))) + olp_pool = collect(skipmissing(vec(-2 .* Array(olp)))) + omh_emp_q = quantile(omh_pool, calibration_check_quantiles) + olp_emp_q = quantile(olp_pool, calibration_check_quantiles) + + println("\n" * "="^72) + println("Output-space M and L (chi-sq ref dim = $(n_out))") + println("="^72) + println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_out)):") + println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_o", "-2*lp_o")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_o[i], omh_emp_q[i], olp_emp_q[i])) + end + + omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good_o) + for j in axes(omh, 3), i in axes(omh, 2) + nm = omh_nonmiss[i, j]; nm == 0 && continue + omh_scores[idx, i, j] = sum(skipmissing(omh[:, i, j]) .<= qg; init=0) / nm + end + for j in axes(olp, 3), i in axes(olp, 2) + nm = olp_nonmiss[i, j]; nm == 0 && continue + olp_scores[idx, i, j] = sum(skipmissing(-2 .* olp[:, i, j]) .<= qg; init=0) / nm + end + end + + omh_coverage = omh_nonmiss ./ n_rng + olp_coverage = olp_nonmiss ./ n_rng + df_oscores = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_coverage=Float64[], lp_coverage=Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_oscores, (q_prob=p, ens_size=ev, k_iter=kv, + mh_score=omh_scores[qi,ei,ki], lp_score=olp_scores[qi,ei,ki], + mh_coverage=omh_coverage[ei,ki], lp_coverage=olp_coverage[ei,ki])) + end + + println("\nOutput-space scores by ens_size and q_prob:") + println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") + for ens in sort(unique(df_oscores.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_oscores), :k_iter) + println(" q = $(p) [chi-sq($(n_out)) bound = $(round(qq_good_o[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + println() + end + end +end + +# === 6. Perturbed-low-rank Mahalanobis (C ≈ aI + U*D*U', k = n_lowrank_modes top modes) === +# Each space yields three terms with separate reference distributions: +# top ~ Chisq(k) — miscalibration in the k principal directions +# residual ~ Chisq(n-k) — miscalibration in the n-k noise-floor directions (a = mean of bottom n-k eigenvalues) +# total ~ Chisq(n) — full-space calibration (top + residual, identical to Woodbury inverse applied to full diff) +for (space_label, top_sym, res_sym, n_dim_key) in [ + ("Forcing", :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), + ("Output", :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), +] + (haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue + n_dim = ncd.dim[n_dim_key] + top_raw = Array(ncd[top_sym]) # (n_rng, n_ens_size, n_k_iter), Union{Float64,Missing} + res_raw = Array(ncd[res_sym]) + tot_raw = top_raw .+ res_raw # missing propagates where either component is missing + + println("\n" * "="^72) + println("$(space_label)-space perturbed-low-rank M (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") + println(" Covariance model: C ≈ a*I + U*D*U' (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of empirical C)") + println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") + println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") + println(" total ~ Chisq($(n_dim)) — full-space calibration (top + residual)") + println("="^72) + + for (part_label, arr, ref_dof) in [ + ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), + ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), + ("total (n=$(n_dim))", tot_raw, n_dim), + ] + nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] + pool = collect(skipmissing(vec(arr))) + isempty(pool) && continue + qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) + emp_q = quantile(pool, calibration_check_quantiles) + coverage = nonmiss ./ n_rng + + println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") + println(" Calibration check — pooled empirical quantiles vs chi-sq($(ref_dof)):") + println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) + end + + scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_ref) + for j in axes(arr, 3), i in axes(arr, 2) + nm = nonmiss[i, j]; nm == 0 && continue + scores[idx, i, j] = sum(skipmissing(arr[:, i, j]) .<= qg; init=0) / nm + end + end + + df_plr = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + score = Union{Float64,Missing}[], + coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, + score=scores[qi,ei,ki], coverage=coverage[ei,ki])) + end + + println("\n Scores by ens_size and q_prob (expected score ≈ q_prob):") + for ens in sort(unique(df_plr.ens_size)) + println("\n === ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) + println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :score, :coverage), allrows=true) + println() + end + end + end +end + +# === 7. Marginal coverage fractions (forcing and output) === +for (space_label, cov_sym) in [("Forcing", :forcing_coverage), ("Output", :output_coverage)] + haskey(ncd, cov_sym) || continue + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs = ncd["coverage_quantile"][:] + + println("\n" * "="^72) + println("$(space_label)-space marginal coverage fractions") + println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("="^72) + println("Pooled mean coverage across all experiments:") + for (qi, qp) in enumerate(cov_qs) + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) + println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + df_cov = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mean_coverage = Union{Float64,Missing}[], + n_valid = Int[], + ) + for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = collect(skipmissing(cov[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + n_valid = length(vals))) + end + + println("\n$(space_label)-space coverage by ens_size and q_prob:") + println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + for ens in sort(unique(df_cov.ens_size)) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(cov_qs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :n_valid), allrows=true) + println() + end + end +end diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index 76cd0302d..115bf9501 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -14,7 +14,14 @@ include("Lorenz96.jl") EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg -n_pushforward_samples = 1000 +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +marginal_coverage_quantiles = [0.1, 0.5, 0.9] # quantile levels for marginal coverage fraction +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -76,6 +83,7 @@ nx = length(x0) ic_cov_sqrt = prelim_data["ic_cov_sqrt"] lorenz_config_settings = prelim_data["lorenz_config_settings"] observation_config = prelim_data["observation_config"] +y = prelim_data["y"] n_output = 2 * nx first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) @@ -102,6 +110,12 @@ forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -125,7 +139,6 @@ for (N_ens, rng_idx) in valid_file_items emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) truth_forcing_vec = forcing(emc_truth, x0) - truth_output_vec = lorenz_forward(emc_truth, x0, lorenz_config_settings, observation_config) for k in k_values post_dist = posteriors_by_k[k] @@ -175,6 +188,16 @@ for (N_ens, rng_idx) in valid_file_items f_diff = fm - truth_forcing_vec forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + for (qi, qp) in enumerate(marginal_coverage_quantiles) + forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) + end # output-space metrics (empirical distribution of pushforward samples vs truth output) os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) @@ -183,9 +206,19 @@ for (N_ens, rng_idx) in valid_file_items o_normal = MvNormal(om, oc) o_cols = Matrix(os') o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] - o_diff = om - truth_output_vec + o_diff = om - y output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) - output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, truth_output_vec) - logpdf(o_normal, o_mode) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end end end @@ -202,6 +235,7 @@ defDim(ds, "param_dim_2", n_params) defDim(ds, "pushforward_sample", n_pushforward_samples) defDim(ds, "forcing_dim", n_forcing) defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -215,6 +249,10 @@ k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + # Data variables n_evals_v = defVar( @@ -302,7 +340,7 @@ output_mahal_v = defVar( ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) -output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - truth_output)' Co^{-1} (om - truth_output), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples" +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - y)' Co^{-1} (om - y), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples, and y is the observations the posterior was fit to" output_mahal_v[:, :, :] = output_mahal_arr output_logpdf_true_v_map_v = defVar( @@ -310,8 +348,32 @@ output_logpdf_true_v_map_v = defVar( ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) -output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), truth_output) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution" +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution, with y the observations the posterior was fit to" output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr +forcing_plr_top_v = defVar(ds, "forcing_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +forcing_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in forcing space: sum((Vf'*(fm-truth_forcing))^2 / λf_i), top $(n_lowrank_modes) eigenvectors/values of Cf ≈ a_f*I + Uf*Df*Uf'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add forcing_plr_mahalanobis_residual for the full Chisq($(n_forcing)) statistic." +forcing_plr_top_v[:, :, :] = forcing_plr_mahal_top_arr + +forcing_plr_res_v = defVar(ds, "forcing_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +forcing_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in forcing space: (||fm-truth_forcing||^2 - ||Vf'*(fm-truth_forcing)||^2) / a_f, where a_f = mean of the bottom $(n_forcing - n_lowrank_modes) eigenvalues of Cf. Reference distribution: Chisq($(n_forcing - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add forcing_plr_mahalanobis_top for the full Chisq($(n_forcing)) statistic." +forcing_plr_res_v[:, :, :] = forcing_plr_mahal_residual_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) eigenvectors/values of Co ≈ a_o*I + Uo*Do*Uo'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add output_plr_mahalanobis_residual for the full Chisq($(n_output)) statistic." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o, where a_o = mean of the bottom $(n_output - n_lowrank_modes) eigenvalues of Co. Reference distribution: Chisq($(n_output - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add output_plr_mahalanobis_top for the full Chisq($(n_output)) statistic." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +forcing_coverage_v = defVar(ds, "forcing_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +forcing_coverage_v.attrib["description"] = "Marginal coverage in forcing space: fraction of forcing dims where truth_forcing[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +forcing_coverage_v[:, :, :, :] = forcing_coverage_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +output_coverage_v[:, :, :, :] = output_coverage_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl index 715ac7e85..4615b60a2 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -81,6 +81,7 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ posteriors_by_k = loaded_p["posteriors_by_k"] priors = loaded_p["priors"] k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] (truth_params_obj, _) = loaded_c["truth_params_structure"] @@ -140,7 +141,7 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings truth_frc_m = is_const ? truth_forcing[1:1] : truth_forcing - param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) + param_mah, param_lp = pushforward_metrics(push_ensemble, truth_params) forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) output_mah, output_lp = pushforward_metrics(G_ens, y) n_frc = length(truth_frc_m) diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 95a88ab46..b734eac6a 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -12,7 +12,14 @@ using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") include("Lorenz96.jl") -n_pushforward_samples = 1000 +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +marginal_coverage_quantiles = [0.1, 0.5, 0.9] # quantile levels for marginal coverage fraction +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -74,6 +81,7 @@ nx = length(x0) ic_cov_sqrt = prelim_data["ic_cov_sqrt"] lorenz_config_settings = prelim_data["lorenz_config_settings"] observation_config = prelim_data["observation_config"] +y = prelim_data["y"] n_output = 2 * nx first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) @@ -100,6 +108,12 @@ forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -123,7 +137,6 @@ for (N_ens, rng_idx) in valid_file_items emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) truth_forcing_vec = forcing(emc_truth, x0) - truth_output_vec = lorenz_forward(emc_truth, x0, lorenz_config_settings, observation_config) for k in k_values post_dist = posteriors_by_k[k] @@ -173,6 +186,16 @@ for (N_ens, rng_idx) in valid_file_items f_diff = fm - truth_forcing_vec forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + for (qi, qp) in enumerate(marginal_coverage_quantiles) + forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) + end # output-space metrics (empirical distribution of pushforward samples vs truth output) os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) @@ -181,9 +204,19 @@ for (N_ens, rng_idx) in valid_file_items o_normal = MvNormal(om, oc) o_cols = Matrix(os') o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] - o_diff = om - truth_output_vec + o_diff = om - y output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) - output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, truth_output_vec) - logpdf(o_normal, o_mode) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end end end @@ -200,6 +233,7 @@ defDim(ds, "param_dim_2", n_params) defDim(ds, "pushforward_sample", n_pushforward_samples) defDim(ds, "forcing_dim", n_forcing) defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -213,6 +247,10 @@ k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + # Data variables n_evals_v = defVar( @@ -300,7 +338,7 @@ output_mahal_v = defVar( ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) -output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - truth_output)' Co^{-1} (om - truth_output), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples" +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - y)' Co^{-1} (om - y), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples, and y is the observations the posterior was fit to" output_mahal_v[:, :, :] = output_mahal_arr output_logpdf_true_v_map_v = defVar( @@ -308,8 +346,32 @@ output_logpdf_true_v_map_v = defVar( ("random_seed", "ensemble_size", "k_iter"); fillvalue = NaN, ) -output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), truth_output) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution" +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution, with y the observations the posterior was fit to" output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr +forcing_plr_top_v = defVar(ds, "forcing_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +forcing_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in forcing space: sum((Vf'*(fm-truth_forcing))^2 / λf_i), top $(n_lowrank_modes) eigenvectors/values of Cf ≈ a_f*I + Uf*Df*Uf'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add forcing_plr_mahalanobis_residual for the full Chisq($(n_forcing)) statistic." +forcing_plr_top_v[:, :, :] = forcing_plr_mahal_top_arr + +forcing_plr_res_v = defVar(ds, "forcing_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +forcing_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in forcing space: (||fm-truth_forcing||^2 - ||Vf'*(fm-truth_forcing)||^2) / a_f, where a_f = mean of the bottom $(n_forcing - n_lowrank_modes) eigenvalues of Cf. Reference distribution: Chisq($(n_forcing - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add forcing_plr_mahalanobis_top for the full Chisq($(n_forcing)) statistic." +forcing_plr_res_v[:, :, :] = forcing_plr_mahal_residual_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) eigenvectors/values of Co ≈ a_o*I + Uo*Do*Uo'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add output_plr_mahalanobis_residual for the full Chisq($(n_output)) statistic." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o, where a_o = mean of the bottom $(n_output - n_lowrank_modes) eigenvalues of Co. Reference distribution: Chisq($(n_output - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add output_plr_mahalanobis_top for the full Chisq($(n_output)) statistic." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +forcing_coverage_v = defVar(ds, "forcing_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +forcing_coverage_v.attrib["description"] = "Marginal coverage in forcing space: fraction of forcing dims where truth_forcing[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +forcing_coverage_v[:, :, :, :] = forcing_coverage_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +output_coverage_v[:, :, :, :] = output_coverage_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl index 6e8686165..60ff932e8 100644 --- a/examples/EKIRace/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -142,7 +142,6 @@ for (N_ens, rng_idx) in valid_file_items truth_params_obj.sample_range ) end - truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) is_const = force_case == "const-force" @@ -191,7 +190,7 @@ for (N_ens, rng_idx) in valid_file_items # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings truth_frc_m = is_const ? truth_forcing[1:1] : truth_forcing - param_mah, param_lp = pushforward_metrics(constrained_push_ensemble, truth_params_constrained) + param_mah, param_lp = pushforward_metrics(push_ensemble, truth_params) forcing_mah, forcing_lp = pushforward_metrics(frc_samples_m, truth_frc_m) output_mah, output_lp = pushforward_metrics(G_ens, y) n_frc = length(truth_frc_m) From 3c59c8883fc05ebbaaf7bd01fc63ed07ea083835 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 4 Jun 2026 12:05:58 -0700 Subject: [PATCH 66/86] change limits --- .../EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index 964696335..27a2db488 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -11,7 +11,7 @@ #SBATCH --job-name=post_diag #SBATCH --output=output/slurm/post_diag_%A_%a.out #SBATCH --error=output/slurm/post_diag_%A_%a.err -#SBATCH --array=1-60%20 +#SBATCH --array=1-60%100 #SBATCH --time=00:30:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 From 4c15b9e2b244f3e0aabc42144c41baf69a502be0 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 8 Jun 2026 15:18:24 -0700 Subject: [PATCH 67/86] protections for l96 --- .../l96_exp_to_leaderboard_utilities.jl | 36 +++++++++++-------- .../l96_exp_to_leaderboard_utilities.jl | 36 +++++++++++-------- 2 files changed, 44 insertions(+), 28 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index 115bf9501..beca40cd0 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -188,13 +188,17 @@ for (N_ens, rng_idx) in valid_file_items f_diff = fm - truth_forcing_vec forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) - Ff = eigen(fc) - a_f = mean(Ff.values[1:end-n_lowrank_modes]) - V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] - λ_f = Ff.values[end-n_lowrank_modes+1:end] - proj_f = V_f' * f_diff - forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) - forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + if a_f < 1e-8 + @warn "Forcing-space PLR skipped: noise floor a_f=$(a_f) ≈ regularization level (covariance near-degenerate, e.g. const-force). Entries left as NaN." + else + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + end for (qi, qp) in enumerate(marginal_coverage_quantiles) forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) end @@ -209,13 +213,17 @@ for (N_ens, rng_idx) in valid_file_items o_diff = om - y output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) - Fo = eigen(oc) - a_o = mean(Fo.values[1:end-n_lowrank_modes]) - V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] - λ_o = Fo.values[end-n_lowrank_modes+1:end] - proj_o = V_o' * o_diff - output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) - output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level (covariance near-degenerate). Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end for (qi, qp) in enumerate(marginal_coverage_quantiles) output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) end diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index b734eac6a..a96d001a4 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -186,13 +186,17 @@ for (N_ens, rng_idx) in valid_file_items f_diff = fm - truth_forcing_vec forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) - Ff = eigen(fc) - a_f = mean(Ff.values[1:end-n_lowrank_modes]) - V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] - λ_f = Ff.values[end-n_lowrank_modes+1:end] - proj_f = V_f' * f_diff - forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) - forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + if a_f < 1e-8 + @warn "Forcing-space PLR skipped: noise floor a_f=$(a_f) ≈ regularization level (covariance near-degenerate, e.g. const-force). Entries left as NaN." + else + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + end for (qi, qp) in enumerate(marginal_coverage_quantiles) forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) end @@ -207,13 +211,17 @@ for (N_ens, rng_idx) in valid_file_items o_diff = om - y output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) - Fo = eigen(oc) - a_o = mean(Fo.values[1:end-n_lowrank_modes]) - V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] - λ_o = Fo.values[end-n_lowrank_modes+1:end] - proj_o = V_o' * o_diff - output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) - output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level (covariance near-degenerate). Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end for (qi, qp) in enumerate(marginal_coverage_quantiles) output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) end From 678e913fcfac29e511ca2294c461f2a27dbfc531 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 8 Jun 2026 15:18:51 -0700 Subject: [PATCH 68/86] toggle what is displayed in metrics --- .../EKIRace/compute_leaderboard_metrics.jl | 527 +++++++++++------- .../compute_leaderboard_metrics.jl | 527 +++++++++++------- 2 files changed, 656 insertions(+), 398 deletions(-) diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl index 18f66b190..f38fd760c 100644 --- a/examples/EKIRace/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -20,97 +20,166 @@ filename = joinpath(indir, filename) calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +########################################################################### +#################### Display filters #################################### +########################################################################### +# +# TWO independent filter axes — set either to `nothing` to show everything +# available in the file, or to a subset vector to restrict output. +# +# display_spaces — which spaces to evaluate metrics in: +# :param — parameter space (posterior mean vs true parameter) +# :forcing — forcing space (pushforward posterior vs true forcing) +# :output — output space (pushforward posterior vs observations) +# +# display_metrics — which metric types to display: +# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) +# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) +# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) +# :coverage — marginal coverage fractions (forcing and output only) +# +# display_ens_sizes — which ensemble-size values (integers) to show in tables. + +display_spaces = [:output] # e.g. [:param, :forcing] +display_metrics = [:coverage] # e.g. [:mahalanobis, :plr_mahalanobis] +display_ens_sizes = nothing # e.g. [10, 50] + +show_metric(m) = isnothing(display_metrics) || m in display_metrics +show_space(s) = isnothing(display_spaces) || s in display_spaces +show_ens(e) = isnothing(display_ens_sizes) || e in display_ens_sizes + +########################################################################### +#################### Load file and report available options ############## +########################################################################### + ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] -# Gaussian: mh = -2lp (n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) +ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end +kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end -mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] -lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] - -qq_good = quantile(Chisq(n_par), calibration_check_quantiles) +# Discover what is actually present in this file +avail_spaces = Symbol[:param] +avail_metrics = Symbol[:mahalanobis, :logratio] +(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage")) && push!(avail_spaces, :forcing) +(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage")) && push!(avail_spaces, :output) +(haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "output_plr_mahalanobis_top")) && push!(avail_metrics, :plr_mahalanobis) +(haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) + +println("="^72) +println("File: $(filename)") +println() +println("All available options:") +println(" display_spaces — $(avail_spaces)") +println(" display_metrics — $(avail_metrics)") +println(" display_ens_sizes — $(ens_vals)") +println() +println("Currently selected (edit display_* variables above to filter):") +println(" display_spaces = $(isnothing(display_spaces) ? "nothing → $(avail_spaces)" : display_spaces)") +println(" display_metrics = $(isnothing(display_metrics) ? "nothing → $(avail_metrics)" : display_metrics)") +println(" display_ens_sizes = $(isnothing(display_ens_sizes) ? "nothing → $(ens_vals)" : display_ens_sizes)") +println("="^72) + +mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] +mh_coverage = mh_nonmiss ./ n_rng +lp_coverage = lp_nonmiss ./ n_rng +qq_good = quantile(Chisq(n_par), calibration_check_quantiles) -# === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === -# If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. -mh_pool = collect(skipmissing(vec(Array(mh)))) -lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) +########################################################################### +#################### Param-space: Mahalanobis and log-ratio ############## +########################################################################### -mh_emp_q = quantile(mh_pool, calibration_check_quantiles) -lp_emp_q = quantile(lp_pool, calibration_check_quantiles) +if show_space(:param) && (show_metric(:mahalanobis) || show_metric(:logratio)) -println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") -println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) -for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) -end + println("\n" * "="^72) + println("Param-space (chi-sq ref dim = $(n_par))") + println("="^72) -# === 2. Per-experiment scores === -# missing where no valid seeds; * in printout where coverage < n_rng (less trusted) -mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) -lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + mh_pool = collect(skipmissing(vec(Array(mh)))) + lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) + mh_emp_q = quantile(mh_pool, calibration_check_quantiles) + lp_emp_q = quantile(lp_pool, calibration_check_quantiles) -for (idx, qg) in enumerate(qq_good) - for j in axes(mh, 3), i in axes(mh, 2) - nm = mh_nonmiss[i, j] - nm == 0 && continue - mh_scores[idx, i, j] = sum(skipmissing(mh[:, i, j]) .<= qg; init=0) / nm + println("Pooled-empirical calibration check vs chi-sq($(n_par)) (should match if metric is correct):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp")) + println(hdr) + for (i, p) in enumerate(calibration_check_quantiles) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", mh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", lp_emp_q[i])) + println(row) end - for j in axes(lp, 3), i in axes(lp, 2) - nm = lp_nonmiss[i, j] - nm == 0 && continue - lp_scores[idx, i, j] = sum(skipmissing(-2 .* lp[:, i, j]) .<= qg; init=0) / nm + + mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + mh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + lp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good) + for j in axes(mh, 3), i in axes(mh, 2) + nm = mh_nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(mh[:, i, j])) .<= qg + mh_scores[idx, i, j] = mean(ind) + mh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end + for j in axes(lp, 3), i in axes(lp, 2) + nm = lp_nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(-2 .* lp[:, i, j])) .<= qg + lp_scores[idx, i, j] = mean(ind) + lp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end end -end -mh_coverage = mh_nonmiss ./ n_rng # 1.0 = all n_rng seeds valid; < 1 = partial -lp_coverage = lp_nonmiss ./ n_rng + df_param = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mh_score = Union{Float64,Missing}[], + mh_score_std = Union{Float64,Missing}[], + lp_score = Union{Float64,Missing}[], + lp_score_std = Union{Float64,Missing}[], + mh_coverage = Float64[], + lp_coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_param, ( + q_prob = p, ens_size = ev, k_iter = kv, + mh_score = mh_scores[qi, ei, ki], + mh_score_std = mh_scores_std[qi, ei, ki], + lp_score = lp_scores[qi, ei, ki], + lp_score_std = lp_scores_std[qi, ei, ki], + mh_coverage = mh_coverage[ei, ki], + lp_coverage = lp_coverage[ei, ki], + )) + end -# === 3. Store scores as DataFrame === -ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end -kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) -df_scores = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mh_score = Union{Float64,Missing}[], - lp_score = Union{Float64,Missing}[], - mh_coverage = Float64[], - lp_coverage = Float64[], -) - -for (qi, p) in enumerate(calibration_check_quantiles) - for (ei, ev) in enumerate(ens_vals) - for (ki, kv) in enumerate(kiter_vals) - push!(df_scores, ( - q_prob = p, - ens_size = ev, - k_iter = kv, - mh_score = mh_scores[qi, ei, ki], - lp_score = lp_scores[qi, ei, ki], - mh_coverage = mh_coverage[ei, ki], - lp_coverage = lp_coverage[ei, ki], - )) + println("\nScores by ens_size and k_iter (missing = no valid seeds; _coverage < 1 = less trusted)") + println(" Expected: score at each q ≈ q for a well-calibrated posterior") + for ens in filter(show_ens, sort(unique(df_param.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_param), :k_iter) + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, score_cols), allrows=true) + println() end end end -# rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) -println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") -println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") -for ens in sort(unique(df_scores.ens_size)) - println("\n=== ens_size = $(ens) ===") - for (qi, p) in enumerate(calibration_check_quantiles) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) - println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) - println() - end -end +########################################################################### +#################### Forcing-space: Mahalanobis and log-ratio ############ +########################################################################### + +if show_space(:forcing) && (show_metric(:mahalanobis) || show_metric(:logratio)) && + haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") -# === 4. Forcing-space M and L scores === -if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") n_for = ncd.dim["forcing_dim"] fmh = ncd[:forcing_mahalanobis] flp = ncd[:forcing_logpdf_true_v_map] @@ -124,24 +193,36 @@ if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map" flp_emp_q = quantile(flp_pool, calibration_check_quantiles) println("\n" * "="^72) - println("Forcing-space M and L (chi-sq ref dim = $(n_for))") + println("Forcing-space (chi-sq ref dim = $(n_for))") println("="^72) - println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_for)):") - println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_f", "-2*lp_f")) + println("Pooled-empirical calibration check vs chi-sq($(n_for)):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh_f")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp_f")) + println(hdr) for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_f[i], fmh_emp_q[i], flp_emp_q[i])) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good_f[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", fmh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", flp_emp_q[i])) + println(row) end - fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + fmh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good_f) for j in axes(fmh, 3), i in axes(fmh, 2) nm = fmh_nonmiss[i, j]; nm == 0 && continue - fmh_scores[idx, i, j] = sum(skipmissing(fmh[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(fmh[:, i, j])) .<= qg + fmh_scores[idx, i, j] = mean(ind) + fmh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end for j in axes(flp, 3), i in axes(flp, 2) nm = flp_nonmiss[i, j]; nm == 0 && continue - flp_scores[idx, i, j] = sum(skipmissing(-2 .* flp[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(-2 .* flp[:, i, j])) .<= qg + flp_scores[idx, i, j] = mean(ind) + flp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end end @@ -149,30 +230,40 @@ if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map" flp_coverage = flp_nonmiss ./ n_rng df_fscores = DataFrame( q_prob=Float64[], ens_size=Int[], k_iter=Int[], - mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_score=Union{Float64,Missing}[], mh_score_std=Union{Float64,Missing}[], + lp_score=Union{Float64,Missing}[], lp_score_std=Union{Float64,Missing}[], mh_coverage=Float64[], lp_coverage=Float64[], ) for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) push!(df_fscores, (q_prob=p, ens_size=ev, k_iter=kv, - mh_score=fmh_scores[qi,ei,ki], lp_score=flp_scores[qi,ei,ki], + mh_score=fmh_scores[qi,ei,ki], mh_score_std=fmh_scores_std[qi,ei,ki], + lp_score=flp_scores[qi,ei,ki], lp_score_std=flp_scores_std[qi,ei,ki], mh_coverage=fmh_coverage[ei,ki], lp_coverage=flp_coverage[ei,ki])) end - println("\nForcing-space scores by ens_size and q_prob:") - println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") - for ens in sort(unique(df_fscores.ens_size)) + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) + + println("\nScores by ens_size and k_iter (expected score ≈ q_prob for calibrated posterior)") + for ens in filter(show_ens, sort(unique(df_fscores.ens_size))) println("\n=== ens_size = $(ens) ===") for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_fscores), :k_iter) println(" q = $(p) [chi-sq($(n_for)) bound = $(round(qq_good_f[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + show(stdout, select(sub, score_cols), allrows=true) println() end end end -# === 5. Output-space M and L scores === -if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") +########################################################################### +#################### Output-space: Mahalanobis and log-ratio ############# +########################################################################### + +if show_space(:output) && (show_metric(:mahalanobis) || show_metric(:logratio)) && + haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") + n_out = ncd.dim["output_dim"] omh = ncd[:output_mahalanobis] olp = ncd[:output_logpdf_true_v_map] @@ -186,24 +277,36 @@ if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") olp_emp_q = quantile(olp_pool, calibration_check_quantiles) println("\n" * "="^72) - println("Output-space M and L (chi-sq ref dim = $(n_out))") + println("Output-space (chi-sq ref dim = $(n_out))") println("="^72) - println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_out)):") - println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_o", "-2*lp_o")) + println("Pooled-empirical calibration check vs chi-sq($(n_out)):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh_o")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp_o")) + println(hdr) for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_o[i], omh_emp_q[i], olp_emp_q[i])) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good_o[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", omh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", olp_emp_q[i])) + println(row) end - omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + omh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good_o) for j in axes(omh, 3), i in axes(omh, 2) nm = omh_nonmiss[i, j]; nm == 0 && continue - omh_scores[idx, i, j] = sum(skipmissing(omh[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(omh[:, i, j])) .<= qg + omh_scores[idx, i, j] = mean(ind) + omh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end for j in axes(olp, 3), i in axes(olp, 2) nm = olp_nonmiss[i, j]; nm == 0 && continue - olp_scores[idx, i, j] = sum(skipmissing(-2 .* olp[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(-2 .* olp[:, i, j])) .<= qg + olp_scores[idx, i, j] = mean(ind) + olp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end end @@ -211,142 +314,168 @@ if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") olp_coverage = olp_nonmiss ./ n_rng df_oscores = DataFrame( q_prob=Float64[], ens_size=Int[], k_iter=Int[], - mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_score=Union{Float64,Missing}[], mh_score_std=Union{Float64,Missing}[], + lp_score=Union{Float64,Missing}[], lp_score_std=Union{Float64,Missing}[], mh_coverage=Float64[], lp_coverage=Float64[], ) for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) push!(df_oscores, (q_prob=p, ens_size=ev, k_iter=kv, - mh_score=omh_scores[qi,ei,ki], lp_score=olp_scores[qi,ei,ki], + mh_score=omh_scores[qi,ei,ki], mh_score_std=omh_scores_std[qi,ei,ki], + lp_score=olp_scores[qi,ei,ki], lp_score_std=olp_scores_std[qi,ei,ki], mh_coverage=omh_coverage[ei,ki], lp_coverage=olp_coverage[ei,ki])) end - println("\nOutput-space scores by ens_size and q_prob:") - println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") - for ens in sort(unique(df_oscores.ens_size)) + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) + + println("\nScores by ens_size and k_iter (expected score ≈ q_prob for calibrated posterior)") + for ens in filter(show_ens, sort(unique(df_oscores.ens_size))) println("\n=== ens_size = $(ens) ===") for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_oscores), :k_iter) println(" q = $(p) [chi-sq($(n_out)) bound = $(round(qq_good_o[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + show(stdout, select(sub, score_cols), allrows=true) println() end end end -# === 6. Perturbed-low-rank Mahalanobis (C ≈ aI + U*D*U', k = n_lowrank_modes top modes) === -# Each space yields three terms with separate reference distributions: +########################################################################### +#################### Perturbed-low-rank Mahalanobis ###################### +########################################################################### +# C ≈ aI + U*D*U' (k = n_lowrank_modes top modes). Three terms per space: # top ~ Chisq(k) — miscalibration in the k principal directions -# residual ~ Chisq(n-k) — miscalibration in the n-k noise-floor directions (a = mean of bottom n-k eigenvalues) -# total ~ Chisq(n) — full-space calibration (top + residual, identical to Woodbury inverse applied to full diff) -for (space_label, top_sym, res_sym, n_dim_key) in [ - ("Forcing", :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), - ("Output", :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), -] - (haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue - n_dim = ncd.dim[n_dim_key] - top_raw = Array(ncd[top_sym]) # (n_rng, n_ens_size, n_k_iter), Union{Float64,Missing} - res_raw = Array(ncd[res_sym]) - tot_raw = top_raw .+ res_raw # missing propagates where either component is missing +# residual ~ Chisq(n-k) — miscalibration in the noise-floor directions +# total ~ Chisq(n) — full-space (= top + residual) - println("\n" * "="^72) - println("$(space_label)-space perturbed-low-rank M (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") - println(" Covariance model: C ≈ a*I + U*D*U' (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of empirical C)") - println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") - println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") - println(" total ~ Chisq($(n_dim)) — full-space calibration (top + residual)") - println("="^72) - - for (part_label, arr, ref_dof) in [ - ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), - ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), - ("total (n=$(n_dim))", tot_raw, n_dim), +if show_metric(:plr_mahalanobis) + for (space_label, space_sym, top_sym, res_sym, n_dim_key) in [ + ("Forcing", :forcing, :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), + ("Output", :output, :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), ] - nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] - pool = collect(skipmissing(vec(arr))) - isempty(pool) && continue - qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) - emp_q = quantile(pool, calibration_check_quantiles) - coverage = nonmiss ./ n_rng - - println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") - println(" Calibration check — pooled empirical quantiles vs chi-sq($(ref_dof)):") - println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) - for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) - end + (show_space(space_sym) && haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue + + n_dim = ncd.dim[n_dim_key] + top_raw = Array(ncd[top_sym]) + res_raw = Array(ncd[res_sym]) + tot_raw = top_raw .+ res_raw # missing propagates where either component is missing + + println("\n" * "="^72) + println("$(space_label)-space perturbed-low-rank Mahalanobis (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") + println(" Covariance model: C ≈ a·I + U·D·Uᵀ (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of C)") + println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") + println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") + println(" total ~ Chisq($(n_dim)) — full-space (top + residual)") + println("="^72) + + for (part_label, arr, ref_dof) in [ + ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), + ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), + ("total (n=$(n_dim))", tot_raw, n_dim), + ] + nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] + pool = collect(skipmissing(vec(arr))) + isempty(pool) && continue + qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) + emp_q = quantile(pool, calibration_check_quantiles) + coverage = nonmiss ./ n_rng + + println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") + println(" Pooled-empirical calibration check vs chi-sq($(ref_dof)):") + println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) + end - scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - for (idx, qg) in enumerate(qq_ref) - for j in axes(arr, 3), i in axes(arr, 2) - nm = nonmiss[i, j]; nm == 0 && continue - scores[idx, i, j] = sum(skipmissing(arr[:, i, j]) .<= qg; init=0) / nm + scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_ref) + for j in axes(arr, 3), i in axes(arr, 2) + nm = nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(arr[:, i, j])) .<= qg + scores[idx, i, j] = mean(ind) + scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end end - end - df_plr = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - score = Union{Float64,Missing}[], - coverage = Float64[], - ) - for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, - score=scores[qi,ei,ki], coverage=coverage[ei,ki])) - end + df_plr = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + score = Union{Float64,Missing}[], + score_std = Union{Float64,Missing}[], + coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, + score=scores[qi,ei,ki], score_std=scores_std[qi,ei,ki], coverage=coverage[ei,ki])) + end - println("\n Scores by ens_size and q_prob (expected score ≈ q_prob):") - for ens in sort(unique(df_plr.ens_size)) - println("\n === ens_size = $(ens) ===") - for (qi, p) in enumerate(calibration_check_quantiles) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) - println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :score, :coverage), allrows=true) - println() + println("\n Scores by ens_size and k_iter (expected score ≈ q_prob):") + for ens in filter(show_ens, sort(unique(df_plr.ens_size))) + println("\n === ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) + println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :score, :score_std, :coverage), allrows=true) + println() + end end end end end -# === 7. Marginal coverage fractions (forcing and output) === -for (space_label, cov_sym) in [("Forcing", :forcing_coverage), ("Output", :output_coverage)] - haskey(ncd, cov_sym) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs = ncd["coverage_quantile"][:] +########################################################################### +#################### Marginal coverage fractions ######################### +########################################################################### - println("\n" * "="^72) - println("$(space_label)-space marginal coverage fractions") - println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") - println("="^72) - println("Pooled mean coverage across all experiments:") - for (qi, qp) in enumerate(cov_qs) - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) - println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) - end +if show_metric(:coverage) + for (space_label, space_sym, cov_sym) in [ + ("Forcing", :forcing, :forcing_coverage), + ("Output", :output, :output_coverage), + ] + (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - df_cov = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mean_coverage = Union{Float64,Missing}[], - n_valid = Int[], - ) - for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - vals = collect(skipmissing(cov[:, ei, ki, qi])) - push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, - mean_coverage = isempty(vals) ? missing : mean(vals), - n_valid = length(vals))) - end + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs = ncd["coverage_quantile"][:] - println("\n$(space_label)-space coverage by ens_size and q_prob:") - println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") - for ens in sort(unique(df_cov.ens_size)) - println("\n=== ens_size = $(ens) ===") + println("\n" * "="^72) + println("$(space_label)-space marginal coverage fractions") + println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("="^72) + println("Pooled mean coverage across all experiments:") for (qi, qp) in enumerate(cov_qs) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) - println(" q = $(qp) (expected ≈ $(qp))") - show(stdout, select(sub, :k_iter, :mean_coverage, :n_valid), allrows=true) - println() + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) + println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + df_cov = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mean_coverage = Union{Float64,Missing}[], + std_coverage = Union{Float64,Missing}[], + n_valid = Int[], + ) + for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = collect(skipmissing(cov[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + std_coverage = length(vals) > 1 ? std(vals) : missing, + n_valid = length(vals))) + end + + println("\n$(space_label)-space coverage by ens_size and k_iter:") + println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + for ens in filter(show_ens, sort(unique(df_cov.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(cov_qs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) + println() + end end end end diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index 61a6209da..d8c8258b3 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -22,97 +22,166 @@ filename = joinpath(indir, filename) calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +########################################################################### +#################### Display filters #################################### +########################################################################### +# +# TWO independent filter axes — set either to `nothing` to show everything +# available in the file, or to a subset vector to restrict output. +# +# display_spaces — which spaces to evaluate metrics in: +# :param — parameter space (posterior mean vs true parameter) +# :forcing — forcing space (pushforward posterior vs true forcing) +# :output — output space (pushforward posterior vs observations) +# +# display_metrics — which metric types to display: +# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) +# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) +# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) +# :coverage — marginal coverage fractions (forcing and output only) +# +# display_ens_sizes — which ensemble-size values (integers) to show in tables. + +display_spaces = [:output] # e.g. [:param, :forcing] +display_metrics = [:coverage] # e.g. [:mahalanobis, :plr_mahalanobis] +display_ens_sizes = nothing # e.g. [10, 50] + +show_metric(m) = isnothing(display_metrics) || m in display_metrics +show_space(s) = isnothing(display_spaces) || s in display_spaces +show_ens(e) = isnothing(display_ens_sizes) || e in display_ens_sizes + +########################################################################### +#################### Load file and report available options ############## +########################################################################### + ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] -# Gaussian: mh = -2lp (n_rng, n_ens_size, n_k_iter, n_par) = (ncd.dim[k] for k in ["random_seed","ensemble_size", "k_iter", "param_dim"]) +ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end +kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end -mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] -lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] - -qq_good = quantile(Chisq(n_par), calibration_check_quantiles) +# Discover what is actually present in this file +avail_spaces = Symbol[:param] +avail_metrics = Symbol[:mahalanobis, :logratio] +(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage")) && push!(avail_spaces, :forcing) +(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage")) && push!(avail_spaces, :output) +(haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "output_plr_mahalanobis_top")) && push!(avail_metrics, :plr_mahalanobis) +(haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) + +println("="^72) +println("File: $(filename)") +println() +println("All available options:") +println(" display_spaces — $(avail_spaces)") +println(" display_metrics — $(avail_metrics)") +println(" display_ens_sizes — $(ens_vals)") +println() +println("Currently selected (edit display_* variables above to filter):") +println(" display_spaces = $(isnothing(display_spaces) ? "nothing → $(avail_spaces)" : display_spaces)") +println(" display_metrics = $(isnothing(display_metrics) ? "nothing → $(avail_metrics)" : display_metrics)") +println(" display_ens_sizes = $(isnothing(display_ens_sizes) ? "nothing → $(ens_vals)" : display_ens_sizes)") +println("="^72) + +mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] +lp_nonmiss = sum(.!ismissing.(lp), dims=1)[1,:,:] +mh_coverage = mh_nonmiss ./ n_rng +lp_coverage = lp_nonmiss ./ n_rng +qq_good = quantile(Chisq(n_par), calibration_check_quantiles) -# === 1. Validate: pooled empirical quantiles should match chi-sq(n_par) === -# If mh is correctly computed it follows chi-sq(n_par), so empirical ≈ theoretical. -mh_pool = collect(skipmissing(vec(Array(mh)))) -lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) +########################################################################### +#################### Param-space: Mahalanobis and log-ratio ############## +########################################################################### -mh_emp_q = quantile(mh_pool, calibration_check_quantiles) -lp_emp_q = quantile(lp_pool, calibration_check_quantiles) +if show_space(:param) && (show_metric(:mahalanobis) || show_metric(:logratio)) -println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_par)) (should match if metric is correct):") -println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh", "-2*lp")) -for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good[i], mh_emp_q[i], lp_emp_q[i])) -end + println("\n" * "="^72) + println("Param-space (chi-sq ref dim = $(n_par))") + println("="^72) -# === 2. Per-experiment scores === -# missing where no valid seeds; * in printout where coverage < n_rng (less trusted) -mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) -lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + mh_pool = collect(skipmissing(vec(Array(mh)))) + lp_pool = collect(skipmissing(vec(-2 .* Array(lp)))) + mh_emp_q = quantile(mh_pool, calibration_check_quantiles) + lp_emp_q = quantile(lp_pool, calibration_check_quantiles) -for (idx, qg) in enumerate(qq_good) - for j in axes(mh, 3), i in axes(mh, 2) - nm = mh_nonmiss[i, j] - nm == 0 && continue - mh_scores[idx, i, j] = sum(skipmissing(mh[:, i, j]) .<= qg; init=0) / nm + println("Pooled-empirical calibration check vs chi-sq($(n_par)) (should match if metric is correct):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp")) + println(hdr) + for (i, p) in enumerate(calibration_check_quantiles) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", mh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", lp_emp_q[i])) + println(row) end - for j in axes(lp, 3), i in axes(lp, 2) - nm = lp_nonmiss[i, j] - nm == 0 && continue - lp_scores[idx, i, j] = sum(skipmissing(-2 .* lp[:, i, j]) .<= qg; init=0) / nm + + mh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + lp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + mh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + lp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_good) + for j in axes(mh, 3), i in axes(mh, 2) + nm = mh_nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(mh[:, i, j])) .<= qg + mh_scores[idx, i, j] = mean(ind) + mh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end + for j in axes(lp, 3), i in axes(lp, 2) + nm = lp_nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(-2 .* lp[:, i, j])) .<= qg + lp_scores[idx, i, j] = mean(ind) + lp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end end -end -mh_coverage = mh_nonmiss ./ n_rng # 1.0 = all n_rng seeds valid; < 1 = partial -lp_coverage = lp_nonmiss ./ n_rng + df_param = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mh_score = Union{Float64,Missing}[], + mh_score_std = Union{Float64,Missing}[], + lp_score = Union{Float64,Missing}[], + lp_score_std = Union{Float64,Missing}[], + mh_coverage = Float64[], + lp_coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_param, ( + q_prob = p, ens_size = ev, k_iter = kv, + mh_score = mh_scores[qi, ei, ki], + mh_score_std = mh_scores_std[qi, ei, ki], + lp_score = lp_scores[qi, ei, ki], + lp_score_std = lp_scores_std[qi, ei, ki], + mh_coverage = mh_coverage[ei, ki], + lp_coverage = lp_coverage[ei, ki], + )) + end -# === 3. Store scores as DataFrame === -ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end -kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) -df_scores = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mh_score = Union{Float64,Missing}[], - lp_score = Union{Float64,Missing}[], - mh_coverage = Float64[], - lp_coverage = Float64[], -) - -for (qi, p) in enumerate(calibration_check_quantiles) - for (ei, ev) in enumerate(ens_vals) - for (ki, kv) in enumerate(kiter_vals) - push!(df_scores, ( - q_prob = p, - ens_size = ev, - k_iter = kv, - mh_score = mh_scores[qi, ei, ki], - lp_score = lp_scores[qi, ei, ki], - mh_coverage = mh_coverage[ei, ki], - lp_coverage = lp_coverage[ei, ki], - )) + println("\nScores by ens_size and k_iter (missing = no valid seeds; _coverage < 1 = less trusted)") + println(" Expected: score at each q ≈ q for a well-calibrated posterior") + for ens in filter(show_ens, sort(unique(df_param.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_param), :k_iter) + println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, score_cols), allrows=true) + println() end end end -# rows with mh_coverage < 1 used fewer than n_rng seeds (less trusted) -println("\nScores by ens_size and q_prob (missing = no valid seeds; coverage < 1 = less trusted):") -println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") -for ens in sort(unique(df_scores.ens_size)) - println("\n=== ens_size = $(ens) ===") - for (qi, p) in enumerate(calibration_check_quantiles) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_scores), :k_iter) - println(" q = $(p) [chi-sq($(n_par)) bound = $(round(qq_good[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) - println() - end -end +########################################################################### +#################### Forcing-space: Mahalanobis and log-ratio ############ +########################################################################### + +if show_space(:forcing) && (show_metric(:mahalanobis) || show_metric(:logratio)) && + haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") -# === 4. Forcing-space M and L scores === -if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map") n_for = ncd.dim["forcing_dim"] fmh = ncd[:forcing_mahalanobis] flp = ncd[:forcing_logpdf_true_v_map] @@ -126,24 +195,36 @@ if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map" flp_emp_q = quantile(flp_pool, calibration_check_quantiles) println("\n" * "="^72) - println("Forcing-space M and L (chi-sq ref dim = $(n_for))") + println("Forcing-space (chi-sq ref dim = $(n_for))") println("="^72) - println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_for)):") - println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_f", "-2*lp_f")) + println("Pooled-empirical calibration check vs chi-sq($(n_for)):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh_f")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp_f")) + println(hdr) for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_f[i], fmh_emp_q[i], flp_emp_q[i])) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good_f[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", fmh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", flp_emp_q[i])) + println(row) end - fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + fmh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + fmh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + flp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good_f) for j in axes(fmh, 3), i in axes(fmh, 2) nm = fmh_nonmiss[i, j]; nm == 0 && continue - fmh_scores[idx, i, j] = sum(skipmissing(fmh[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(fmh[:, i, j])) .<= qg + fmh_scores[idx, i, j] = mean(ind) + fmh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end for j in axes(flp, 3), i in axes(flp, 2) nm = flp_nonmiss[i, j]; nm == 0 && continue - flp_scores[idx, i, j] = sum(skipmissing(-2 .* flp[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(-2 .* flp[:, i, j])) .<= qg + flp_scores[idx, i, j] = mean(ind) + flp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end end @@ -151,30 +232,40 @@ if haskey(ncd, "forcing_mahalanobis") && haskey(ncd, "forcing_logpdf_true_v_map" flp_coverage = flp_nonmiss ./ n_rng df_fscores = DataFrame( q_prob=Float64[], ens_size=Int[], k_iter=Int[], - mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_score=Union{Float64,Missing}[], mh_score_std=Union{Float64,Missing}[], + lp_score=Union{Float64,Missing}[], lp_score_std=Union{Float64,Missing}[], mh_coverage=Float64[], lp_coverage=Float64[], ) for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) push!(df_fscores, (q_prob=p, ens_size=ev, k_iter=kv, - mh_score=fmh_scores[qi,ei,ki], lp_score=flp_scores[qi,ei,ki], + mh_score=fmh_scores[qi,ei,ki], mh_score_std=fmh_scores_std[qi,ei,ki], + lp_score=flp_scores[qi,ei,ki], lp_score_std=flp_scores_std[qi,ei,ki], mh_coverage=fmh_coverage[ei,ki], lp_coverage=flp_coverage[ei,ki])) end - println("\nForcing-space scores by ens_size and q_prob:") - println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") - for ens in sort(unique(df_fscores.ens_size)) + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) + + println("\nScores by ens_size and k_iter (expected score ≈ q_prob for calibrated posterior)") + for ens in filter(show_ens, sort(unique(df_fscores.ens_size))) println("\n=== ens_size = $(ens) ===") for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_fscores), :k_iter) println(" q = $(p) [chi-sq($(n_for)) bound = $(round(qq_good_f[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + show(stdout, select(sub, score_cols), allrows=true) println() end end end -# === 5. Output-space M and L scores === -if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") +########################################################################### +#################### Output-space: Mahalanobis and log-ratio ############# +########################################################################### + +if show_space(:output) && (show_metric(:mahalanobis) || show_metric(:logratio)) && + haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") + n_out = ncd.dim["output_dim"] omh = ncd[:output_mahalanobis] olp = ncd[:output_logpdf_true_v_map] @@ -188,24 +279,36 @@ if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") olp_emp_q = quantile(olp_pool, calibration_check_quantiles) println("\n" * "="^72) - println("Output-space M and L (chi-sq ref dim = $(n_out))") + println("Output-space (chi-sq ref dim = $(n_out))") println("="^72) - println("Metric calibration check — pooled empirical quantiles vs chi-sq($(n_out)):") - println(@sprintf(" %-8s %-12s %-12s %-12s", "prob", "chi-sq ref", "mh_o", "-2*lp_o")) + println("Pooled-empirical calibration check vs chi-sq($(n_out)):") + hdr = @sprintf(" %-8s %-12s", "prob", "chi-sq ref") + show_metric(:mahalanobis) && (hdr *= @sprintf(" %-12s", "mh_o")) + show_metric(:logratio) && (hdr *= @sprintf(" %-12s", "-2*lp_o")) + println(hdr) for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f %-12.3f", p, qq_good_o[i], omh_emp_q[i], olp_emp_q[i])) + row = @sprintf(" %-8.2f %-12.3f", p, qq_good_o[i]) + show_metric(:mahalanobis) && (row *= @sprintf(" %-12.3f", omh_emp_q[i])) + show_metric(:logratio) && (row *= @sprintf(" %-12.3f", olp_emp_q[i])) + println(row) end - omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + omh_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + omh_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + olp_scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) for (idx, qg) in enumerate(qq_good_o) for j in axes(omh, 3), i in axes(omh, 2) nm = omh_nonmiss[i, j]; nm == 0 && continue - omh_scores[idx, i, j] = sum(skipmissing(omh[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(omh[:, i, j])) .<= qg + omh_scores[idx, i, j] = mean(ind) + omh_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end for j in axes(olp, 3), i in axes(olp, 2) nm = olp_nonmiss[i, j]; nm == 0 && continue - olp_scores[idx, i, j] = sum(skipmissing(-2 .* olp[:, i, j]) .<= qg; init=0) / nm + ind = collect(skipmissing(-2 .* olp[:, i, j])) .<= qg + olp_scores[idx, i, j] = mean(ind) + olp_scores_std[idx, i, j] = nm > 1 ? std(ind) : missing end end @@ -213,142 +316,168 @@ if haskey(ncd, "output_mahalanobis") && haskey(ncd, "output_logpdf_true_v_map") olp_coverage = olp_nonmiss ./ n_rng df_oscores = DataFrame( q_prob=Float64[], ens_size=Int[], k_iter=Int[], - mh_score=Union{Float64,Missing}[], lp_score=Union{Float64,Missing}[], + mh_score=Union{Float64,Missing}[], mh_score_std=Union{Float64,Missing}[], + lp_score=Union{Float64,Missing}[], lp_score_std=Union{Float64,Missing}[], mh_coverage=Float64[], lp_coverage=Float64[], ) for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) push!(df_oscores, (q_prob=p, ens_size=ev, k_iter=kv, - mh_score=omh_scores[qi,ei,ki], lp_score=olp_scores[qi,ei,ki], + mh_score=omh_scores[qi,ei,ki], mh_score_std=omh_scores_std[qi,ei,ki], + lp_score=olp_scores[qi,ei,ki], lp_score_std=olp_scores_std[qi,ei,ki], mh_coverage=omh_coverage[ei,ki], lp_coverage=olp_coverage[ei,ki])) end - println("\nOutput-space scores by ens_size and q_prob:") - println(" (Expected: score at each q ≈ q for a well-calibrated posterior)") - for ens in sort(unique(df_oscores.ens_size)) + score_cols = Symbol[:k_iter] + show_metric(:mahalanobis) && append!(score_cols, [:mh_score, :mh_score_std, :mh_coverage]) + show_metric(:logratio) && append!(score_cols, [:lp_score, :lp_score_std, :lp_coverage]) + + println("\nScores by ens_size and k_iter (expected score ≈ q_prob for calibrated posterior)") + for ens in filter(show_ens, sort(unique(df_oscores.ens_size))) println("\n=== ens_size = $(ens) ===") for (qi, p) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_oscores), :k_iter) println(" q = $(p) [chi-sq($(n_out)) bound = $(round(qq_good_o[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :mh_score, :lp_score, :mh_coverage), allrows=true) + show(stdout, select(sub, score_cols), allrows=true) println() end end end -# === 6. Perturbed-low-rank Mahalanobis (C ≈ aI + U*D*U', k = n_lowrank_modes top modes) === -# Each space yields three terms with separate reference distributions: +########################################################################### +#################### Perturbed-low-rank Mahalanobis ###################### +########################################################################### +# C ≈ aI + U*D*U' (k = n_lowrank_modes top modes). Three terms per space: # top ~ Chisq(k) — miscalibration in the k principal directions -# residual ~ Chisq(n-k) — miscalibration in the n-k noise-floor directions (a = mean of bottom n-k eigenvalues) -# total ~ Chisq(n) — full-space calibration (top + residual, identical to Woodbury inverse applied to full diff) -for (space_label, top_sym, res_sym, n_dim_key) in [ - ("Forcing", :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), - ("Output", :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), -] - (haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue - n_dim = ncd.dim[n_dim_key] - top_raw = Array(ncd[top_sym]) # (n_rng, n_ens_size, n_k_iter), Union{Float64,Missing} - res_raw = Array(ncd[res_sym]) - tot_raw = top_raw .+ res_raw # missing propagates where either component is missing +# residual ~ Chisq(n-k) — miscalibration in the noise-floor directions +# total ~ Chisq(n) — full-space (= top + residual) - println("\n" * "="^72) - println("$(space_label)-space perturbed-low-rank M (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") - println(" Covariance model: C ≈ a*I + U*D*U' (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of empirical C)") - println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") - println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") - println(" total ~ Chisq($(n_dim)) — full-space calibration (top + residual)") - println("="^72) - - for (part_label, arr, ref_dof) in [ - ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), - ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), - ("total (n=$(n_dim))", tot_raw, n_dim), +if show_metric(:plr_mahalanobis) + for (space_label, space_sym, top_sym, res_sym, n_dim_key) in [ + ("Forcing", :forcing, :forcing_plr_mahalanobis_top, :forcing_plr_mahalanobis_residual, "forcing_dim"), + ("Output", :output, :output_plr_mahalanobis_top, :output_plr_mahalanobis_residual, "output_dim"), ] - nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] - pool = collect(skipmissing(vec(arr))) - isempty(pool) && continue - qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) - emp_q = quantile(pool, calibration_check_quantiles) - coverage = nonmiss ./ n_rng - - println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") - println(" Calibration check — pooled empirical quantiles vs chi-sq($(ref_dof)):") - println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) - for (i, p) in enumerate(calibration_check_quantiles) - println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) - end + (show_space(space_sym) && haskey(ncd, top_sym) && haskey(ncd, res_sym)) || continue + + n_dim = ncd.dim[n_dim_key] + top_raw = Array(ncd[top_sym]) + res_raw = Array(ncd[res_sym]) + tot_raw = top_raw .+ res_raw # missing propagates where either component is missing + + println("\n" * "="^72) + println("$(space_label)-space perturbed-low-rank Mahalanobis (k=$(n_lowrank_modes) modes, n=$(n_dim) total dims)") + println(" Covariance model: C ≈ a·I + U·D·Uᵀ (a = mean of bottom $(n_dim-n_lowrank_modes) eigenvalues of C)") + println(" top ~ Chisq($(n_lowrank_modes)) — miscalibration in the k principal directions") + println(" residual ~ Chisq($(n_dim-n_lowrank_modes)) — miscalibration in the noise-floor directions") + println(" total ~ Chisq($(n_dim)) — full-space (top + residual)") + println("="^72) + + for (part_label, arr, ref_dof) in [ + ("top (k=$(n_lowrank_modes))", top_raw, n_lowrank_modes), + ("residual (n-k=$(n_dim-n_lowrank_modes))", res_raw, n_dim - n_lowrank_modes), + ("total (n=$(n_dim))", tot_raw, n_dim), + ] + nonmiss = sum(.!ismissing.(arr), dims=1)[1,:,:] + pool = collect(skipmissing(vec(arr))) + isempty(pool) && continue + qq_ref = quantile(Chisq(ref_dof), calibration_check_quantiles) + emp_q = quantile(pool, calibration_check_quantiles) + coverage = nonmiss ./ n_rng + + println("\n $(space_label)-space PLR-M $(part_label) [ref: Chisq($(ref_dof))]:") + println(" Pooled-empirical calibration check vs chi-sq($(ref_dof)):") + println(@sprintf(" %-8s %-12s %-12s", "prob", "chi-sq ref", "plr_mh")) + for (i, p) in enumerate(calibration_check_quantiles) + println(@sprintf(" %-8.2f %-12.3f %-12.3f", p, qq_ref[i], emp_q[i])) + end - scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) - for (idx, qg) in enumerate(qq_ref) - for j in axes(arr, 3), i in axes(arr, 2) - nm = nonmiss[i, j]; nm == 0 && continue - scores[idx, i, j] = sum(skipmissing(arr[:, i, j]) .<= qg; init=0) / nm + scores = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + scores_std = Array{Union{Float64,Missing}}(missing, length(calibration_check_quantiles), n_ens_size, n_k_iter) + for (idx, qg) in enumerate(qq_ref) + for j in axes(arr, 3), i in axes(arr, 2) + nm = nonmiss[i, j]; nm == 0 && continue + ind = collect(skipmissing(arr[:, i, j])) .<= qg + scores[idx, i, j] = mean(ind) + scores_std[idx, i, j] = nm > 1 ? std(ind) : missing + end end - end - df_plr = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - score = Union{Float64,Missing}[], - coverage = Float64[], - ) - for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, - score=scores[qi,ei,ki], coverage=coverage[ei,ki])) - end + df_plr = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + score = Union{Float64,Missing}[], + score_std = Union{Float64,Missing}[], + coverage = Float64[], + ) + for (qi, p) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + push!(df_plr, (q_prob=p, ens_size=ev, k_iter=kv, + score=scores[qi,ei,ki], score_std=scores_std[qi,ei,ki], coverage=coverage[ei,ki])) + end - println("\n Scores by ens_size and q_prob (expected score ≈ q_prob):") - for ens in sort(unique(df_plr.ens_size)) - println("\n === ens_size = $(ens) ===") - for (qi, p) in enumerate(calibration_check_quantiles) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) - println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") - show(stdout, select(sub, :k_iter, :score, :coverage), allrows=true) - println() + println("\n Scores by ens_size and k_iter (expected score ≈ q_prob):") + for ens in filter(show_ens, sort(unique(df_plr.ens_size))) + println("\n === ens_size = $(ens) ===") + for (qi, p) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == p, df_plr), :k_iter) + println(" q = $(p) [chi-sq($(ref_dof)) bound = $(round(qq_ref[qi], digits=2)), expected score ≈ $(p)]") + show(stdout, select(sub, :k_iter, :score, :score_std, :coverage), allrows=true) + println() + end end end end end -# === 7. Marginal coverage fractions (forcing and output) === -for (space_label, cov_sym) in [("Forcing", :forcing_coverage), ("Output", :output_coverage)] - haskey(ncd, cov_sym) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs = ncd["coverage_quantile"][:] +########################################################################### +#################### Marginal coverage fractions ######################### +########################################################################### - println("\n" * "="^72) - println("$(space_label)-space marginal coverage fractions") - println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") - println("="^72) - println("Pooled mean coverage across all experiments:") - for (qi, qp) in enumerate(cov_qs) - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) - println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) - end +if show_metric(:coverage) + for (space_label, space_sym, cov_sym) in [ + ("Forcing", :forcing, :forcing_coverage), + ("Output", :output, :output_coverage), + ] + (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - df_cov = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mean_coverage = Union{Float64,Missing}[], - n_valid = Int[], - ) - for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - vals = collect(skipmissing(cov[:, ei, ki, qi])) - push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, - mean_coverage = isempty(vals) ? missing : mean(vals), - n_valid = length(vals))) - end + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs = ncd["coverage_quantile"][:] - println("\n$(space_label)-space coverage by ens_size and q_prob:") - println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") - for ens in sort(unique(df_cov.ens_size)) - println("\n=== ens_size = $(ens) ===") + println("\n" * "="^72) + println("$(space_label)-space marginal coverage fractions") + println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("="^72) + println("Pooled mean coverage across all experiments:") for (qi, qp) in enumerate(cov_qs) - sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) - println(" q = $(qp) (expected ≈ $(qp))") - show(stdout, select(sub, :k_iter, :mean_coverage, :n_valid), allrows=true) - println() + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) + println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + df_cov = DataFrame( + q_prob = Float64[], + ens_size = Int[], + k_iter = Int[], + mean_coverage = Union{Float64,Missing}[], + std_coverage = Union{Float64,Missing}[], + n_valid = Int[], + ) + for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = collect(skipmissing(cov[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + std_coverage = length(vals) > 1 ? std(vals) : missing, + n_valid = length(vals))) + end + + println("\n$(space_label)-space coverage by ens_size and k_iter:") + println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + for ens in filter(show_ens, sort(unique(df_cov.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(cov_qs) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) + println() + end end end end From 31ed3fb89eb7fdbf3c47401d6f42a395b8f31956 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 8 Jun 2026 16:07:17 -0700 Subject: [PATCH 69/86] compute more coverage quantiles in netcdf --- .../EKIRace/compute_leaderboard_metrics.jl | 18 +++++++++------ .../compute_leaderboard_metrics.jl | 22 ++++++++++++------- .../l96_exp_to_leaderboard_utilities.jl | 12 +++++++++- .../l96_exp_to_leaderboard_utilities.jl | 12 +++++++++- 4 files changed, 47 insertions(+), 17 deletions(-) diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl index f38fd760c..cdd9404a5 100644 --- a/examples/EKIRace/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -437,17 +437,20 @@ if show_metric(:coverage) ] (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs = ncd["coverage_quantile"][:] + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs_all = Float64.(ncd["coverage_quantile"][:]) + qi_keep = findall(qp -> any(isapprox(qp, p, atol=1e-8) for p in calibration_check_quantiles), cov_qs_all) + cov_qs = cov_qs_all[qi_keep] println("\n" * "="^72) println("$(space_label)-space marginal coverage fractions") println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") println("="^72) println("Pooled mean coverage across all experiments:") - for (qi, qp) in enumerate(cov_qs) - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) - println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + for (qi_local, qi_file) in enumerate(qi_keep) + qp = cov_qs[qi_local] + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi_file])))) + println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) end df_cov = DataFrame( @@ -458,8 +461,9 @@ if show_metric(:coverage) std_coverage = Union{Float64,Missing}[], n_valid = Int[], ) - for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - vals = collect(skipmissing(cov[:, ei, ki, qi])) + for (qi_local, qi_file) in enumerate(qi_keep), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + qp = cov_qs[qi_local] + vals = collect(skipmissing(cov[:, ei, ki, qi_file])) push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, mean_coverage = isempty(vals) ? missing : mean(vals), std_coverage = length(vals) > 1 ? std(vals) : missing, diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index d8c8258b3..a11902a22 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -19,7 +19,7 @@ filename = joinpath(indir, filename) #################### Metric parameters ################################### ########################################################################### -calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage +calibration_check_quantiles = [0.2, 0.5, 0.8] # quantile levels for calibration check scores and coverage n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis ########################################################################### @@ -439,17 +439,20 @@ if show_metric(:coverage) ] (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs = ncd["coverage_quantile"][:] + cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) + cov_qs_all = Float64.(ncd["coverage_quantile"][:]) + qi_keep = findall(qp -> any(isapprox(qp, p, atol=1e-8) for p in calibration_check_quantiles), cov_qs_all) + cov_qs = cov_qs_all[qi_keep] println("\n" * "="^72) println("$(space_label)-space marginal coverage fractions") println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") println("="^72) println("Pooled mean coverage across all experiments:") - for (qi, qp) in enumerate(cov_qs) - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi])))) - println(@sprintf(" q = %.1f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + for (qi_local, qi_file) in enumerate(qi_keep) + qp = cov_qs[qi_local] + pool = collect(skipmissing(vec(Array(cov[:, :, :, qi_file])))) + println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) end df_cov = DataFrame( @@ -460,8 +463,9 @@ if show_metric(:coverage) std_coverage = Union{Float64,Missing}[], n_valid = Int[], ) - for (qi, qp) in enumerate(cov_qs), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - vals = collect(skipmissing(cov[:, ei, ki, qi])) + for (qi_local, qi_file) in enumerate(qi_keep), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + qp = cov_qs[qi_local] + vals = collect(skipmissing(cov[:, ei, ki, qi_file])) push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, mean_coverage = isempty(vals) ? missing : mean(vals), std_coverage = length(vals) > 1 ? std(vals) : missing, @@ -481,3 +485,5 @@ if show_metric(:coverage) end end end + +@info calibration_check_quantiles diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index beca40cd0..3a877d678 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -20,7 +20,7 @@ EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis -marginal_coverage_quantiles = [0.1, 0.5, 0.9] # quantile levels for marginal coverage fraction +marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -116,6 +116,7 @@ output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) +truth_forcing_arr = fill(NaN, n_rng, n_ens, n_forcing) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -139,6 +140,7 @@ for (N_ens, rng_idx) in valid_file_items emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) truth_forcing_vec = forcing(emc_truth, x0) + truth_forcing_arr[i, j, :] = truth_forcing_vec for k in k_values post_dist = posteriors_by_k[k] @@ -261,6 +263,14 @@ cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" cov_q_var[:] = marginal_coverage_quantiles +truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble_size", "forcing_dim"); fillvalue=NaN) +truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair. Use with forcing_samples to recompute marginal coverage at arbitrary quantile levels." +truth_forcing_v[:, :, :] = truth_forcing_arr + +truth_output_v = defVar(ds, "truth_output", Float64, ("output_dim",); fillvalue=NaN) +truth_output_v.attrib["description"] = "Observation vector y used as truth in output space. Constant across seeds and ensemble sizes. Use with output_samples to recompute marginal coverage at arbitrary quantile levels." +truth_output_v[:] = y + # Data variables n_evals_v = defVar( diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index a96d001a4..0ac4bc09e 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -18,7 +18,7 @@ include("Lorenz96.jl") n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis -marginal_coverage_quantiles = [0.1, 0.5, 0.9] # quantile levels for marginal coverage fraction +marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -114,6 +114,7 @@ output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) +truth_forcing_arr = fill(NaN, n_rng, n_ens, n_forcing) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -137,6 +138,7 @@ for (N_ens, rng_idx) in valid_file_items emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) truth_forcing_vec = forcing(emc_truth, x0) + truth_forcing_arr[i, j, :] = truth_forcing_vec for k in k_values post_dist = posteriors_by_k[k] @@ -259,6 +261,14 @@ cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" cov_q_var[:] = marginal_coverage_quantiles +truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble_size", "forcing_dim"); fillvalue=NaN) +truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair. Use with forcing_samples to recompute marginal coverage at arbitrary quantile levels." +truth_forcing_v[:, :, :] = truth_forcing_arr + +truth_output_v = defVar(ds, "truth_output", Float64, ("output_dim",); fillvalue=NaN) +truth_output_v.attrib["description"] = "Observation vector y used as truth in output space. Constant across seeds and ensemble sizes. Use with output_samples to recompute marginal coverage at arbitrary quantile levels." +truth_output_v[:] = y + # Data variables n_evals_v = defVar( From 00c2302dc4b872956b2e0fef13f326877f331808 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 10 Jun 2026 11:13:33 -0700 Subject: [PATCH 70/86] update l63 and l96 metrics -> focus on the output metrics, coverage in particular --- .../l63_exp_to_leaderboard_utilities.jl | 139 +++++++++++++++++- .../l96_exp_to_leaderboard_utilities.jl | 4 - .../l63_exp_to_leaderboard_utilities.jl | 139 +++++++++++++++++- .../l96_exp_to_leaderboard_utilities.jl | 4 - 4 files changed, 268 insertions(+), 18 deletions(-) diff --git a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl index d9a64582f..aed68e657 100644 --- a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl @@ -3,11 +3,22 @@ using Dates using JLD2 using Distributions using LinearAlgebra +using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz63.jl") + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_pushforward_samples = 1000 # pushforward ensemble size for output metrics +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -58,6 +69,31 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) +# Reconstruct deterministic l63 setup for pushforward (same seeds as calibrate_l63.jl) +nx = 3 +ny = 9 +t_step = 0.01 +T = 40.0 +lorenz_config_settings = LorenzConfig(t_step, T) +T_start = 30.0 +T_end = T +observation_config = ObservationConfig(T_start, T_end) +truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) +rng_seed_init = 11 +rng_i = MersenneTwister(rng_seed_init) +T_long = 1000.0 +x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) +x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) +x0 = x_spun_up[:, end] +covT = 2000.0 +cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) +ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) +n_output = ny + +# Load y from first valid results file (y is deterministic, same across all cases) +first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) +y = first_results["y"] + # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) n_evals_arr = fill(NaN, n_rng, n_ens) @@ -65,6 +101,12 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -95,12 +137,56 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(n_params)." + end + post_mean_arr[i, j, k, :] = pm post_cov_arr[i, j, k, :, :] = pc # (m - truth)' C^{-1} (m - truth) mahal_arr[i, j, k] = diff' * (C_reg \ diff) # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + # pushforward ensemble for output space + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" + for s in 1:n_pushforward_samples + output_samples_arr[i, j, k, s, :] = lorenz_forward( + EnsembleMemberConfig(push_con[:, s]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + + # output-space metrics (empirical distribution of pushforward samples vs truth output) + os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - y + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level (covariance near-degenerate). Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end end end @@ -109,11 +195,14 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -127,6 +216,10 @@ k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + # Data variables n_evals_v = defVar( @@ -177,5 +270,41 @@ posterior_logpdf_true_v_map_v = defVar( posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr +output_samples_v = defVar( + ds, "output_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); + fillvalue = NaN, +) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz63 output space (9-dim statistics: 3 means, 3 variances, 3 covariances)" +output_samples_v[:, :, :, :, :] = output_samples_arr + +output_mahal_v = defVar( + ds, "output_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - y)' Co^{-1} (om - y), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples, and y is the observations the posterior was fit to" +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar( + ds, "output_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution, with y the observations the posterior was fit to" +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) eigenvectors/values of Co ≈ a_o*I + Uo*Do*Uo'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add output_plr_mahalanobis_residual for the full Chisq($(n_output)) statistic." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o, where a_o = mean of the bottom $(n_output - n_lowrank_modes) eigenvalues of Co. Reference distribution: Chisq($(n_output - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add output_plr_mahalanobis_top for the full Chisq($(n_output)) statistic." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +output_coverage_v[:, :, :, :] = output_coverage_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index 3a877d678..229ebdc70 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -267,10 +267,6 @@ truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair. Use with forcing_samples to recompute marginal coverage at arbitrary quantile levels." truth_forcing_v[:, :, :] = truth_forcing_arr -truth_output_v = defVar(ds, "truth_output", Float64, ("output_dim",); fillvalue=NaN) -truth_output_v.attrib["description"] = "Observation vector y used as truth in output space. Constant across seeds and ensemble sizes. Use with output_samples to recompute marginal coverage at arbitrary quantile levels." -truth_output_v[:] = y - # Data variables n_evals_v = defVar( diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index d9a64582f..aed68e657 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -3,11 +3,22 @@ using Dates using JLD2 using Distributions using LinearAlgebra +using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") +include("Lorenz63.jl") + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_pushforward_samples = 1000 # pushforward ensemble size for output metrics +n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) # Step 1, Read and extract the available experiments @@ -58,6 +69,31 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) +# Reconstruct deterministic l63 setup for pushforward (same seeds as calibrate_l63.jl) +nx = 3 +ny = 9 +t_step = 0.01 +T = 40.0 +lorenz_config_settings = LorenzConfig(t_step, T) +T_start = 30.0 +T_end = T +observation_config = ObservationConfig(T_start, T_end) +truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) +rng_seed_init = 11 +rng_i = MersenneTwister(rng_seed_init) +T_long = 1000.0 +x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) +x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) +x0 = x_spun_up[:, end] +covT = 2000.0 +cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) +ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) +n_output = ny + +# Load y from first valid results file (y is deterministic, same across all cases) +first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) +y = first_results["y"] + # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) n_evals_arr = fill(NaN, n_rng, n_ens) @@ -65,6 +101,12 @@ post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) mahal_arr = fill(NaN, n_rng, n_ens, n_k) logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) for (N_ens, rng_idx) in valid_file_items post_fn = posterior_filename(cfg, N_ens, rng_idx) @@ -95,12 +137,56 @@ for (N_ens, rng_idx) in valid_file_items pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] diff = pm - truth_params + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(n_params)." + end + post_mean_arr[i, j, k, :] = pm post_cov_arr[i, j, k, :, :] = pc # (m - truth)' C^{-1} (m - truth) mahal_arr[i, j, k] = diff' * (C_reg \ diff) # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + # pushforward ensemble for output space + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" + for s in 1:n_pushforward_samples + output_samples_arr[i, j, k, s, :] = lorenz_forward( + EnsembleMemberConfig(push_con[:, s]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + + # output-space metrics (empirical distribution of pushforward samples vs truth output) + os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - y + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level (covariance near-degenerate). Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end end end @@ -109,11 +195,14 @@ end ds = NCDataset(nc_save_filename, "c") -defDim(ds, "random_seed", n_rng) -defDim(ds, "ensemble_size", n_ens) -defDim(ds, "k_iter", n_k) -defDim(ds, "param_dim", n_params) -defDim(ds, "param_dim_2", n_params) +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) # Coordinate variables rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) @@ -127,6 +216,10 @@ k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" k_var[:] = collect(1:n_k) +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + # Data variables n_evals_v = defVar( @@ -177,5 +270,41 @@ posterior_logpdf_true_v_map_v = defVar( posterior_logpdf_true_v_map_v.attrib["description"] = "P = posterior_logpdf(truth_param) - posterior_logpdf(mode(posterior)). Ideally -2P Ideally M ∈ quantile(Chisq(dims), [0.05,0.95]) 90% of the time. It asks `How much less probable is the true value compared to the MAP`. For gaussian -2P = M, if M > -2P => heavy tails" posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr +output_samples_v = defVar( + ds, "output_samples", Float64, + ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); + fillvalue = NaN, +) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples mapped to Lorenz63 output space (9-dim statistics: 3 means, 3 variances, 3 covariances)" +output_samples_v[:, :, :, :, :] = output_samples_arr + +output_mahal_v = defVar( + ds, "output_mahalanobis", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om - y)' Co^{-1} (om - y), where om and Co are the empirical mean/cov of the $(n_pushforward_samples) posterior pushforward output samples, and y is the observations the posterior was fit to" +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar( + ds, "output_logpdf_true_v_map", Float64, + ("random_seed", "ensemble_size", "k_iter"); + fillvalue = NaN, +) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode). Analogue of posterior_logpdf_true_v_map but for the pushforward output distribution, with y the observations the posterior was fit to" +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (top term) in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) eigenvectors/values of Co ≈ a_o*I + Uo*Do*Uo'. Reference distribution: Chisq($(n_lowrank_modes)). Captures miscalibration in the k principal directions. Add output_plr_mahalanobis_residual for the full Chisq($(n_output)) statistic." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "Perturbed-low-rank Mahalanobis (residual term) in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o, where a_o = mean of the bottom $(n_output - n_lowrank_modes) eigenvalues of Co. Reference distribution: Chisq($(n_output - n_lowrank_modes)). Captures miscalibration in the noise-floor directions. Add output_plr_mahalanobis_top for the full Chisq($(n_output)) statistic." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples. Expected ≈ p for calibrated marginals." +output_coverage_v[:, :, :, :] = output_coverage_arr + close(ds) @info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 0ac4bc09e..55498525f 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -265,10 +265,6 @@ truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair. Use with forcing_samples to recompute marginal coverage at arbitrary quantile levels." truth_forcing_v[:, :, :] = truth_forcing_arr -truth_output_v = defVar(ds, "truth_output", Float64, ("output_dim",); fillvalue=NaN) -truth_output_v.attrib["description"] = "Observation vector y used as truth in output space. Constant across seeds and ensemble sizes. Use with output_samples to recompute marginal coverage at arbitrary quantile levels." -truth_output_v[:] = y - # Data variables n_evals_v = defVar( From 08e00ff2b5b5a1eda02df4d87dcb0552ba588bd1 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 10 Jun 2026 15:02:17 -0700 Subject: [PATCH 71/86] unify l63 prelims --- examples/EKIRace/hpc-variant/calibrate_l63.jl | 20 ++ .../compute_leaderboard_metrics.jl | 187 ++++++++++++++---- .../l63_exp_to_leaderboard_utilities.jl | 51 +++-- 3 files changed, 202 insertions(+), 56 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl index 5c8c2fe9f..9edcc1b41 100644 --- a/examples/EKIRace/hpc-variant/calibrate_l63.jl +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -248,6 +248,26 @@ function main() setup = build_setup(cfg) write_priors(cfg, setup, output_dir) + prelim_file = joinpath(output_dir, "l63_computed_preliminaries.jld2") + if !isfile(prelim_file) + prelim_tmp = splitext(prelim_file)[1] * ".tmp.$(getpid()).jld2" + JLD2.save( + prelim_tmp, + "x0", setup.x0, "y", setup.y, + "lorenz_config_settings", setup.lorenz_config_settings, + "observation_config", setup.observation_config, + "ic_cov_sqrt", setup.ic_cov_sqrt, + "R", setup.R, "R_inv_var", setup.R_inv_var, + ) + try + mv(prelim_tmp, prelim_file) + @info "Saved computed quantities to $(prelim_file)" + catch + rm(prelim_tmp; force = true) + @info "Prelim file already written by another task; discarding duplicate" + end + end + tasks = flat_tasks(cfg) idx = task_index_from_args() diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index a11902a22..005c089b2 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -4,12 +4,14 @@ using Statistics using Printf using DataFrames using Dates +using JLD2 +using LinearAlgebra calib_date = Date("2026-06-03", "yyyy-mm-dd") indir = joinpath("output","from-hpc_$(calib_date)") -#filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" +filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" #filename = "ces-eki-dmc_l96_ensemble_results_$(calib_date).nc" -filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_$(calib_date).nc" +#filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_$(calib_date).nc" #filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_$(calib_date).nc" filename = joinpath(indir, filename) @@ -19,8 +21,17 @@ filename = joinpath(indir, filename) #################### Metric parameters ################################### ########################################################################### -calibration_check_quantiles = [0.2, 0.5, 0.8] # quantile levels for calibration check scores and coverage -n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +calibration_check_quantiles = [0.15, 0.5, 0.85] # quantile levels for calibration check scores and coverage +n_lowrank_modes = 2 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +R_variance_retain = 0.99 # fraction of R variance to retain for R-whitened PCA coverage + +# Optional: path to the prelim JLD2 written by calibrate_l96.jl. +# Enables coverage recomputation at any quantile and R-whitened PCA coverage. +# e.g. "output/l96_computed_preliminaries_const-force.jld2" +prelim_jld2_file = "output/from-hpc_2026-06-03/l96_computed_preliminaries_flux-force.jld2" + +prelim_filename = joinpath(indir, prelim_jld2_file) +@info "using precomputed quantities from $(prelim_filename)" ########################################################################### #################### Display filters #################################### @@ -43,7 +54,7 @@ n_lowrank_modes = 5 # top EOF modes for perturbed-low # display_ens_sizes — which ensemble-size values (integers) to show in tables. display_spaces = [:output] # e.g. [:param, :forcing] -display_metrics = [:coverage] # e.g. [:mahalanobis, :plr_mahalanobis] +display_metrics = [:coverage, :mahalanobis] # e.g. [:mahalanobis, :plr_mahalanobis] display_ens_sizes = nothing # e.g. [10, 50] show_metric(m) = isnothing(display_metrics) || m in display_metrics @@ -61,13 +72,28 @@ lp = ncd[:posterior_logpdf_true_v_map] ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end +# Load truth and observation-noise covariance from the prelim JLD2 (calibrate_l96.jl output). +# Falls back to truth_output stored in the netcdf (legacy) if JLD2 is not configured. +if !isnothing(prelim_jld2_file) && isfile(prelim_jld2_file) + _pd = JLD2.load(prelim_jld2_file) + y_truth = Float64.(_pd["y"]) + R_obs = Float64.(_pd["R"]) + @info "Loaded y and R from $(prelim_jld2_file)" +elseif haskey(ncd, "truth_output") + y_truth = Float64.(ncd["truth_output"][:]) + R_obs = nothing +else + y_truth = nothing + R_obs = nothing +end + # Discover what is actually present in this file avail_spaces = Symbol[:param] avail_metrics = Symbol[:mahalanobis, :logratio] -(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage")) && push!(avail_spaces, :forcing) -(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage")) && push!(avail_spaces, :output) +(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage") || (haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing"))) && push!(avail_spaces, :forcing) +(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage") || (haskey(ncd, "output_samples") && !isnothing(y_truth))) && push!(avail_spaces, :output) (haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "output_plr_mahalanobis_top")) && push!(avail_metrics, :plr_mahalanobis) -(haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) +((haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing")) || (haskey(ncd, "output_samples") && !isnothing(y_truth)) || haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) println("="^72) println("File: $(filename)") @@ -431,52 +457,143 @@ end ########################################################################### #################### Marginal coverage fractions ######################### ########################################################################### +# Coverage is computed on-the-fly from raw pushforward samples and truth, +# allowing any quantile level. When R_obs is loaded from the prelim JLD2, +# an additional R-whitened (PCA) coverage is computed: each coordinate is +# decorrelated by the eigenvectors of R and scaled by the inverse sqrt of +# the corresponding eigenvalue, making dimensions more independent. if show_metric(:coverage) - for (space_label, space_sym, cov_sym) in [ - ("Forcing", :forcing, :forcing_coverage), - ("Output", :output, :output_coverage), - ] - (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs_all = Float64.(ncd["coverage_quantile"][:]) - qi_keep = findall(qp -> any(isapprox(qp, p, atol=1e-8) for p in calibration_check_quantiles), cov_qs_all) - cov_qs = cov_qs_all[qi_keep] + # ── Output-space coverage ────────────────────────────────────────────── + if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) + + n_out = ncd.dim["output_dim"] + os_all = coalesce.(Array(ncd["output_samples"]), NaN) # (n_rng, n_ens_size, n_k_iter, n_ps, n_out) + actual_out_dim = size(os_all, 5) + actual_out_dim != n_out && @warn "output_samples trailing dim ($actual_out_dim) ≠ output_dim ($n_out); raw coverage uses actual dim, R-whitened coverage disabled" + + do_whiten = !isnothing(R_obs) + if do_whiten + eig_R = eigen(Symmetric(R_obs)) + ord = sortperm(eig_R.values; rev=true) # descending eigenvalue order + λ_all = eig_R.values[ord] + V_all = eig_R.vectors[:, ord] + cum_var = cumsum(λ_all) ./ sum(λ_all) + k_R = something(findfirst(>=(R_variance_retain), cum_var), length(λ_all)) + V_R = V_all[:, 1:k_R] + λ_R = λ_all[1:k_R] + yw = V_R' * y_truth ./ sqrt.(λ_R) + println(@sprintf("\nR eigenvalue truncation: %d / %d modes retained (threshold %.1f%% → %.4f%% variance)", + k_R, n_out, 100*R_variance_retain, 100*cum_var[k_R])) + end + + for (label, do_white) in [("raw", false), ("R-whitened PCA", true)] + (do_white && !do_whiten) && continue + + local cov_arr = fill(NaN, n_rng, n_ens_size, n_k_iter, length(calibration_check_quantiles)) + for ri in 1:n_rng, ei in 1:n_ens_size, ki in 1:n_k_iter + os = os_all[ri, ei, ki, :, :] # (n_ps, n_out) + all(isnan.(os)) && continue + if do_white + size(os, 2) != size(V_R, 1) && continue # dim mismatch; leave cov_arr[ri,ei,ki,:] as NaN + sw = Matrix((V_R' * Matrix(os')) ./ sqrt.(λ_R))' # (n_ps, k_R) + truth_vec = yw + n_cov_dim = k_R + else + sw = os + truth_vec = y_truth + n_cov_dim = actual_out_dim + end + for (qi, qp) in enumerate(calibration_check_quantiles) + cov_arr[ri, ei, ki, qi] = mean(truth_vec[d] <= quantile(sw[:, d], qp) for d in 1:n_cov_dim) + end + end + + println("\n" * "="^72) + println("Output-space marginal coverage fractions [$(label)]") + if do_white + println(" x̃_d = (Vᵀx)_d / √λ_d (R = VΛVᵀ, top $(k_R)/$(n_out) modes, $(round(Int, 100*R_variance_retain))% var threshold)") + println(" Coverage evaluated over effective dimension = $(k_R) (full output dim = $(n_out))") + end + println(" (dims as trials; expected mean_coverage ≈ q_prob for calibrated posterior)") + println("="^72) + println("Pooled mean coverage across all experiments:") + for (qi, qp) in enumerate(calibration_check_quantiles) + pool = collect(filter(!isnan, vec(cov_arr[:, :, :, qi]))) + isempty(pool) && continue + println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + local df_cov = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mean_coverage=Union{Float64,Missing}[], std_coverage=Union{Float64,Missing}[], n_valid=Int[], + ) + for (qi, qp) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = filter(!isnan, vec(cov_arr[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + std_coverage = length(vals) > 1 ? std(vals) : missing, + n_valid = length(vals))) + end + + println("\nOutput-space coverage by ens_size and k_iter:") + for ens in filter(show_ens, sort(unique(df_cov.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) + println() + end + end + end + end + + # ── Forcing-space coverage ───────────────────────────────────────────── + if show_space(:forcing) && haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing") + + n_for = ncd.dim["forcing_dim"] + fs_all = coalesce.(Array(ncd["forcing_samples"]), NaN) # (n_rng, n_ens_size, n_k_iter, n_ps, n_for) + tf_all = coalesce.(Array(ncd["truth_forcing"]), NaN) # (n_rng, n_ens_size, n_for) + + cov_arr = fill(NaN, n_rng, n_ens_size, n_k_iter, length(calibration_check_quantiles)) + for ri in 1:n_rng, ei in 1:n_ens_size, ki in 1:n_k_iter + fs = fs_all[ri, ei, ki, :, :] + truth_f = tf_all[ri, ei, :] + all(isnan.(fs)) && continue + for (qi, qp) in enumerate(calibration_check_quantiles) + cov_arr[ri, ei, ki, qi] = mean(truth_f[d] <= quantile(fs[:, d], qp) for d in 1:n_for) + end + end println("\n" * "="^72) - println("$(space_label)-space marginal coverage fractions") - println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("Forcing-space marginal coverage fractions [raw]") + println(" (dims as trials; expected mean_coverage ≈ q_prob for calibrated posterior)") println("="^72) println("Pooled mean coverage across all experiments:") - for (qi_local, qi_file) in enumerate(qi_keep) - qp = cov_qs[qi_local] - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi_file])))) + for (qi, qp) in enumerate(calibration_check_quantiles) + pool = collect(filter(!isnan, vec(cov_arr[:, :, :, qi]))) + isempty(pool) && continue println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) end df_cov = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mean_coverage = Union{Float64,Missing}[], - std_coverage = Union{Float64,Missing}[], - n_valid = Int[], + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mean_coverage=Union{Float64,Missing}[], std_coverage=Union{Float64,Missing}[], n_valid=Int[], ) - for (qi_local, qi_file) in enumerate(qi_keep), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - qp = cov_qs[qi_local] - vals = collect(skipmissing(cov[:, ei, ki, qi_file])) + for (qi, qp) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = filter(!isnan, vec(cov_arr[:, ei, ki, qi])) push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, mean_coverage = isempty(vals) ? missing : mean(vals), std_coverage = length(vals) > 1 ? std(vals) : missing, n_valid = length(vals))) end - println("\n$(space_label)-space coverage by ens_size and k_iter:") - println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + println("\nForcing-space coverage by ens_size and k_iter:") for ens in filter(show_ens, sort(unique(df_cov.ens_size))) println("\n=== ens_size = $(ens) ===") - for (qi, qp) in enumerate(cov_qs) + for (qi, qp) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) println(" q = $(qp) (expected ≈ $(qp))") show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) @@ -485,5 +602,3 @@ if show_metric(:coverage) end end end - -@info calibration_check_quantiles diff --git a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl index aed68e657..a76dea745 100644 --- a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl @@ -69,30 +69,41 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) -# Reconstruct deterministic l63 setup for pushforward (same seeds as calibrate_l63.jl) nx = 3 ny = 9 -t_step = 0.01 -T = 40.0 -lorenz_config_settings = LorenzConfig(t_step, T) -T_start = 30.0 -T_end = T -observation_config = ObservationConfig(T_start, T_end) -truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) -rng_seed_init = 11 -rng_i = MersenneTwister(rng_seed_init) -T_long = 1000.0 -x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) -x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) -x0 = x_spun_up[:, end] -covT = 2000.0 -cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) -ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) n_output = ny -# Load y from first valid results file (y is deterministic, same across all cases) -first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) -y = first_results["y"] +# Load prelim quantities from calibrate_l63.jl output; fall back to recomputation if absent. +prelim_file = joinpath(homedir, "output", "l63_computed_preliminaries.jld2") +if isfile(prelim_file) + prelim_data = JLD2.load(prelim_file) + x0 = prelim_data["x0"] + y = prelim_data["y"] + ic_cov_sqrt = prelim_data["ic_cov_sqrt"] + lorenz_config_settings = prelim_data["lorenz_config_settings"] + observation_config = prelim_data["observation_config"] + @info "Loaded L63 preliminaries from $(prelim_file)" +else + @warn "Prelim file not found at $(prelim_file); recomputing deterministically (run calibrate_l63.jl first)" + t_step = 0.01 + T = 40.0 + lorenz_config_settings = LorenzConfig(t_step, T) + T_start = 30.0 + T_end = T + observation_config = ObservationConfig(T_start, T_end) + truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) + rng_seed_init = 11 + rng_i = MersenneTwister(rng_seed_init) + T_long = 1000.0 + x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) + x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) + x0 = x_spun_up[:, end] + covT = 2000.0 + cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) + ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) + first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) + y = first_results["y"] +end # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) From bb1f2a2396c61dd7bc264c2f356e60b877ab4763 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 10 Jun 2026 15:22:13 -0700 Subject: [PATCH 72/86] add comment to run the leaderboard for l63 --- examples/EKIRace/hpc-variant/submit_l63.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 1e9c9ac17..30b137104 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -45,3 +45,5 @@ LB_JID=$(sbatch --parsable \ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" + +# sbatch -A esm --job-name="leaderboard_l63" --export=ALL, exp_to_leaderboard.sbatch From 32e595d014919b3c20a1cd7eef4388b53ae34af8 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 11 Jun 2026 15:39:36 -0700 Subject: [PATCH 73/86] add more easily changeable n_ens and n_repeat, consistent for the sake of sbatch --- .../EKIRace/hpc-variant/experiment_config.jl | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl index d1f9d354f..388b8674e 100644 --- a/examples/EKIRace/hpc-variant/experiment_config.jl +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -10,7 +10,7 @@ EXPERIMENT = experiments[4] # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). # PIN this before submitting an array job so all tasks use the same directory. -#calibrate_date = Date("2026-05-29", "yyyy-mm-dd") +#calibrate_date = Date("2026-06-03", "yyyy-mm-dd") calibrate_date = today() ######################################################################## @@ -47,15 +47,19 @@ forcing_cases_key = Dict( # N_ens_sizes: Vector of experiments # n_repeats: number of rng seeds + function experiment_config(case::Symbol) + n_ens_step = 8 ## MUST CHANGE SBATCH ARRAY SIZE TO ((N_ENS_STEP + 1) * N_REPEATS) + n_repeats = 20 if case == :l63 + ens_step = 2 return ( model = "l63", force_case = nothing, - N_ens_sizes = [10, 25, 40], + N_ens_sizes = collect(4:ens_step:4+n_ens_step*ens_step), N_iter = 20, terminate_at = 2.0, # DataMisfitController end time - n_repeats = 20, + n_repeats = n_repeats, max_iter = 10, retain_var = 0.99, n_features = 100, @@ -63,13 +67,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_const + ens_step_const = 2 return ( model = "l96", force_case = "const-force", - N_ens_sizes = [5, 15, 30], + N_ens_sizes = collect(4:ens_step_const:4+n_ens_step*ens_step_const), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, @@ -77,13 +82,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_vec + ens_step_vec = 5 return ( model = "l96", force_case = "vec-force", - N_ens_sizes = [50, 75, 100], + N_ens_sizes = collect(40:ens_step_vec:40+n_ens_step*ens_step_vec), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, @@ -91,13 +97,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_flux + ens_step_flux = 5 return ( model = "l96", force_case = "flux-force", - N_ens_sizes = [50, 75, 100], + N_ens_sizes = collect(30:ens_step_flux:30+n_ens_step*ens_step_flux), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, From 8fba8c55a1916d6549f09411eb62e216d7d206be Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 11 Jun 2026 15:41:18 -0700 Subject: [PATCH 74/86] update array consistent with exp config --- examples/EKIRace/hpc-variant/calibrate_array.sbatch | 2 +- examples/EKIRace/hpc-variant/emulate_sample_array.sbatch | 2 +- .../EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 7ca5ae8c9..9a33a2b56 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -20,7 +20,7 @@ #SBATCH --job-name=calib #SBATCH --output=output/slurm/calib_%A_%a.out #SBATCH --error=output/slurm/calib_%A_%a.err -#SBATCH --array=1-60%100 +#SBATCH --array=1-180%100 #SBATCH --time=02:00:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index e4dd99503..156b5a07b 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -18,7 +18,7 @@ #SBATCH --job-name=emusample #SBATCH --output=output/slurm/emusample_%A_%a.out #SBATCH --error=output/slurm/emusample_%A_%a.err -#SBATCH --array=1-60%100 +#SBATCH --array=1-180%100 #SBATCH --time=12:00:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index 27a2db488..27239931b 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -11,7 +11,7 @@ #SBATCH --job-name=post_diag #SBATCH --output=output/slurm/post_diag_%A_%a.out #SBATCH --error=output/slurm/post_diag_%A_%a.err -#SBATCH --array=1-60%100 +#SBATCH --array=1-180%100 #SBATCH --time=00:30:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 From f22032df7e523377cf50a276cb5691d162e739b5 Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 12 Jun 2026 16:32:47 -0700 Subject: [PATCH 75/86] add pushforward from posterior as separate parallel script --- .claude/settings.json | 10 ++ .../compute_leaderboard_metrics.jl | 48 +++++--- .../l63_exp_to_leaderboard_utilities.jl | 76 ++++-------- .../l96_exp_to_leaderboard_utilities.jl | 78 +++++------- .../posterior_diagnostic_plots_l96.jl | 29 ++--- .../pushforward_from_posterior.sbatch | 39 ++++++ .../pushforward_from_posterior_l63.jl | 86 ++++++++++++++ .../pushforward_from_posterior_l96.jl | 111 ++++++++++++++++++ examples/EKIRace/hpc-variant/submit_l63.sh | 12 +- .../EKIRace/hpc-variant/submit_l96_const.sh | 17 ++- .../EKIRace/hpc-variant/submit_l96_flux.sh | 17 ++- .../EKIRace/hpc-variant/submit_l96_vec.sh | 17 ++- .../l63_exp_to_leaderboard_utilities.jl | 67 ++++------- .../l96_exp_to_leaderboard_utilities.jl | 78 +++++------- .../EKIRace/posterior_diagnostic_plots_l96.jl | 30 ++--- 15 files changed, 468 insertions(+), 247 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch create mode 100644 examples/EKIRace/hpc-variant/pushforward_from_posterior_l63.jl create mode 100644 examples/EKIRace/hpc-variant/pushforward_from_posterior_l96.jl diff --git a/.claude/settings.json b/.claude/settings.json index 4990e57e8..e8fcb47ec 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -7,6 +7,16 @@ "FileWrite(src/*, test/*, docs/*, examples/*, ai/*)", "Edit(src/*)", "Write(src/*)", + "Grep", + "Glob", + "Bash(find *)", + "Bash(ls *)", + "Bash(tree *)", + "Bash(wc *)", + "Bash(head *)", + "Bash(tail *)", + "Bash(stat *)", + "Bash(grep *)", "Bash(git log*)", "Bash(git diff*)", "Bash(git status*)", diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index 005c089b2..a6ab370e3 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -7,15 +7,34 @@ using Dates using JLD2 using LinearAlgebra -calib_date = Date("2026-06-03", "yyyy-mm-dd") +calib_date = Date("2026-06-11", "yyyy-mm-dd") indir = joinpath("output","from-hpc_$(calib_date)") -filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" -#filename = "ces-eki-dmc_l96_ensemble_results_$(calib_date).nc" -#filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_$(calib_date).nc" -#filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_$(calib_date).nc" + +experiment_list = [:l63, :l96_const, :l96_vec, :l96_flux ] +experiment = experiment_list[1] +if experiment == :l63 + filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" + prelim_jld2_file = "l63_computed_preliminaries.jld2" +elseif experiment == :l96_const + filename = "ces-eki-dmc_l96_ensemble_results_$(calib_date).nc" + prelim_jld2_file = "l96_computed_preliminaries_const-force.jld2" +elseif experiment == :l96_vec + filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_$(calib_date).nc" + prelim_jld2_file = "l96_computed_preliminaries_vec-force.jld2" +elseif experiment == :l96_flux + filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_$(calib_date).nc" + prelim_jld2_file = "l96_computed_preliminaries_flux-force.jld2" +else + throw(ArgumentError("Expected Experiment from $(experiment_list). Got $(experiment).")) +end filename = joinpath(indir, filename) -@info "computing leaderboard metrics from $(filename)" +prelim_filename = joinpath(indir, prelim_jld2_file) +# Optional: path to the prelim JLD2 written by calibrate_l96.jl. +# Enables coverage recomputation at any quantile and R-whitened PCA coverage. +# e.g. "output/l96_computed_preliminaries_const-force.jld2" + + ########################################################################### #################### Metric parameters ################################### @@ -25,13 +44,6 @@ calibration_check_quantiles = [0.15, 0.5, 0.85] # quantile levels for calibratio n_lowrank_modes = 2 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis R_variance_retain = 0.99 # fraction of R variance to retain for R-whitened PCA coverage -# Optional: path to the prelim JLD2 written by calibrate_l96.jl. -# Enables coverage recomputation at any quantile and R-whitened PCA coverage. -# e.g. "output/l96_computed_preliminaries_const-force.jld2" -prelim_jld2_file = "output/from-hpc_2026-06-03/l96_computed_preliminaries_flux-force.jld2" - -prelim_filename = joinpath(indir, prelim_jld2_file) -@info "using precomputed quantities from $(prelim_filename)" ########################################################################### #################### Display filters #################################### @@ -65,6 +77,7 @@ show_ens(e) = isnothing(display_ens_sizes) || e in display_ens_sizes #################### Load file and report available options ############## ########################################################################### +@info "computing leaderboard metrics from $(filename)" ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] @@ -74,8 +87,9 @@ kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end # Load truth and observation-noise covariance from the prelim JLD2 (calibrate_l96.jl output). # Falls back to truth_output stored in the netcdf (legacy) if JLD2 is not configured. -if !isnothing(prelim_jld2_file) && isfile(prelim_jld2_file) - _pd = JLD2.load(prelim_jld2_file) +if !isnothing(prelim_jld2_file) && isfile(prelim_filename) + @info "using precomputed quantities from $(prelim_filename)" + _pd = JLD2.load(prelim_filename) y_truth = Float64.(_pd["y"]) R_obs = Float64.(_pd["R"]) @info "Loaded y and R from $(prelim_jld2_file)" @@ -83,6 +97,7 @@ elseif haskey(ncd, "truth_output") y_truth = Float64.(ncd["truth_output"][:]) R_obs = nothing else + @error "precomputed quantities NOT FOUND at $(prelim_filename)" y_truth = nothing R_obs = nothing end @@ -466,6 +481,9 @@ end if show_metric(:coverage) # ── Output-space coverage ────────────────────────────────────────────── + @info show_space(:output) + @info haskey(ncd, "output_samples") + @info !isnothing(y_truth) if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) n_out = ncd.dim["output_dim"] diff --git a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl index a76dea745..db204bd06 100644 --- a/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl @@ -3,19 +3,16 @@ using Dates using JLD2 using Distributions using LinearAlgebra -using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") -include("Lorenz63.jl") ########################################################################### #################### Metric parameters ################################### ########################################################################### -n_pushforward_samples = 1000 # pushforward ensemble size for output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) @@ -57,9 +54,16 @@ end ### Load data -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +# Load first valid file to determine parameter dimension, pushforward dimensions, and max k +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run push_forward_from_posterior.sbatch first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) n_k = maximum( maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) @@ -69,41 +73,10 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) -nx = 3 -ny = 9 -n_output = ny - -# Load prelim quantities from calibrate_l63.jl output; fall back to recomputation if absent. +# y is needed for output-space metrics; load from the shared prelim file prelim_file = joinpath(homedir, "output", "l63_computed_preliminaries.jld2") -if isfile(prelim_file) - prelim_data = JLD2.load(prelim_file) - x0 = prelim_data["x0"] - y = prelim_data["y"] - ic_cov_sqrt = prelim_data["ic_cov_sqrt"] - lorenz_config_settings = prelim_data["lorenz_config_settings"] - observation_config = prelim_data["observation_config"] - @info "Loaded L63 preliminaries from $(prelim_file)" -else - @warn "Prelim file not found at $(prelim_file); recomputing deterministically (run calibrate_l63.jl first)" - t_step = 0.01 - T = 40.0 - lorenz_config_settings = LorenzConfig(t_step, T) - T_start = 30.0 - T_end = T - observation_config = ObservationConfig(T_start, T_end) - truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) - rng_seed_init = 11 - rng_i = MersenneTwister(rng_seed_init) - T_long = 1000.0 - x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) - x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) - x0 = x_spun_up[:, end] - covT = 2000.0 - cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) - ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) - first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) - y = first_results["y"] -end +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l63.jl first.") +y = JLD2.load(prelim_file, "y") # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) @@ -124,10 +97,18 @@ for (N_ens, rng_idx) in valid_file_items @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run push_forward_from_posterior.sbatch first." + continue + end + posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) @@ -161,18 +142,9 @@ for (N_ens, rng_idx) in valid_file_items # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # pushforward ensemble for output space - push_unc = sample(post_dist, n_pushforward_samples) - push_con = transform_unconstrained_to_constrained(post_dist, push_unc) - @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" - for s in 1:n_pushforward_samples - output_samples_arr[i, j, k, s, :] = lorenz_forward( - EnsembleMemberConfig(push_con[:, s]), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) - end + # load precomputed pushforward samples (produced by pushforward_from_posterior_l63.jl) + ki = findfirst(==(k), pf_k_values) + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] # output-space metrics (empirical distribution of pushforward samples vs truth output) os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) diff --git a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl index 229ebdc70..c916040ac 100644 --- a/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -3,14 +3,11 @@ using Dates using JLD2 using Distributions using LinearAlgebra -using Flux -using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") -include("Lorenz96.jl") EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg @@ -18,7 +15,6 @@ EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg #################### Metric parameters ################################### ########################################################################### -n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) @@ -62,9 +58,17 @@ end ### Load data -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +# Load first valid file to determine parameter dimension, pushforward dimensions, and max k +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run push_forward_from_posterior.sbatch first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) +n_forcing = size(first_loaded["pushforward_forcing_samples"], 2) n_k = maximum( maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) @@ -74,28 +78,10 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) -# Load Lorenz preliminaries and truth structure for pushforward +# y is needed for output-space metrics; load once from the shared prelim file prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") -prelim_data = JLD2.load(prelim_file) -x0 = prelim_data["x0"] -nx = length(x0) -ic_cov_sqrt = prelim_data["ic_cov_sqrt"] -lorenz_config_settings = prelim_data["lorenz_config_settings"] -observation_config = prelim_data["observation_config"] -y = prelim_data["y"] -n_output = 2 * nx - -first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) -(truth_phi, _) = first_calib["truth_params_structure"] -(phi_structure, sample_range) = if force_case == "const-force" - (nothing, nothing) -elseif force_case == "vec-force" - (nothing, nothing) -elseif force_case == "flux-force" - (truth_phi.model, truth_phi.sample_range) -end -n_forcing = force_case == "flux-force" ? length(sample_range) : nx +y = JLD2.load(prelim_file, "y") # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) @@ -123,10 +109,19 @@ for (N_ens, rng_idx) in valid_file_items @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run push_forward_from_posterior.sbatch first." + continue + end + posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] - truth_params = loaded["truth_params"] - truth_params_constrained = loaded["truth_params_constrained"] + truth_params = loaded["truth_params"] + + pf_forcing = loaded["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + truth_forcing_vec = loaded["truth_forcing"] # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -135,11 +130,8 @@ for (N_ens, rng_idx) in valid_file_items i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - true_param_arr[i, j, :] = truth_params - n_evals_arr[i, j] = conv_alg_iters * N_ens - - emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) - truth_forcing_vec = forcing(emc_truth, x0) + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens truth_forcing_arr[i, j, :] = truth_forcing_vec for k in k_values @@ -165,20 +157,10 @@ for (N_ens, rng_idx) in valid_file_items # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # pushforward ensemble for forcing and output space - push_unc = sample(post_dist, n_pushforward_samples) - push_con = transform_unconstrained_to_constrained(post_dist, push_unc) - @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" - for s in 1:n_pushforward_samples - emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) - forcing_samples_arr[i, j, k, s, :] = forcing(emc, x0) - output_samples_arr[i, j, k, s, :] = lorenz_forward( - emc, - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) - end + # load precomputed pushforward samples (produced by pushforward_from_posterior_l96.jl) + ki = findfirst(==(k), pf_k_values) + forcing_samples_arr[i, j, k, :, :] = pf_forcing[:, :, ki] + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] # forcing-space metrics (empirical distribution of pushforward samples vs truth forcing) fs = forcing_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_forcing) diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl index 4615b60a2..eb0a1a6d0 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -116,27 +116,28 @@ function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[ ]..., ) + if !haskey(loaded_p, "pushforward_output_samples") + error("Pushforward data not found in $(post_fn). Run push_forward_from_posterior.sbatch first.") + end + pf_output = loaded_p["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_forcing = loaded_p["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + pf_k_values = loaded_p["pushforward_k_values"] + for k in k_values post_dist = posteriors_by_k[k] - push_ensemble = sample(post_dist, n_samples_pushforward) - # note: push_ensemble is unconstrained + + # sample posterior for parameter-space analysis (no forward map needed) + push_ensemble = sample(post_dist, n_samples_pushforward) constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) - G_ens = hcat( - [ - lorenz_forward( - build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) for j in 1:n_samples_pushforward - ]..., - ) + # load precomputed posterior pushforward (produced by pushforward_from_posterior_l96.jl) + ki = findfirst(==(k), pf_k_values) + G_ens = pf_output[:, :, ki]' # (n_output, n_samples) + push_forcings = pf_forcing[:, :, ki]' # (n_forcing, n_samples) ny = length(y) n_par = length(truth_params_constrained) - push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) - param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings diff --git a/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch new file mode 100644 index 000000000..4935401ac --- /dev/null +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch @@ -0,0 +1,39 @@ +#!/bin/bash +# Array job: one task per (N_ens, rng_idx) pair. +# Each task loads the emulate_sample posterior, pushes n_pushforward_samples through +# the forward map for every k-iteration, and saves forcing + output samples back into +# the posterior JLD2 file. Downstream consumers (leaderboard utilities, diagnostic +# plots) load these stored arrays instead of rerunning the expensive forward map. +# +# L63 (180 tasks for default config): +# sbatch --dependency=afterok: push_forward_from_posterior.sbatch +# +# L96 (180 tasks per case): +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_flux push_forward_from_posterior.sbatch + +#SBATCH --job-name=pushfwd +#SBATCH --output=output/slurm/pushfwd_%A_%a.out +#SBATCH --error=output/slurm/pushfwd_%A_%a.err +#SBATCH --array=1-180%100 +#SBATCH --time=01:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +if [[ -z "${EXPERIMENT:-}" || "${EXPERIMENT}" == "l63" ]]; then + julia --project=. pushforward_from_posterior_l63.jl "${SLURM_ARRAY_TASK_ID}" +else + julia --project=. pushforward_from_posterior_l96.jl "${SLURM_ARRAY_TASK_ID}" +fi diff --git a/examples/EKIRace/hpc-variant/pushforward_from_posterior_l63.jl b/examples/EKIRace/hpc-variant/pushforward_from_posterior_l63.jl new file mode 100644 index 000000000..fc17c4da2 --- /dev/null +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior_l63.jl @@ -0,0 +1,86 @@ +using JLD2 +using Distributions +using LinearAlgebra +using Random +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz63.jl") + +const n_pushforward_samples = 1000 + +function pushforward_one(cfg, N_ens, rng_idx; method = method_cases[1]) + calib_dir = calib_directory(method, cfg) + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + post_fn = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + + if !isfile(post_fn) + @warn "No posterior file for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + loaded_p = JLD2.load(post_fn) + if haskey(loaded_p, "pushforward_output_samples") + @info "Pushforward already present in $(post_fn); skipping." + return + end + + prelim_file = joinpath(homedir, "output", "l63_computed_preliminaries.jld2") + isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l63.jl first.") + prelim = JLD2.load(prelim_file) + x0 = prelim["x0"] + nx = length(x0) + ic_cov_sqrt = prelim["ic_cov_sqrt"] + lorenz_config_settings = prelim["lorenz_config_settings"] + observation_config = prelim["observation_config"] + n_output = 9 # 3 means + 3 variances + 3 covariances + + posteriors_by_k = loaded_p["posteriors_by_k"] + k_values = loaded_p["k_values"] + n_k = length(k_values) + + output_arr = Array{Float64}(undef, n_pushforward_samples, n_output, n_k) + + for (ki, k) in enumerate(k_values) + post_dist = posteriors_by_k[k] + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz63 evals" + for s in 1:n_pushforward_samples + output_arr[s, :, ki] = lorenz_forward( + EnsembleMemberConfig(push_con[:, s]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + end + + JLD2.jldopen(post_fn, "r+") do f + f["pushforward_output_samples"] = output_arr # (n_samples, n_output, n_k) + f["pushforward_k_values"] = k_values + f["pushforward_n_samples"] = n_pushforward_samples + end + @info "Saved pushforward to $(post_fn)" +end + +function main() + cfg = experiment_config(:l63) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + pushforward_one(cfg, N_ens, rng_idx) + end + else + idx < 1 || idx > length(tasks) && error("Task index $(idx) out of range 1:$(length(tasks))") + N_ens, rng_idx = tasks[idx] + pushforward_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/pushforward_from_posterior_l96.jl b/examples/EKIRace/hpc-variant/pushforward_from_posterior_l96.jl new file mode 100644 index 000000000..0e6ec4430 --- /dev/null +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior_l96.jl @@ -0,0 +1,111 @@ +using JLD2 +using Distributions +using LinearAlgebra +using Random +using Flux +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz96.jl") + +const n_pushforward_samples = 1000 + +function pushforward_one(cfg, N_ens, rng_idx; method = method_cases[1]) + force_case = cfg.force_case + calib_dir = calib_directory(method, cfg) + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + post_fn = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + calib_fn = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + + if !isfile(post_fn) + @warn "No posterior file for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + loaded_p = JLD2.load(post_fn) + if haskey(loaded_p, "pushforward_output_samples") + @info "Pushforward already present in $(post_fn); skipping." + return + end + + prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") + isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") + prelim = JLD2.load(prelim_file) + x0 = prelim["x0"] + nx = length(x0) + ic_cov_sqrt = prelim["ic_cov_sqrt"] + lorenz_config_settings = prelim["lorenz_config_settings"] + observation_config = prelim["observation_config"] + n_output = 2 * nx + + calib_loaded = JLD2.load(calib_fn) + (truth_phi, _) = calib_loaded["truth_params_structure"] + (phi_structure, sample_range) = if force_case == "const-force" + (nothing, nothing) + elseif force_case == "vec-force" + (nothing, nothing) + elseif force_case == "flux-force" + (truth_phi.model, truth_phi.sample_range) + end + n_forcing = force_case == "flux-force" ? length(sample_range) : nx + + truth_params_constrained = loaded_p["truth_params_constrained"] + truth_emc = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) + truth_forcing_vec = forcing(truth_emc, x0) + + posteriors_by_k = loaded_p["posteriors_by_k"] + k_values = loaded_p["k_values"] + n_k = length(k_values) + + forcing_arr = Array{Float64}(undef, n_pushforward_samples, n_forcing, n_k) + output_arr = Array{Float64}(undef, n_pushforward_samples, n_output, n_k) + + for (ki, k) in enumerate(k_values) + post_dist = posteriors_by_k[k] + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz96 evals" + for s in 1:n_pushforward_samples + emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) + forcing_arr[s, :, ki] = forcing(emc, x0) + output_arr[s, :, ki] = lorenz_forward( + emc, + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + end + + JLD2.jldopen(post_fn, "r+") do f + f["pushforward_forcing_samples"] = forcing_arr # (n_samples, n_forcing, n_k) + f["pushforward_output_samples"] = output_arr # (n_samples, n_output, n_k) + f["pushforward_k_values"] = k_values + f["pushforward_n_samples"] = n_pushforward_samples + f["truth_forcing"] = truth_forcing_vec + end + @info "Saved pushforward to $(post_fn)" +end + +function main() + exp = l96_experiment() + @assert exp in (:l96_const, :l96_vec, :l96_flux) "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $exp)" + cfg = experiment_config(exp) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + pushforward_one(cfg, N_ens, rng_idx) + end + else + idx < 1 || idx > length(tasks) && error("Task index $(idx) out of range 1:$(length(tasks))") + N_ens, rng_idx = tasks[idx] + pushforward_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 30b137104..4477c7bbf 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -36,11 +36,19 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting exp_to_leaderboard (L63, after ${EMU_JID}) ===" +echo "=== Submitting pushforward_from_posterior (L63, after ${EMU_JID}) ===" +PUSHFWD_JID=$(sbatch --parsable \ + -A esm \ + --job-name="pushfwd_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + pushforward_from_posterior.sbatch) +echo " pushforward_from_posterior job ID: ${PUSHFWD_JID}" + +echo "=== Submitting exp_to_leaderboard (L63, after ${PUSHFWD_JID}) ===" LB_JID=$(sbatch --parsable \ -A esm \ --job-name="leaderboard_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ exp_to_leaderboard.sbatch) echo " exp_to_leaderboard job ID: ${LB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 36f8fa455..272d0a379 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -46,20 +46,29 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting posterior_diagnostic_plots (L96 const-force, after ${EMU_JID}) ===" +echo "=== Submitting pushforward_from_posterior (L96 const-force, after ${EMU_JID}) ===" +PUSHFWD_JID=$(sbatch --parsable \ + -A esm \ + --job-name="pushfwd_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + pushforward_from_posterior.sbatch) +echo " pushforward_from_posterior job ID: ${PUSHFWD_JID}" + +echo "=== Submitting posterior_diagnostic_plots (L96 const-force, after ${PUSHFWD_JID}) ===" POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_const \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" -echo "=== Submitting exp_to_leaderboard (L96 const-force, after ${EMU_JID}) ===" +echo "=== Submitting exp_to_leaderboard (L96 const-force, after ${PUSHFWD_JID}) ===" LB_JID=$(sbatch --parsable \ -A esm \ --job-name="leaderboard_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_const \ exp_to_leaderboard.sbatch) echo " exp_to_leaderboard job ID: ${LB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index ed260875d..bc096198c 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -46,20 +46,29 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting posterior_diagnostic_plots (L96 flux-force, after ${EMU_JID}) ===" +echo "=== Submitting pushforward_from_posterior (L96 flux-force, after ${EMU_JID}) ===" +PUSHFWD_JID=$(sbatch --parsable \ + -A esm \ + --job-name="pushfwd_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_flux \ + pushforward_from_posterior.sbatch) +echo " pushforward_from_posterior job ID: ${PUSHFWD_JID}" + +echo "=== Submitting posterior_diagnostic_plots (L96 flux-force, after ${PUSHFWD_JID}) ===" POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_flux \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" -echo "=== Submitting exp_to_leaderboard (L96 flux-force, after ${EMU_JID}) ===" +echo "=== Submitting exp_to_leaderboard (L96 flux-force, after ${PUSHFWD_JID}) ===" LB_JID=$(sbatch --parsable \ -A esm \ --job-name="leaderboard_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_flux \ exp_to_leaderboard.sbatch) echo " exp_to_leaderboard job ID: ${LB_JID}" diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 65bc1ae95..777588fb3 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -46,20 +46,29 @@ EMU_JID=$(sbatch --parsable \ emulate_sample_array.sbatch) echo " emulate_sample job ID: ${EMU_JID}" -echo "=== Submitting posterior_diagnostic_plots (L96 vec-force, after ${EMU_JID}) ===" +echo "=== Submitting pushforward_from_posterior (L96 vec-force, after ${EMU_JID}) ===" +PUSHFWD_JID=$(sbatch --parsable \ + -A esm \ + --job-name="pushfwd_${LABEL}" \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_vec \ + pushforward_from_posterior.sbatch) +echo " pushforward_from_posterior job ID: ${PUSHFWD_JID}" + +echo "=== Submitting posterior_diagnostic_plots (L96 vec-force, after ${PUSHFWD_JID}) ===" POST_DIAG_JID=$(sbatch --parsable \ -A esm \ --job-name="post_diag_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_vec \ posterior_diagnostic_plots_l96.sbatch) echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" -echo "=== Submitting exp_to_leaderboard (L96 vec-force, after ${EMU_JID}) ===" +echo "=== Submitting exp_to_leaderboard (L96 vec-force, after ${PUSHFWD_JID}) ===" LB_JID=$(sbatch --parsable \ -A esm \ --job-name="leaderboard_${LABEL}" \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_vec \ exp_to_leaderboard.sbatch) echo " exp_to_leaderboard job ID: ${LB_JID}" diff --git a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl index aed68e657..c2bef088e 100644 --- a/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -3,19 +3,16 @@ using Dates using JLD2 using Distributions using LinearAlgebra -using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") -include("Lorenz63.jl") ########################################################################### #################### Metric parameters ################################### ########################################################################### -n_pushforward_samples = 1000 # pushforward ensemble size for output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) @@ -57,9 +54,16 @@ end ### Load data -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +# Load first valid file to determine parameter dimension, pushforward dimensions, and max k +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run pushforward_from_posterior_l63.jl first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) n_k = maximum( maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) @@ -69,30 +73,10 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) -# Reconstruct deterministic l63 setup for pushforward (same seeds as calibrate_l63.jl) -nx = 3 -ny = 9 -t_step = 0.01 -T = 40.0 -lorenz_config_settings = LorenzConfig(t_step, T) -T_start = 30.0 -T_end = T -observation_config = ObservationConfig(T_start, T_end) -truth_params_setup = EnsembleMemberConfig([28.0, 8.0 / 3.0]) -rng_seed_init = 11 -rng_i = MersenneTwister(rng_seed_init) -T_long = 1000.0 -x_initial_setup = rand(rng_i, Normal(0.0, 1.0), nx) -x_spun_up = lorenz_solve(truth_params_setup, x_initial_setup, LorenzConfig(t_step, T_long)) -x0 = x_spun_up[:, end] -covT = 2000.0 -cov_solve = lorenz_solve(truth_params_setup, x0, LorenzConfig(t_step, covT)) -ic_cov_sqrt = sqrt(0.1 * cov(cov_solve, dims = 2)) -n_output = ny - -# Load y from first valid results file (y is deterministic, same across all cases) -first_results = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) -y = first_results["y"] +# y is needed for output-space metrics; load from the shared prelim file +prelim_file = joinpath(homedir, "output", "l63_computed_preliminaries.jld2") +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l63.jl first.") +y = JLD2.load(prelim_file, "y") # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) @@ -113,10 +97,18 @@ for (N_ens, rng_idx) in valid_file_items @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run pushforward_from_posterior_l63.jl first." + continue + end + posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] truth_params = loaded["truth_params"] + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) @@ -150,18 +142,9 @@ for (N_ens, rng_idx) in valid_file_items # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # pushforward ensemble for output space - push_unc = sample(post_dist, n_pushforward_samples) - push_con = transform_unconstrained_to_constrained(post_dist, push_unc) - @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" - for s in 1:n_pushforward_samples - output_samples_arr[i, j, k, s, :] = lorenz_forward( - EnsembleMemberConfig(push_con[:, s]), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) - end + # load precomputed pushforward samples (produced by pushforward_from_posterior_l63.jl) + ki = findfirst(==(k), pf_k_values) + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] # output-space metrics (empirical distribution of pushforward samples vs truth output) os = output_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_output) diff --git a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl index 55498525f..bcca90bb1 100644 --- a/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -3,20 +3,16 @@ using Dates using JLD2 using Distributions using LinearAlgebra -using Flux -using Random using CalibrateEmulateSample.ParameterDistributions using CalibrateEmulateSample.DataContainers using CalibrateEmulateSample.EnsembleKalmanProcesses include("experiment_config.jl") -include("Lorenz96.jl") ########################################################################### #################### Metric parameters ################################### ########################################################################### -n_pushforward_samples = 1000 # pushforward ensemble size for forcing/output metrics n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis marginal_coverage_quantiles = collect(0.05:0.05:0.95) # quantile levels for marginal coverage fraction n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) @@ -60,9 +56,17 @@ end ### Load data -# Load first valid file to determine parameter dimension and max k -first_loaded = JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...))) -n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +# Load first valid file to determine parameter dimension, pushforward dimensions, and max k +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run pushforward_from_posterior_l96.jl first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) +n_forcing = size(first_loaded["pushforward_forcing_samples"], 2) n_k = maximum( maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) @@ -72,28 +76,10 @@ n_k = maximum( n_rng = length(rng_idxs) n_ens = length(N_enss) -# Load Lorenz preliminaries and truth structure for pushforward +# y is needed for output-space metrics; load once from the shared prelim file prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") -prelim_data = JLD2.load(prelim_file) -x0 = prelim_data["x0"] -nx = length(x0) -ic_cov_sqrt = prelim_data["ic_cov_sqrt"] -lorenz_config_settings = prelim_data["lorenz_config_settings"] -observation_config = prelim_data["observation_config"] -y = prelim_data["y"] -n_output = 2 * nx - -first_calib = JLD2.load(joinpath(data_save_directory, results_filename(cfg, valid_file_items[1]...))) -(truth_phi, _) = first_calib["truth_params_structure"] -(phi_structure, sample_range) = if force_case == "const-force" - (nothing, nothing) -elseif force_case == "vec-force" - (nothing, nothing) -elseif force_case == "flux-force" - (truth_phi.model, truth_phi.sample_range) -end -n_forcing = force_case == "flux-force" ? length(sample_range) : nx +y = JLD2.load(prelim_file, "y") # Pre-allocate: posterior stats have a k_iter dimension; calibration cost and truth do not true_param_arr = fill(NaN, n_rng, n_ens, n_params) @@ -121,10 +107,19 @@ for (N_ens, rng_idx) in valid_file_items @info "loading case $(post_fn)" loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run pushforward_from_posterior_l96.jl first." + continue + end + posteriors_by_k = loaded["posteriors_by_k"] k_values = loaded["k_values"] - truth_params = loaded["truth_params"] - truth_params_constrained = loaded["truth_params_constrained"] + truth_params = loaded["truth_params"] + + pf_forcing = loaded["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + truth_forcing_vec = loaded["truth_forcing"] # Load ekpobj for total calibration cost ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) @@ -133,11 +128,8 @@ for (N_ens, rng_idx) in valid_file_items i = findfirst(==(rng_idx), rng_idxs) j = findfirst(==(N_ens), N_enss) - true_param_arr[i, j, :] = truth_params - n_evals_arr[i, j] = conv_alg_iters * N_ens - - emc_truth = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) - truth_forcing_vec = forcing(emc_truth, x0) + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens truth_forcing_arr[i, j, :] = truth_forcing_vec for k in k_values @@ -163,20 +155,10 @@ for (N_ens, rng_idx) in valid_file_items # log p(truth | posterior) - log p(mode | posterior) logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) - # pushforward ensemble for forcing and output space - push_unc = sample(post_dist, n_pushforward_samples) - push_con = transform_unconstrained_to_constrained(post_dist, push_unc) - @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz forward maps" - for s in 1:n_pushforward_samples - emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) - forcing_samples_arr[i, j, k, s, :] = forcing(emc, x0) - output_samples_arr[i, j, k, s, :] = lorenz_forward( - emc, - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) - end + # load precomputed pushforward samples (produced by pushforward_from_posterior_l96.jl) + ki = findfirst(==(k), pf_k_values) + forcing_samples_arr[i, j, k, :, :] = pf_forcing[:, :, ki] + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] # forcing-space metrics (empirical distribution of pushforward samples vs truth forcing) fs = forcing_samples_arr[i, j, k, :, :] # (n_pushforward_samples, n_forcing) diff --git a/examples/EKIRace/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl index 60ff932e8..bb64e2c58 100644 --- a/examples/EKIRace/posterior_diagnostic_plots_l96.jl +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -122,11 +122,19 @@ for (N_ens, rng_idx) in valid_file_items loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) loaded_c = JLD2.load(joinpath(data_save_directory, calib_fn)) + if !haskey(loaded_p, "pushforward_output_samples") + error("Pushforward data not found in $(post_fn). Run pushforward_from_posterior_l96.jl first.") + end + posteriors_by_k = loaded_p["posteriors_by_k"] priors = loaded_p["priors"] k_values = loaded_p["k_values"] truth_params = loaded_p["truth_params"] + pf_output = loaded_p["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_forcing = loaded_p["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + pf_k_values = loaded_p["pushforward_k_values"] + (truth_params_obj, _) = loaded_c["truth_params_structure"] (truth_params_constrained, structure, sample_range) = if force_case == "const-force" @@ -167,25 +175,19 @@ for (N_ens, rng_idx) in valid_file_items for k in k_values post_dist = posteriors_by_k[k] - push_ensemble = sample(post_dist, n_samples_pushforward) - # note: push_ensemble is unconstrained + + # sample posterior for parameter-space analysis (no forward map needed) + push_ensemble = sample(post_dist, n_samples_pushforward) constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) - G_ens = hcat( - [ - lorenz_forward( - build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), - x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), - lorenz_config_settings, - observation_config, - ) for j in 1:n_samples_pushforward - ]..., - ) + # load precomputed posterior pushforward (produced by pushforward_from_posterior_l96.jl) + ki = findfirst(==(k), pf_k_values) + G_ens = pf_output[:, :, ki]' # (n_output, n_samples) + push_forcings = pf_forcing[:, :, ki]' # (n_forcing, n_samples) ny = length(y) n_par = length(truth_params_constrained) - push_forcings = hcat([forcing(build_forcing(truth_params_obj, constrained_push_ensemble[:, j], structure, sample_range), x0) for j in 1:n_samples_pushforward]...) - param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) # Mahalanobis and logpdf metrics in parameter, forcing, and output spaces frc_samples_m = is_const ? push_forcings[1:1, :] : push_forcings From 002c95098d0507de259fcaa8c91d15f3475bdecb Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 12 Jun 2026 16:35:43 -0700 Subject: [PATCH 76/86] add readme --- examples/EKIRace/hpc-variant/README.md | 56 +++++++++++++++++--------- 1 file changed, 36 insertions(+), 20 deletions(-) diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index 531f83dc6..32ac01859 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -26,21 +26,24 @@ submission time. ### L63 ``` -calibrate_array ──afterok──► emulate_sample_array ──afterany──► exp_to_leaderboard +calibrate_array ──afterok──► emulate_sample_array ──afterany──► pushforward_from_posterior ──afterany──► exp_to_leaderboard ``` ### L96 (const / vec / flux) ``` ┌──afterok──► calibration_diagnostic_plots_l96 -calibrate_array ──afterok──► emulate_sample_array ──afterany──► posterior_diagnostic_plots_l96 - ──afterany──► exp_to_leaderboard +calibrate_array ──afterok──► emulate_sample_array ──afterany──► pushforward_from_posterior ──afterany──► posterior_diagnostic_plots_l96 + ──afterany──► exp_to_leaderboard ``` `calibration_diagnostic_plots` and `emulate_sample` both start once calibrate -succeeds (they run in parallel). `posterior_diagnostic_plots` and -`exp_to_leaderboard` both start once `emulate_sample` finishes (whether or not -it succeeded, so partial results are still processed). +succeeds (they run in parallel). `pushforward_from_posterior` starts once +`emulate_sample` finishes and runs the Lorenz forward map for every posterior +cell in parallel, saving forcing and output samples back into each posterior +JLD2 file. `posterior_diagnostic_plots` and `exp_to_leaderboard` both start +once `pushforward_from_posterior` finishes — they simply load the precomputed +samples rather than re-running the forward map. ## Standalone (serial) @@ -50,11 +53,13 @@ Run with no arguments to sweep all `(N_ens, rng_idx)` cells sequentially. # L63 julia --project=. calibrate_l63.jl julia --project=. emulate_sample_l63.jl +julia --project=. pushforward_from_posterior_l63.jl # L96 — set EXPERIMENT env var or edit the toggle in experiment_config.jl EXPERIMENT=l96_const julia --project=. calibrate_l96.jl EXPERIMENT=l96_const julia --project=. calibration_diagnostic_plots_l96.jl EXPERIMENT=l96_const julia --project=. emulate_sample_l96.jl +EXPERIMENT=l96_const julia --project=. pushforward_from_posterior_l96.jl EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl ``` @@ -62,12 +67,15 @@ You can also run a single cell by passing its 1-based task index for the array-capable scripts: ```bash -julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only -julia --project=. emulate_sample_l63.jl 5 # fifth cell only +julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only +julia --project=. emulate_sample_l63.jl 5 # fifth cell only +julia --project=. pushforward_from_posterior_l63.jl 5 # fifth cell only +EXPERIMENT=l96_const julia --project=. pushforward_from_posterior_l96.jl 3 EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl 3 ``` -Leaderboard conversion runs all cells serially in a single call: +Leaderboard conversion runs all cells serially in a single call (requires +pushforward to have been run first): ```bash julia --project=. l63_exp_to_leaderboard_utilities.jl @@ -115,8 +123,11 @@ CALIB_JID=$(sbatch --parsable -A esm \ EMU_JID=$(sbatch --parsable -A esm \ --dependency=afterok:${CALIB_JID} --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l63.jl emulate_sample_array.sbatch) +PUSHFWD_JID=$(sbatch --parsable -A esm \ + --dependency=afterany:${EMU_JID} \ + pushforward_from_posterior.sbatch) sbatch -A esm \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ exp_to_leaderboard.sbatch # L96 @@ -131,12 +142,16 @@ EMU_JID=$(sbatch --parsable -A esm \ --dependency=afterok:${CALIB_JID} --kill-on-invalid-dep=yes \ --export=ALL,SCRIPT=emulate_sample_l96.jl,EXPERIMENT=l96_const \ emulate_sample_array.sbatch) +PUSHFWD_JID=$(sbatch --parsable -A esm \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + pushforward_from_posterior.sbatch) sbatch -A esm \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_const \ posterior_diagnostic_plots_l96.sbatch sbatch -A esm \ - --dependency=afterany:${EMU_JID} \ + --dependency=afterany:${PUSHFWD_JID} \ --export=ALL,EXPERIMENT=l96_const \ exp_to_leaderboard.sbatch ``` @@ -145,20 +160,21 @@ sbatch -A esm \ | File | Type | Description | |------|------|-------------| -| `calibrate_array.sbatch` | array (1–60) | One task per `(N_ens, rng_idx)` cell | -| `emulate_sample_array.sbatch` | array (1–60) | One task per `(N_ens, rng_idx)` cell | +| `calibrate_array.sbatch` | array (1–180) | One task per `(N_ens, rng_idx)` cell | +| `emulate_sample_array.sbatch` | array (1–180) | One task per `(N_ens, rng_idx)` cell | +| `pushforward_from_posterior.sbatch` | array (1–180) | Posterior pushforward (forcing + output), one task per cell; saves results into the posterior JLD2 | | `calibration_diagnostic_plots_l96.sbatch` | single job | Calibration figures, all cells serially (L96) | -| `posterior_diagnostic_plots_l96.sbatch` | array (1–60) | Pushforward figures, one task per cell (L96) | -| `exp_to_leaderboard.sbatch` | single job | NetCDF leaderboard file, all cells serially | +| `posterior_diagnostic_plots_l96.sbatch` | array (1–180) | Posterior ribbon/scatter figures, one task per cell (L96); loads pushforward from JLD2 | +| `exp_to_leaderboard.sbatch` | single job | NetCDF leaderboard file, all cells serially; loads pushforward from JLD2 | | `precompile.sbatch` | single job | `Pkg.instantiate()` + `Pkg.precompile()` | ### Adjusting array size -The sbatch files default to `--array=1-60` (3 ensemble sizes × 20 repeats). +The sbatch files default to `--array=1-180` (9 ensemble sizes × 20 repeats). If you change `N_ens_sizes` or `n_repeats` in `experiment_config.jl`, update -the upper bound to `length(N_ens_sizes) * n_repeats`. The `%20` suffix on -`posterior_diagnostic_plots_l96.sbatch` caps concurrent tasks to 20 at a time -as a cluster-courtesy limit; raise or remove it if you want faster turnaround. +the upper bound to `length(N_ens_sizes) * n_repeats` in every array sbatch file, +including `pushforward_from_posterior.sbatch`. The `%100` suffix caps concurrent +tasks as a cluster-courtesy limit; raise or remove it if you want faster turnaround. ### Smoke test From 0ec918315821dd2da19035c922102a7df636afdd Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 12 Jun 2026 17:04:39 -0700 Subject: [PATCH 77/86] add sbatch commands individually --- examples/EKIRace/hpc-variant/submit_l63.sh | 3 +++ examples/EKIRace/hpc-variant/submit_l96_const.sh | 2 ++ examples/EKIRace/hpc-variant/submit_l96_flux.sh | 5 ++++- examples/EKIRace/hpc-variant/submit_l96_vec.sh | 3 +++ 4 files changed, 12 insertions(+), 1 deletion(-) diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 4477c7bbf..9406447f1 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -54,4 +54,7 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" +# sbatch --parsable -A esm --job-name="pushfwd_l63" --export=ALL pushforward_from_posterior.sbatch + + # sbatch -A esm --job-name="leaderboard_l63" --export=ALL, exp_to_leaderboard.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_const.sh b/examples/EKIRace/hpc-variant/submit_l96_const.sh index 272d0a379..e9a18d56f 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_const.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -75,6 +75,8 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" +# sbatch --parsable -A esm --job-name="pushfwd_l96_const" --export=ALL,EXPERIMENT=l96_const pushforward_from_posterior.sbatch + # sbatch -A esm --job-name="post_diag_l96_const" --export=ALL,EXPERIMENT=l96_const posterior_diagnostic_plots_l96.sbatch # sbatch -A esm --job-name="leaderboard_l96_const" --export=ALL,EXPERIMENT=l96_const exp_to_leaderboard.sbatch diff --git a/examples/EKIRace/hpc-variant/submit_l96_flux.sh b/examples/EKIRace/hpc-variant/submit_l96_flux.sh index bc096198c..58edc2910 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_flux.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -75,6 +75,9 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" -#sbatch --parsable -A esm --job-name="post_diag_l96_flux" --export=ALL,EXPERIMENT=l96_flux posterior_diagnostic_plots_l96.sbatch +# sbatch --parsable -A esm --job-name="pushfwd_l96_flux" --export=ALL,EXPERIMENT=l96_flux pushforward_from_posterior.sbatch + +# sbatch --parsable -A esm --job-name="post_diag_l96_flux" --export=ALL,EXPERIMENT=l96_flux posterior_diagnostic_plots_l96.sbatch # sbatch -A esm --job-name="leaderboard_l96_flux" --export=ALL,EXPERIMENT=l96_flux exp_to_leaderboard.sbatch + diff --git a/examples/EKIRace/hpc-variant/submit_l96_vec.sh b/examples/EKIRace/hpc-variant/submit_l96_vec.sh index 777588fb3..9ee7e2d8f 100755 --- a/examples/EKIRace/hpc-variant/submit_l96_vec.sh +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -75,6 +75,9 @@ echo " exp_to_leaderboard job ID: ${LB_JID}" echo "=== Done. Monitor with: squeue -u \$USER ===" + +# sbatch --parsable -A esm --job-name="pushfwd_l96_vec" --export=ALL,EXPERIMENT=l96_vec pushforward_from_posterior.sbatch + # sbatch -A esm --job-name="post_diag_l96_vec" --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch # sbatch -A esm --job-name="leaderboard_l96_vec" --export=ALL,EXPERIMENT=l96_vec exp_to_leaderboard.sbatch From e50132e84755169333647b1a0aadb6e13b55a70b Mon Sep 17 00:00:00 2001 From: odunbar Date: Fri, 12 Jun 2026 17:05:08 -0700 Subject: [PATCH 78/86] reduce time lims --- examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch | 2 +- examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch index cbf5cc297..ab37c23bc 100644 --- a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch @@ -11,7 +11,7 @@ #SBATCH --job-name=leaderboard #SBATCH --output=output/slurm/leaderboard_%j.out #SBATCH --error=output/slurm/leaderboard_%j.err -#SBATCH --time=04:00:00 +#SBATCH --time=1:00:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 diff --git a/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch index 4935401ac..4a1baec41 100644 --- a/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch @@ -16,7 +16,7 @@ #SBATCH --output=output/slurm/pushfwd_%A_%a.out #SBATCH --error=output/slurm/pushfwd_%A_%a.err #SBATCH --array=1-180%100 -#SBATCH --time=01:00:00 +#SBATCH --time=00:30:00 #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 From 17329ee5065f6c5775345fb9f7fcb153a18b9397 Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 15 Jun 2026 21:38:02 -0700 Subject: [PATCH 79/86] add graphs for race metric, and metric computation. reduce exp_to_leaderboard files --- .../compute_leaderboard_metrics.jl | 301 +++++++++++- .../EKIRace/hpc-variant/exp_to_leaderboard.jl | 427 ++++++++++++++++++ .../hpc-variant/exp_to_leaderboard.sbatch | 6 +- examples/EKIRace/hpc-variant/submit_l63.sh | 3 +- 4 files changed, 722 insertions(+), 15 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/exp_to_leaderboard.jl diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index a6ab370e3..6ed8c13c7 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -11,7 +11,7 @@ calib_date = Date("2026-06-11", "yyyy-mm-dd") indir = joinpath("output","from-hpc_$(calib_date)") experiment_list = [:l63, :l96_const, :l96_vec, :l96_flux ] -experiment = experiment_list[1] +experiment = experiment_list[4] if experiment == :l63 filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" prelim_jld2_file = "l63_computed_preliminaries.jld2" @@ -43,6 +43,7 @@ prelim_filename = joinpath(indir, prelim_jld2_file) calibration_check_quantiles = [0.15, 0.5, 0.85] # quantile levels for calibration check scores and coverage n_lowrank_modes = 2 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis R_variance_retain = 0.99 # fraction of R variance to retain for R-whitened PCA coverage +budget_target_scalings = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5] # c values for α_c(q) = c·√(q(1−q)/N_y) ########################################################################### @@ -58,15 +59,17 @@ R_variance_retain = 0.99 # fraction of R variance to retai # :output — output space (pushforward posterior vs observations) # # display_metrics — which metric types to display: -# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) -# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) -# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) -# :coverage — marginal coverage fractions (forcing and output only) +# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) +# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) +# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) +# :coverage — marginal coverage fractions (forcing and output only) +# :budget_for_coverage — per-repeat budget (N_ens·k_iter) to reach c·√(q(1-q)/N_y) tolerance +# for each (c, N_ens) pair; reports mean/median/IQR/min/max across repeats # # display_ens_sizes — which ensemble-size values (integers) to show in tables. display_spaces = [:output] # e.g. [:param, :forcing] -display_metrics = [:coverage, :mahalanobis] # e.g. [:mahalanobis, :plr_mahalanobis] +display_metrics = [:coverage, :budget_for_coverage]#, :mahalanobis] # e.g. [:mahalanobis, :plr_mahalanobis] display_ens_sizes = nothing # e.g. [10, 50] show_metric(m) = isnothing(display_metrics) || m in display_metrics @@ -109,6 +112,7 @@ avail_metrics = Symbol[:mahalanobis, :logratio] (haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage") || (haskey(ncd, "output_samples") && !isnothing(y_truth))) && push!(avail_spaces, :output) (haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "output_plr_mahalanobis_top")) && push!(avail_metrics, :plr_mahalanobis) ((haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing")) || (haskey(ncd, "output_samples") && !isnothing(y_truth)) || haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) +(haskey(ncd, "output_coverage") || haskey(ncd, "forcing_coverage")) && push!(avail_metrics, :budget_for_coverage) println("="^72) println("File: $(filename)") @@ -481,9 +485,6 @@ end if show_metric(:coverage) # ── Output-space coverage ────────────────────────────────────────────── - @info show_space(:output) - @info haskey(ncd, "output_samples") - @info !isnothing(y_truth) if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) n_out = ncd.dim["output_dim"] @@ -620,3 +621,285 @@ if show_metric(:coverage) end end end + +########################################################################### +#################### Budget for coverage ################################# +########################################################################### +# For each (c, N_ens) pair: smallest k_iter (per repeat) such that +# |S(q) − q| ≤ α_c(q) = c·√(q(1−q)/N_y) for ALL quantiles q. +# Budget = N_ens·k_iter. NaN = target never reached within the k_iter range. +# Prints mean/median/IQR/min/max; reports 0 for all stats when n_conv = 0. +# Saves one figure per space with one panel per c, box-and-whisker per N_ens. + +if show_metric(:budget_for_coverage) + + ENV["GKSwstype"] = "100" + using Plots + using Plots.Measures + + bq_vals = calibration_check_quantiles + + # ── Table printer ────────────────────────────────────────────────────── + # When a (N_ens, c) pair achieves the target in zero repeats, all statistics + # are reported as 0 (consistent sentinel for "not achieved"). + function _bfcov_print_table(budgets_mat, label, n_dim_label) + println("\n$(label) (N_y = $(n_dim_label))") + println(@sprintf(" %-10s %-6s %8s %8s %14s %8s %8s %s", + "N_ens", "c", "mean", "median", "IQR[25,75]", "min", "max", "conv")) + for (ei, ev) in enumerate(ens_vals) + show_ens(ev) || continue + for (si, c) in enumerate(budget_target_scalings) + b = filter(!isnan, vec(budgets_mat[:, ei, si])) + n_conv = length(b) + if n_conv == 0 + println(@sprintf(" %-10d %-6.2f %8.1f %8.1f [%5.0f,%5.0f] %8.1f %8.1f %d/%d", + ev, c, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0, n_rng)) + else + q25, q75 = quantile(b, [0.25, 0.75]) + println(@sprintf(" %-10d %-6.2f %8.1f %8.1f [%5.0f,%5.0f] %8.1f %8.1f %d/%d", + ev, c, mean(b), median(b), q25, q75, minimum(b), maximum(b), n_conv, n_rng)) + end + end + end + end + + # ── Figure builder ───────────────────────────────────────────────────── + # Layout: (2*n_rows, n_cols). For each c value at subplot index si: + # Top tier (si): budget whisker per N_ens column + # Bottom tier (si + n_rows*n_cols): failure-fraction bar per N_ens column + # + # Whisker conventions (top tier): + # n_conv = 0 → red ✕ at y = 0 + # n_conv = 1 → marker at the single budget value + # n_conv = 2 → vertical line (min → max) with end caps + # n_conv ≥ 3 → vertical line (min → max), end caps, filled median marker + # Annotation "n_conv/n_rng" floats above each whisker column. + # + # Fraction bars (bottom tier): tomato bar height = (n_rng − n_conv) / n_rng, + # y-axis fixed [0, 1]. Taller bar = more failures. + function _bfcov_make_figure(budgets_mat, label, n_dim_label; + ylabel_str = "budget (N_ens · k_iter)", + title_prefix = "Budget for coverage") + n_c = length(budget_target_scalings) + n_cols = n_c <= 4 ? n_c : 4 + n_rows = ceil(Int, n_c / n_cols) + dx = 0.20 # half-width of whisker cap ticks in x-index space + + filt_ens = [(ei, ev) for (ei, ev) in enumerate(ens_vals) if show_ens(ev)] + n_x = length(filt_ens) + + p = plot(layout = (2 * n_rows, n_cols), + size = (max(560 * n_cols, 1120), max(640 * n_rows, 700)), + plot_title = "$(title_prefix) — $(label) (N_y = $(n_dim_label))", + plot_titlefontsize = 10, + margin = 6mm) + + # ── Top tier: budget whiskers ────────────────────────────────────── + for (si, c) in enumerate(budget_target_scalings) + b_all_panel = filter(!isnan, vec(budgets_mat[:, :, si])) + y_ref = isempty(b_all_panel) ? 1.0 : max(maximum(b_all_panel), 1.0) + + for (xi, (ei, ev)) in enumerate(filt_ens) + b_all = vec(budgets_mat[:, ei, si]) + b_conv = filter(!isnan, b_all) + n_conv = length(b_conv) + n_tot = n_rng + + if n_conv == 0 + scatter!(p, [xi], [0.0]; subplot = si, color = :red, + marker = :x, markersize = 8, markerstrokewidth = 2, label = false) + y_ann = 0.06 * y_ref + + elseif n_conv == 1 + scatter!(p, [xi], b_conv; subplot = si, color = :steelblue, + marker = :circle, markersize = 6, label = false) + y_ann = b_conv[1] + + elseif n_conv == 2 + mn, mx = minimum(b_conv), maximum(b_conv) + plot!(p, [xi, xi], [mn, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mn, mn]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mx, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + y_ann = mx + + else # n_conv ≥ 3 + mn, mx, med = minimum(b_conv), maximum(b_conv), median(b_conv) + plot!(p, [xi, xi], [mn, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mn, mn]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mx, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + scatter!(p, [xi], [med]; subplot = si, color = :steelblue, + marker = :circle, markersize = 7, markerstrokecolor = :white, label = false) + y_ann = mx + end + + ann_y = y_ann + 0.07 * y_ref + annotate!(p, [xi], [ann_y], + [Plots.text("$(n_conv)/$(n_tot)", 7, :center, :black)]; subplot = si) + end + + is_left = mod(si - 1, n_cols) == 0 + plot!(p; subplot = si, + title = @sprintf("c = %.2f", c), + titlefontsize = 9, + xlabel = "N_ens", + ylabel = is_left ? ylabel_str : "", + xticks = (collect(1:n_x), string.(last.(filt_ens))), + xlims = (0.5, n_x + 0.5), + xrotation = 30, + legend = false) + end + + # ── Bottom tier: failure-fraction bars ───────────────────────────── + for (si, c) in enumerate(budget_target_scalings) + si_frac = si + n_rows * n_cols + is_left = mod(si - 1, n_cols) == 0 + + frac_fail_vec = [(n_rng - count(!isnan, budgets_mat[:, ei, si])) / n_rng + for (ei, _ev) in filt_ens] + + bar!(p, collect(1:n_x), frac_fail_vec; subplot = si_frac, + bar_width = 0.6, + color = :tomato, + alpha = 0.75, + label = false, + xlabel = "N_ens", + ylabel = is_left ? "frac. failed" : "", + xticks = (collect(1:n_x), string.(last.(filt_ens))), + xlims = (0.5, n_x + 0.5), + ylims = (0.0, 1.05), + yticks = [0.0, 0.5, 1.0], + xrotation = 30, + legend = false) + end + + # ── Hide unused subplots in both tiers ──────────────────────────── + for si in (n_c + 1):(n_rows * n_cols) + plot!(p; subplot = si, axis = false, grid = false, showaxis = false, framestyle = :none) + plot!(p; subplot = si + n_rows * n_cols, axis = false, grid = false, showaxis = false, framestyle = :none) + end + + return p + end + + println("\n" * "="^72) + println("Budget for coverage") + println(" α_c(q) = c·√(q(1−q)/N_y), budget = N_ens·k_iter") + println(" Find min k_iter per repeat where |S(q) − q| ≤ α_c(q) for ALL q") + println(" Quantile grid: $(length(bq_vals)) levels [$(bq_vals[1]), …, $(bq_vals[end])]") + println(" target_scalings c = $(budget_target_scalings)") + println("="^72) + + budget_plot_data = [] # [(label, n_dim_label, budgets_mat), ...] + kiter_plot_data = [] # [(label, n_dim_label, kiters_mat), ...] + + # ── Raw coverage from stored NC variables ────────────────────────────── + for (space_label, space_sym, cov_key, dim_key) in [ + ("Output-space raw", :output, "output_coverage", "output_dim"), + ("Forcing-space raw", :forcing, "forcing_coverage", "forcing_dim"), + ] + (show_space(space_sym) && haskey(ncd, cov_key) && haskey(ncd, dim_key)) || continue + + n_dim_bud = ncd.dim[dim_key] + cov_nc_full = Array(ncd[cov_key]) # (n_rng, n_ens_size, n_k_iter, n_cov_q) + # Subset to calibration_check_quantiles indices within the stored grid + if haskey(ncd, "coverage_quantile") + bq_nc = Float64.(ncd["coverage_quantile"][:]) + bq_idx = [something(findfirst(q -> abs(q - qc) < 1e-10, bq_nc), 0) for qc in bq_vals] + cov_nc = all(>(0), bq_idx) ? cov_nc_full[:, :, :, bq_idx] : cov_nc_full + else + cov_nc = cov_nc_full + end + + budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = cov_nc[ri, ei, ki, :] + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_raw[ri, ei, si] = ev * kv + kiters_raw[ri, ei, si] = kv + break + end + end + end + end + _bfcov_print_table(budgets_raw, space_label, n_dim_bud) + push!(budget_plot_data, (space_label, string(n_dim_bud), budgets_raw)) + push!(kiter_plot_data, (space_label, string(n_dim_bud), kiters_raw)) + end + + # ── R-whitened output coverage (recomputed from raw samples) ────────── + if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) && !isnothing(R_obs) + n_out_bud = ncd.dim["output_dim"] + os_all_bud = coalesce.(Array(ncd["output_samples"]), NaN) + + eig_Rb = eigen(Symmetric(R_obs)) + ord_b = sortperm(eig_Rb.values; rev=true) + λ_all_b = eig_Rb.values[ord_b] + V_all_b = eig_Rb.vectors[:, ord_b] + cum_var_b = cumsum(λ_all_b) ./ sum(λ_all_b) + k_Rb = something(findfirst(>=(R_variance_retain), cum_var_b), length(λ_all_b)) + V_Rb = V_all_b[:, 1:k_Rb] + λ_Rb = λ_all_b[1:k_Rb] + yw_b = V_Rb' * y_truth ./ sqrt.(λ_Rb) + n_dim_wb = k_Rb + + n_bq = length(bq_vals) + wh_cov_b = fill(NaN, n_rng, n_ens_size, n_k_iter, n_bq) + for ri in 1:n_rng, (ei, _ev) in enumerate(ens_vals), ki in 1:n_k_iter + os = os_all_bud[ri, ei, ki, :, :] + all(isnan.(os)) && continue + size(os, 2) != size(V_Rb, 1) && continue + sw = Matrix((V_Rb' * Matrix(os')) ./ sqrt.(λ_Rb))' # (n_ps, k_Rb) + for (qi, qp) in enumerate(bq_vals) + wh_cov_b[ri, ei, ki, qi] = mean(yw_b[d] <= quantile(sw[:, d], qp) for d in 1:n_dim_wb) + end + end + + budgets_wh = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_wh = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_wb) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = wh_cov_b[ri, ei, ki, :] + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_wh[ri, ei, si] = ev * kv + kiters_wh[ri, ei, si] = kv + break + end + end + end + end + wh_label = "Output-space R-whitened" + wh_ndim = "$(n_dim_wb) eff. dims ($(round(Int, 100*R_variance_retain))% var of R)" + _bfcov_print_table(budgets_wh, wh_label, wh_ndim) + push!(budget_plot_data, (wh_label, wh_ndim, budgets_wh)) + push!(kiter_plot_data, (wh_label, wh_ndim, kiters_wh)) + end + + # ── Generate and save figures ────────────────────────────────────────── + for (label, n_dim_label, bmat) in budget_plot_data + p_fig = _bfcov_make_figure(bmat, label, n_dim_label) + tag = replace(lowercase(label), r"[^a-z0-9]+" => "_") + base = joinpath(indir, "budget_for_coverage_$(tag)_$(experiment)_$(calib_date)") + savefig(p_fig, base * ".pdf") + savefig(p_fig, base * ".png") + @info "Saved: $(base).pdf" + end + + for (label, n_dim_label, kmat) in kiter_plot_data + p_fig = _bfcov_make_figure(kmat, label, n_dim_label; + ylabel_str = "k_iter", + title_prefix = "Iterations for coverage") + tag = replace(lowercase(label), r"[^a-z0-9]+" => "_") + base = joinpath(indir, "iters_for_coverage_$(tag)_$(experiment)_$(calib_date)") + savefig(p_fig, base * ".pdf") + savefig(p_fig, base * ".png") + @info "Saved: $(base).pdf" + end +end diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl b/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl new file mode 100644 index 000000000..b0aa6ec8f --- /dev/null +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl @@ -0,0 +1,427 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") + +# Allow EXPERIMENT env var to override the toggle in experiment_config.jl +if haskey(ENV, "EXPERIMENT") + EXPERIMENT = Symbol(ENV["EXPERIMENT"]) +end + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_lowrank_modes = 5 +marginal_coverage_quantiles = collect(0.05:0.05:0.95) +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) +budget_target_scalings = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5] +n_target_scalings = length(budget_target_scalings) + +########################################################################### +#################### Config setup ######################################## +########################################################################### + +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) +has_forcing = cfg.force_case !== nothing + +nc_save_filename = nc_filename(cfg, method) + +########################################################################### +#################### Locate valid posterior files ######################## +########################################################################### + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss, rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample first.") +end + +########################################################################### +#################### Determine dimensions ################################ +########################################################################### + +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run pushforward_from_posterior.sbatch first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) +n_forcing = has_forcing ? size(first_loaded["pushforward_forcing_samples"], 2) : 0 + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +########################################################################### +#################### Load truth observation vector ####################### +########################################################################### + +prelim_file = if cfg.force_case === nothing + joinpath(homedir, "output", "$(cfg.model)_computed_preliminaries.jld2") +else + joinpath(homedir, "output", "$(cfg.model)_computed_preliminaries_$(cfg.force_case).jld2") +end +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate first.") +y = JLD2.load(prelim_file, "y") + +########################################################################### +#################### Pre-allocate arrays ################################# +########################################################################### + +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) + +if has_forcing + forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) + forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) + truth_forcing_arr = fill(NaN, n_rng, n_ens, n_forcing) +end + +########################################################################### +#################### Main loop over posterior files ###################### +########################################################################### + +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run pushforward_from_posterior.sbatch first." + continue + end + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + + if has_forcing + pf_forcing = loaded["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + truth_forcing_vec = loaded["truth_forcing"] + end + + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens + if has_forcing + truth_forcing_arr[i, j, :] = truth_forcing_vec + end + + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate; recommend num_samples > $(n_params)." + end + + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + ki = findfirst(==(k), pf_k_values) + + # ── Output-space metrics ────────────────────────────────────────── + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] + os = output_samples_arr[i, j, k, :, :] + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - y + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level. Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end + + # ── Forcing-space metrics (L96 only) ───────────────────────────── + if has_forcing + forcing_samples_arr[i, j, k, :, :] = pf_forcing[:, :, ki] + fs = forcing_samples_arr[i, j, k, :, :] + fm = vec(mean(fs, dims=1)) + fc = Symmetric(cov(fs) + 1e-10 * I) + f_normal = MvNormal(fm, fc) + f_cols = Matrix(fs') + f_mode = f_cols[:, argmax(logpdf(f_normal, f_cols))] + f_diff = fm - truth_forcing_vec + forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) + forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + if a_f < 1e-8 + @warn "Forcing-space PLR skipped: noise floor a_f=$(a_f) ≈ regularization level (e.g. const-force). Entries left as NaN." + else + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) + end + end + end +end + +########################################################################### +#################### Budget for coverage ################################# +########################################################################### +# For each (target_scaling c, N_ens, rng seed): smallest k such that +# |S(q) − q| ≤ c·√(q(1−q)/N_y) for ALL quantile levels q. +# budget_to_target = N_ens · k, iters_to_target = k. +# NaN when the target was not reached within the k_iter range. + +output_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) +output_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) +for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_output) + for ri in 1:n_rng, ei in 1:n_ens + for k in 1:n_k + s_q = output_coverage_arr[ri, ei, k, :] + all(isnan.(s_q)) && continue + if all(abs.(s_q .- marginal_coverage_quantiles) .<= tol) + output_budget_to_target[ri, ei, si] = N_enss[ei] * k + output_iters_to_target[ri, ei, si] = k + break + end + end + end +end + +if has_forcing + forcing_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) + forcing_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_forcing) + for ri in 1:n_rng, ei in 1:n_ens + for k in 1:n_k + s_q = forcing_coverage_arr[ri, ei, k, :] + all(isnan.(s_q)) && continue + if all(abs.(s_q .- marginal_coverage_quantiles) .<= tol) + forcing_budget_to_target[ri, ei, si] = N_enss[ei] * k + forcing_iters_to_target[ri, ei, si] = k + break + end + end + end + end +end + +########################################################################### +#################### Save to NetCDF ##################################### +########################################################################### + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) +defDim(ds, "target_scaling", n_target_scalings) +if has_forcing + defDim(ds, "forcing_dim", n_forcing) +end + +# ── Coordinate variables ────────────────────────────────────────────── +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + +ts_var = defVar(ds, "target_scaling", Float64, ("target_scaling",)) +ts_var.attrib["description"] = "Scaling c in α_c(q) = c·√(q(1−q)/N_y); tolerance for budget_to_target / iters_to_target" +ts_var[:] = budget_target_scalings + +# ── Calibration cost and truth ──────────────────────────────────────── +n_evals_v = defVar(ds, "n_evals_to_target", Float64, ("random_seed", "ensemble_size"); fillvalue=NaN) +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." +n_evals_v[:, :] = n_evals_arr + +true_param_v = defVar(ds, "true_param", Float64, ("random_seed", "ensemble_size", "param_dim"); fillvalue=NaN) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :] = true_param_arr + +if has_forcing + truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble_size", "forcing_dim"); fillvalue=NaN) + truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair." + truth_forcing_v[:, :, :] = truth_forcing_arr +end + +# ── Param-space metrics ─────────────────────────────────────────────── +post_mean_v = defVar(ds, "post_mean", Float64, ("random_seed", "ensemble_size", "k_iter", "param_dim"); fillvalue=NaN) +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :] = post_mean_arr + +post_cov_v = defVar(ds, "post_cov", Float64, ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); fillvalue=NaN) +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar(ds, "mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +mahalanobis_v.attrib["description"] = "M = (m-truth)' C^{-1} (m-truth) in parameter space." +mahalanobis_v[:, :, :] = mahal_arr + +posterior_logpdf_true_v_map_v = defVar(ds, "posterior_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +posterior_logpdf_true_v_map_v.attrib["description"] = "P = logpdf(posterior, truth_param) - logpdf(posterior, mode). For Gaussian, -2P = M." +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr + +# ── Output-space metrics ────────────────────────────────────────────── +output_samples_v = defVar(ds, "output_samples", Float64, ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); fillvalue=NaN) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples in output space" +output_samples_v[:, :, :, :, :] = output_samples_arr + +output_mahal_v = defVar(ds, "output_mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om-y)' Co^{-1} (om-y)." +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar(ds, "output_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode)." +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "PLR Mahalanobis top term in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) modes. Ref: Chisq($(n_lowrank_modes))." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "PLR Mahalanobis residual term in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o. Ref: Chisq($(n_output - n_lowrank_modes))." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." +output_coverage_v[:, :, :, :] = output_coverage_arr + +output_budget_v = defVar(ds, "output_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) +output_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) for all q in output space. NaN = not reached." +output_budget_v[:, :, :] = output_budget_to_target + +output_iters_v = defVar(ds, "output_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) +output_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target in output space. NaN = not reached." +output_iters_v[:, :, :] = output_iters_to_target + +# ── Forcing-space metrics (L96 only) ───────────────────────────────── +if has_forcing + forcing_samples_v = defVar(ds, "forcing_samples", Float64, ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "forcing_dim"); fillvalue=NaN) + forcing_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples in forcing space" + forcing_samples_v[:, :, :, :, :] = forcing_samples_arr + + forcing_mahal_v = defVar(ds, "forcing_mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_mahal_v.attrib["description"] = "Mahalanobis distance in forcing space: (fm-truth_forcing)' Cf^{-1} (fm-truth_forcing)." + forcing_mahal_v[:, :, :] = forcing_mahal_arr + + forcing_logpdf_true_v_map_v = defVar(ds, "forcing_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in forcing space: logpdf(N(fm,Cf), truth_forcing) - logpdf(N(fm,Cf), mode)." + forcing_logpdf_true_v_map_v[:, :, :] = forcing_logpdf_true_v_map_arr + + forcing_plr_top_v = defVar(ds, "forcing_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_plr_top_v.attrib["description"] = "PLR Mahalanobis top term in forcing space: sum((Vf'*(fm-truth))^2 / λf_i), top $(n_lowrank_modes) modes. Ref: Chisq($(n_lowrank_modes))." + forcing_plr_top_v[:, :, :] = forcing_plr_mahal_top_arr + + forcing_plr_res_v = defVar(ds, "forcing_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_plr_res_v.attrib["description"] = "PLR Mahalanobis residual term in forcing space: (||fm-truth||^2 - ||Vf'*(fm-truth)||^2) / a_f. Ref: Chisq($(n_forcing - n_lowrank_modes))." + forcing_plr_res_v[:, :, :] = forcing_plr_mahal_residual_arr + + forcing_coverage_v = defVar(ds, "forcing_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) + forcing_coverage_v.attrib["description"] = "Marginal coverage in forcing space: fraction of forcing dims where truth_forcing[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." + forcing_coverage_v[:, :, :, :] = forcing_coverage_arr + + forcing_budget_v = defVar(ds, "forcing_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) + forcing_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) for all q in forcing space. NaN = not reached." + forcing_budget_v[:, :, :] = forcing_budget_to_target + + forcing_iters_v = defVar(ds, "forcing_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) + forcing_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target in forcing space. NaN = not reached." + forcing_iters_v[:, :, :] = forcing_iters_to_target +end + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch index ab37c23bc..de9855f1f 100644 --- a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch @@ -27,8 +27,4 @@ cd "${SLURM_SUBMIT_DIR}" export JULIA_PKG_PRECOMPILE_AUTO=0 -if [[ -z "${EXPERIMENT:-}" || "${EXPERIMENT}" == "l63" ]]; then - julia --project=. l63_exp_to_leaderboard_utilities.jl -else - julia --project=. l96_exp_to_leaderboard_utilities.jl -fi +julia --project=. exp_to_leaderboard.jl diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 9406447f1..49604546c 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -49,6 +49,7 @@ LB_JID=$(sbatch --parsable \ -A esm \ --job-name="leaderboard_${LABEL}" \ --dependency=afterany:${PUSHFWD_JID} \ + --export=ALL,EXPERIMENT=l63 \ exp_to_leaderboard.sbatch) echo " exp_to_leaderboard job ID: ${LB_JID}" @@ -57,4 +58,4 @@ echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch --parsable -A esm --job-name="pushfwd_l63" --export=ALL pushforward_from_posterior.sbatch -# sbatch -A esm --job-name="leaderboard_l63" --export=ALL, exp_to_leaderboard.sbatch +# sbatch -A esm --job-name="leaderboard_l63" --export=ALL,EXPERIMENT=l63 exp_to_leaderboard.sbatch From 73d1b461f14764bf9b934f9757c56912d215d123 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 16 Jun 2026 19:42:39 -0700 Subject: [PATCH 80/86] change the budget and iter to be per quantile --- .../EKIRace/hpc-variant/exp_to_leaderboard.jl | 56 +++++++++---------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl b/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl index b0aa6ec8f..3918c4bc6 100644 --- a/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl @@ -248,17 +248,17 @@ end # budget_to_target = N_ens · k, iters_to_target = k. # NaN when the target was not reached within the k_iter range. -output_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) -output_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) +output_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) +output_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) for (si, c) in enumerate(budget_target_scalings) tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_output) - for ri in 1:n_rng, ei in 1:n_ens + for ri in 1:n_rng, ei in 1:n_ens, (qi, qp) in enumerate(marginal_coverage_quantiles) for k in 1:n_k - s_q = output_coverage_arr[ri, ei, k, :] - all(isnan.(s_q)) && continue - if all(abs.(s_q .- marginal_coverage_quantiles) .<= tol) - output_budget_to_target[ri, ei, si] = N_enss[ei] * k - output_iters_to_target[ri, ei, si] = k + s = output_coverage_arr[ri, ei, k, qi] + isnan(s) && continue + if abs(s - qp) <= tol[qi] + output_budget_to_target[ri, ei, si, qi] = N_enss[ei] * k + output_iters_to_target[ri, ei, si, qi] = k break end end @@ -266,17 +266,17 @@ for (si, c) in enumerate(budget_target_scalings) end if has_forcing - forcing_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) - forcing_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings) + forcing_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) + forcing_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) for (si, c) in enumerate(budget_target_scalings) tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_forcing) - for ri in 1:n_rng, ei in 1:n_ens + for ri in 1:n_rng, ei in 1:n_ens, (qi, qp) in enumerate(marginal_coverage_quantiles) for k in 1:n_k - s_q = forcing_coverage_arr[ri, ei, k, :] - all(isnan.(s_q)) && continue - if all(abs.(s_q .- marginal_coverage_quantiles) .<= tol) - forcing_budget_to_target[ri, ei, si] = N_enss[ei] * k - forcing_iters_to_target[ri, ei, si] = k + s = forcing_coverage_arr[ri, ei, k, qi] + isnan(s) && continue + if abs(s - qp) <= tol[qi] + forcing_budget_to_target[ri, ei, si, qi] = N_enss[ei] * k + forcing_iters_to_target[ri, ei, si, qi] = k break end end @@ -380,13 +380,13 @@ output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ense output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." output_coverage_v[:, :, :, :] = output_coverage_arr -output_budget_v = defVar(ds, "output_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) -output_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) for all q in output space. NaN = not reached." -output_budget_v[:, :, :] = output_budget_to_target +output_budget_v = defVar(ds, "output_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) +output_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) per quantile q in output space. NaN = not reached. Take max over quantiles for the all-quantile condition." +output_budget_v[:, :, :, :] = output_budget_to_target -output_iters_v = defVar(ds, "output_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) -output_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target in output space. NaN = not reached." -output_iters_v[:, :, :] = output_iters_to_target +output_iters_v = defVar(ds, "output_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) +output_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target per quantile in output space. NaN = not reached." +output_iters_v[:, :, :, :] = output_iters_to_target # ── Forcing-space metrics (L96 only) ───────────────────────────────── if has_forcing @@ -414,13 +414,13 @@ if has_forcing forcing_coverage_v.attrib["description"] = "Marginal coverage in forcing space: fraction of forcing dims where truth_forcing[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." forcing_coverage_v[:, :, :, :] = forcing_coverage_arr - forcing_budget_v = defVar(ds, "forcing_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) - forcing_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) for all q in forcing space. NaN = not reached." - forcing_budget_v[:, :, :] = forcing_budget_to_target + forcing_budget_v = defVar(ds, "forcing_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) + forcing_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) per quantile q in forcing space. NaN = not reached. Take max over quantiles for the all-quantile condition." + forcing_budget_v[:, :, :, :] = forcing_budget_to_target - forcing_iters_v = defVar(ds, "forcing_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling"); fillvalue=NaN) - forcing_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target in forcing space. NaN = not reached." - forcing_iters_v[:, :, :] = forcing_iters_to_target + forcing_iters_v = defVar(ds, "forcing_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) + forcing_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target per quantile in forcing space. NaN = not reached." + forcing_iters_v[:, :, :, :] = forcing_iters_to_target end close(ds) From 5ae808b5521c1b284b8c43704b0fb34d3619719b Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 18 Jun 2026 16:38:35 -0700 Subject: [PATCH 81/86] add plots for l63 --- .../compute_leaderboard_metrics.jl | 75 ++++-- .../posterior_diagnostic_plots_l63.jl | 255 ++++++++++++++++++ .../posterior_diagnostic_plots_l63.sbatch | 29 ++ examples/EKIRace/hpc-variant/submit_l63.sh | 9 + 4 files changed, 344 insertions(+), 24 deletions(-) create mode 100644 examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.jl create mode 100644 examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch diff --git a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl index 6ed8c13c7..cb8ba37a4 100644 --- a/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -7,11 +7,11 @@ using Dates using JLD2 using LinearAlgebra -calib_date = Date("2026-06-11", "yyyy-mm-dd") +calib_date = Date("2026-06-15", "yyyy-mm-dd") indir = joinpath("output","from-hpc_$(calib_date)") experiment_list = [:l63, :l96_const, :l96_vec, :l96_flux ] -experiment = experiment_list[4] +experiment = experiment_list[2] if experiment == :l63 filename = "ces-eki-dmc_l63_ensemble_results_$(calib_date).nc" prelim_jld2_file = "l63_computed_preliminaries.jld2" @@ -114,6 +114,8 @@ avail_metrics = Symbol[:mahalanobis, :logratio] ((haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing")) || (haskey(ncd, "output_samples") && !isnothing(y_truth)) || haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) (haskey(ncd, "output_coverage") || haskey(ncd, "forcing_coverage")) && push!(avail_metrics, :budget_for_coverage) +avail_quantiles = haskey(ncd, "coverage_quantile") ? Float64.(ncd["coverage_quantile"][:]) : nothing + println("="^72) println("File: $(filename)") println() @@ -121,11 +123,13 @@ println("All available options:") println(" display_spaces — $(avail_spaces)") println(" display_metrics — $(avail_metrics)") println(" display_ens_sizes — $(ens_vals)") +isnothing(avail_quantiles) || println(" coverage_quantile — $(avail_quantiles) ($(length(avail_quantiles)) levels)") println() println("Currently selected (edit display_* variables above to filter):") println(" display_spaces = $(isnothing(display_spaces) ? "nothing → $(avail_spaces)" : display_spaces)") println(" display_metrics = $(isnothing(display_metrics) ? "nothing → $(avail_metrics)" : display_metrics)") println(" display_ens_sizes = $(isnothing(display_ens_sizes) ? "nothing → $(ens_vals)" : display_ens_sizes)") +isnothing(avail_quantiles) || println(" calibration_check_quantiles = $(calibration_check_quantiles) (used to select quantile subset for budget metric)") println("="^72) mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] @@ -793,35 +797,58 @@ if show_metric(:budget_for_coverage) kiter_plot_data = [] # [(label, n_dim_label, kiters_mat), ...] # ── Raw coverage from stored NC variables ────────────────────────────── - for (space_label, space_sym, cov_key, dim_key) in [ - ("Output-space raw", :output, "output_coverage", "output_dim"), - ("Forcing-space raw", :forcing, "forcing_coverage", "forcing_dim"), + for (space_label, space_sym, cov_key, dim_key, budget_nc_key, iters_nc_key) in [ + ("Output-space raw", :output, "output_coverage", "output_dim", "output_budget_to_target", "output_iters_to_target"), + ("Forcing-space raw", :forcing, "forcing_coverage", "forcing_dim", "forcing_budget_to_target", "forcing_iters_to_target"), ] (show_space(space_sym) && haskey(ncd, cov_key) && haskey(ncd, dim_key)) || continue n_dim_bud = ncd.dim[dim_key] - cov_nc_full = Array(ncd[cov_key]) # (n_rng, n_ens_size, n_k_iter, n_cov_q) - # Subset to calibration_check_quantiles indices within the stored grid - if haskey(ncd, "coverage_quantile") + + # Subset coverage to the requested quantiles; fall back to full grid when + # coverage_quantile coordinate is absent (old files). + can_load = haskey(ncd, "coverage_quantile") + bq_idx = nothing + if can_load bq_nc = Float64.(ncd["coverage_quantile"][:]) - bq_idx = [something(findfirst(q -> abs(q - qc) < 1e-10, bq_nc), 0) for qc in bq_vals] - cov_nc = all(>(0), bq_idx) ? cov_nc_full[:, :, :, bq_idx] : cov_nc_full - else - cov_nc = cov_nc_full + bq_idx = [findfirst(q -> abs(q - qc) < 1e-10, bq_nc) for qc in bq_vals] + can_load = all(!isnothing, bq_idx) end - budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) - kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) - for (si, c) in enumerate(budget_target_scalings) - tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) - for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng - for (ki, kv) in enumerate(kiter_vals) - s_q = cov_nc[ri, ei, ki, :] - all(isnan.(s_q)) && continue - if all(abs.(s_q .- bq_vals) .<= tol) - budgets_raw[ri, ei, si] = ev * kv - kiters_raw[ri, ei, si] = kv - break + if can_load + @info "Computing $(space_label) budget from NC (quantile subset: $(bq_vals))" + cov_bq = Array(ncd[cov_key])[:, :, :, bq_idx] # (n_rng, n_ens, n_k, n_bq) + budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = coalesce.(cov_bq[ri, ei, ki, :], NaN) + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_raw[ri, ei, si] = ev * kv + kiters_raw[ri, ei, si] = kv + break + end + end + end + end + else + cov_nc = Array(ncd[cov_key]) # (n_rng, n_ens_size, n_k_iter, n_cov_q) + budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = coalesce.(cov_nc[ri, ei, ki, :], NaN) + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_raw[ri, ei, si] = ev * kv + kiters_raw[ri, ei, si] = kv + break + end end end end diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.jl new file mode 100644 index 000000000..570d58729 --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.jl @@ -0,0 +1,255 @@ +# Import modules +using Distributions +using LinearAlgebra +using Random +using JLD2 +using Statistics +ENV["GKSwstype"] = "100" +using Plots +using Plots.Measures + + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz63.jl") +include("experiment_config.jl") + +######################################################################## +############### Per-cell pushforward ################################### +######################################################################## + +function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) + m = vec(mean(samples, dims=2)) + C_raw = cov(samples, dims=2) + num_samples = size(samples, 2) + dim_samples = size(samples, 1) + r = rank(C_raw) + if r == num_samples - 1 && r < dim_samples - 1 + @warn "Covariance rank $(r) = num_samples-1 = $(num_samples-1) < dim-1 = $(dim_samples-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(dim_samples)." + end + λ = max(1e-10, 1e-4 * mean(diag(C_raw))) + C = Symmetric(C_raw + λ * I) + dist = MvNormal(m, C) + pmode = samples[:, argmax(logpdf(dist, samples))] + diff = m - truth + mah = diff' * (C \ diff) + lp = logpdf(dist, truth) - logpdf(dist, pmode) + return mah, lp +end + +function ensemble_from_posterior_one(cfg, N_ens, rng_idx; method = method_cases[1]) + n_samples_pushforward = 1000 + + calib_dir = calib_directory(method, cfg) + post_fn = posterior_filename(cfg, N_ens, rng_idx) + + # load preliminaries + prelim_dir = joinpath(@__DIR__, "output") + prelim_file = joinpath(prelim_dir, "l63_computed_preliminaries.jld2") + if !isfile(prelim_file) + throw(ErrorException("preliminaries file not found. \n First run: \n > julia --project calibrate_l63.jl")) + end + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + if !isfile(joinpath(data_save_directory, post_fn)) + @warn "No posterior file found for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] + truth_params_constrained = loaded_p["truth_params_constrained"] + + n_par = length(truth_params_constrained) + ny = length(y) + + # --- sample from prior for grey background --- + prior_samples_unconstrained = sample(priors, n_samples_pushforward) + prior_samples_constrained = transform_unconstrained_to_constrained(priors, prior_samples_unconstrained) + prior_param_diffs = reduce(hcat, [prior_samples_constrained[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + G_prior = hcat( + [ + lorenz_forward( + EnsembleMemberConfig(prior_samples_constrained[:, j]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + if !haskey(loaded_p, "pushforward_output_samples") + error("Pushforward data not found in $(post_fn). Run pushforward_from_posterior.sbatch first.") + end + pf_output = loaded_p["pushforward_output_samples"] # (n_samples, n_output, n_k) + pf_k_values = loaded_p["pushforward_k_values"] + + for k in k_values + post_dist = posteriors_by_k[k] + + # sample posterior for parameter-space analysis (no forward map needed) + push_ensemble = sample(post_dist, n_samples_pushforward) + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + # load precomputed posterior pushforward (produced by pushforward_from_posterior_l63.jl) + ki = findfirst(==(k), pf_k_values) + G_ens = pf_output[:, :, ki]' # (n_output, n_samples) + + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + + # Mahalanobis and logpdf metrics in parameter and output spaces + param_mah, param_lp = pushforward_metrics(push_ensemble, truth_params) + output_mah, output_lp = pushforward_metrics(G_ens, y) + lowbd_par, upbd_par = round.(quantile(Chisq(n_par), [0.01, 0.99]), digits=2) + lowbd_out, upbd_out = round.(quantile(Chisq(ny), [0.01, 0.99]), digits=2) + @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" + @info " param [d=$(n_par)]: target: [$(lowbd_par), $(upbd_par)] mahal=$(round(param_mah, digits=2)) -2*logpdf_ratio=$(round(-2*param_lp, digits=2))" + @info " output [d=$(ny)]: target: [$(lowbd_out), $(upbd_out)] mahal=$(round(output_mah, digits=2)) -2*logpdf_ratio=$(round(-2*output_lp, digits=2))" + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameters (differences from truth) + p2 = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Panel (ii): output + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Output index", + ylabel = "State statistics (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Add prior samples (grey) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + + # Add posterior samples (green) + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + l = @layout [a b] + plt = plot(p2, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + # --- posterior ribbons plot --- + post_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(param_diffs)])' + post_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_ens)])' + prior_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_param_diffs)])' + prior_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_prior)])' + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameter ribbons + pa = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + plot!(pa, collect(1:n_par), prior_param_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_param_quantiles[:, 2] - prior_param_quantiles[:, 1] prior_param_quantiles[:, 3] - prior_param_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pa, collect(1:n_par), post_param_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_param_quantiles[:, 2] - post_param_quantiles[:, 1] post_param_quantiles[:, 3] - post_param_quantiles[:, 2]], + fillalpha = 0.2) + + # Panel (ii): output ribbons + pc = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Output index", + ylabel = "State statistics (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + plot!(pc, 1:ny, prior_G_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_G_quantiles[:, 2] - prior_G_quantiles[:, 1] prior_G_quantiles[:, 3] - prior_G_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pc, 1:ny, post_G_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_G_quantiles[:, 2] - post_G_quantiles[:, 1] post_G_quantiles[:, 3] - post_G_quantiles[:, 2]], + fillalpha = 0.2) + + ribbons_plt = plot(pa, pc, layout = @layout [a b]) + ribbons_fn = "posterior_ribbons_$(suffix)_k$(k)" + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".png")) + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".pdf")) + end +end + +######################################################################## +############### Main dispatcher ######################################## +######################################################################## + +function main() + cfg = experiment_config(:l63) + tasks = flat_tasks(cfg) + idx = task_index_from_args() + + if isnothing(idx) + for (N_ens, rng_idx) in tasks + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end + else + if idx < 1 || idx > length(tasks) + error("SLURM_ARRAY_TASK_ID $(idx) out of range 1:$(length(tasks))") + end + N_ens, rng_idx = tasks[idx] + ensemble_from_posterior_one(cfg, N_ens, rng_idx) + end +end + +main() diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch new file mode 100644 index 000000000..6a90cf553 --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch @@ -0,0 +1,29 @@ +#!/bin/bash +# Array job: one task per (N_ens, rng_idx) pair (180 tasks for the L63 case). +# Each task generates pushforward + posterior-ribbon diagnostic figures for L63. +# +# Invoked automatically by submit_l63.sh after pushforward_from_posterior completes. +# Manual use: +# sbatch --dependency=afterok: posterior_diagnostic_plots_l63.sbatch + +#SBATCH --job-name=post_diag_l63 +#SBATCH --output=output/slurm/post_diag_l63_%A_%a.out +#SBATCH --error=output/slurm/post_diag_l63_%A_%a.err +#SBATCH --array=1-180%100 +#SBATCH --time=00:15:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 + +set -euo pipefail + +module load julia/1.12.2 + +export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} +export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} + +cd "${SLURM_SUBMIT_DIR}" + +export JULIA_PKG_PRECOMPILE_AUTO=0 + +julia --project=. posterior_diagnostic_plots_l63.jl "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/submit_l63.sh b/examples/EKIRace/hpc-variant/submit_l63.sh index 49604546c..db5ba198d 100755 --- a/examples/EKIRace/hpc-variant/submit_l63.sh +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -44,6 +44,14 @@ PUSHFWD_JID=$(sbatch --parsable \ pushforward_from_posterior.sbatch) echo " pushforward_from_posterior job ID: ${PUSHFWD_JID}" +echo "=== Submitting posterior_diagnostic_plots (L63, after ${PUSHFWD_JID}) ===" +POST_DIAG_JID=$(sbatch --parsable \ + -A esm \ + --job-name="post_diag_${LABEL}" \ + --dependency=afterany:${PUSHFWD_JID} \ + posterior_diagnostic_plots_l63.sbatch) +echo " posterior_diagnostic_plots job ID: ${POST_DIAG_JID}" + echo "=== Submitting exp_to_leaderboard (L63, after ${PUSHFWD_JID}) ===" LB_JID=$(sbatch --parsable \ -A esm \ @@ -57,5 +65,6 @@ echo "=== Done. Monitor with: squeue -u \$USER ===" # sbatch --parsable -A esm --job-name="pushfwd_l63" --export=ALL pushforward_from_posterior.sbatch +# sbatch -A esm --job-name="post_diag_l63" posterior_diagnostic_plots_l63.sbatch # sbatch -A esm --job-name="leaderboard_l63" --export=ALL,EXPERIMENT=l63 exp_to_leaderboard.sbatch From 3a384a8be4122e8303aca9fc17a76af35c7a7ec0 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 18 Jun 2026 18:00:41 -0700 Subject: [PATCH 82/86] add constraints of precompile to architectures --- .../EKIRace/hpc-variant/calibrate_array.sbatch | 1 + .../calibration_diagnostic_plots_l96.sbatch | 1 + .../hpc-variant/emulate_sample_array.sbatch | 3 ++- .../EKIRace/hpc-variant/exp_to_leaderboard.sbatch | 5 +++-- .../posterior_diagnostic_plots_l63.sbatch | 1 + .../posterior_diagnostic_plots_l96.sbatch | 1 + examples/EKIRace/hpc-variant/precompile.sbatch | 15 +++++++++++++++ .../hpc-variant/pushforward_from_posterior.sbatch | 1 + 8 files changed, 25 insertions(+), 3 deletions(-) diff --git a/examples/EKIRace/hpc-variant/calibrate_array.sbatch b/examples/EKIRace/hpc-variant/calibrate_array.sbatch index 9a33a2b56..6ed860afe 100644 --- a/examples/EKIRace/hpc-variant/calibrate_array.sbatch +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -25,6 +25,7 @@ #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch index b6176cdd3..d0bafc31a 100644 --- a/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch @@ -15,6 +15,7 @@ #SBATCH --mem=8G #SBATCH --cpus-per-task=1 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch index 156b5a07b..b1caba570 100644 --- a/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -23,10 +23,11 @@ #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail -module load julia/1.12.2 # pin version as needed, e.g.: module load julia/1.10.4 +module load julia/1.12.2 export JULIA_NUM_THREADS=${SLURM_CPUS_PER_TASK} export OPENBLAS_NUM_THREADS=${SLURM_CPUS_PER_TASK} diff --git a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch index de9855f1f..96dafcebc 100644 --- a/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch @@ -11,10 +11,11 @@ #SBATCH --job-name=leaderboard #SBATCH --output=output/slurm/leaderboard_%j.out #SBATCH --error=output/slurm/leaderboard_%j.err -#SBATCH --time=1:00:00 -#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --mem=8G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch index 6a90cf553..eb55e318b 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch @@ -14,6 +14,7 @@ #SBATCH --mem=8G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch index 27239931b..e280e6e19 100644 --- a/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -16,6 +16,7 @@ #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail diff --git a/examples/EKIRace/hpc-variant/precompile.sbatch b/examples/EKIRace/hpc-variant/precompile.sbatch index f48c4bebf..67c6906a8 100644 --- a/examples/EKIRace/hpc-variant/precompile.sbatch +++ b/examples/EKIRace/hpc-variant/precompile.sbatch @@ -10,6 +10,7 @@ #SBATCH --mem=8G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail @@ -19,4 +20,18 @@ cd "${SLURM_SUBMIT_DIR}" export JULIA_NUM_THREADS=$SLURM_CPUS_PER_TASK # parallelize +# To see what is available call: sinfo -o "%20N %10c %20f" +# We explicitly compile on one set of nodes, and then run on another +export JULIA_CPU_TARGET="cascadelake" # explicitly target an architecture +export JULIA_PKG_PRECOMPILE_AUTO=1 + +echo "JULIA_CPU_TARGET=$JULIA_CPU_TARGET" + +julia -e ' +println("CPU = ", Sys.CPU_NAME) +println("JULIA_CPU_TARGET = ", get(ENV,"JULIA_CPU_TARGET","")) +' +# these two should be the same + julia --project=. -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()' + diff --git a/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch index 4a1baec41..e58a74fcc 100644 --- a/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch @@ -20,6 +20,7 @@ #SBATCH --mem=16G #SBATCH --cpus-per-task=4 #SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake set -euo pipefail From 26b1f1f0832f5dda1b8abe11b9db33c7eb4c17b3 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 18 Jun 2026 18:06:28 -0700 Subject: [PATCH 83/86] updated readme --- examples/EKIRace/hpc-variant/README.md | 39 +++++++++++++++++--------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/examples/EKIRace/hpc-variant/README.md b/examples/EKIRace/hpc-variant/README.md index 32ac01859..7f28f5393 100644 --- a/examples/EKIRace/hpc-variant/README.md +++ b/examples/EKIRace/hpc-variant/README.md @@ -26,7 +26,9 @@ submission time. ### L63 ``` -calibrate_array ──afterok──► emulate_sample_array ──afterany──► pushforward_from_posterior ──afterany──► exp_to_leaderboard + ┌──afterany──► posterior_diagnostic_plots_l63 +calibrate_array ──afterok──► emulate_sample_array ──afterany──► pushforward_from_posterior ─┤ + └──afterany──► exp_to_leaderboard ``` ### L96 (const / vec / flux) @@ -37,12 +39,12 @@ calibrate_array ──afterok──► emulate_sample_array ──afterany─ ──afterany──► exp_to_leaderboard ``` -`calibration_diagnostic_plots` and `emulate_sample` both start once calibrate -succeeds (they run in parallel). `pushforward_from_posterior` starts once -`emulate_sample` finishes and runs the Lorenz forward map for every posterior -cell in parallel, saving forcing and output samples back into each posterior -JLD2 file. `posterior_diagnostic_plots` and `exp_to_leaderboard` both start -once `pushforward_from_posterior` finishes — they simply load the precomputed +For L96, `calibration_diagnostic_plots` and `emulate_sample` both start once +calibrate succeeds (they run in parallel). For all cases, `pushforward_from_posterior` +starts once `emulate_sample` finishes and runs the Lorenz forward map for every +posterior cell in parallel, saving forcing and output samples back into each +posterior JLD2 file. `posterior_diagnostic_plots` and `exp_to_leaderboard` both +start once `pushforward_from_posterior` finishes — they load the precomputed samples rather than re-running the forward map. ## Standalone (serial) @@ -54,6 +56,8 @@ Run with no arguments to sweep all `(N_ens, rng_idx)` cells sequentially. julia --project=. calibrate_l63.jl julia --project=. emulate_sample_l63.jl julia --project=. pushforward_from_posterior_l63.jl +julia --project=. posterior_diagnostic_plots_l63.jl +julia --project=. exp_to_leaderboard.jl # L96 — set EXPERIMENT env var or edit the toggle in experiment_config.jl EXPERIMENT=l96_const julia --project=. calibrate_l96.jl @@ -61,25 +65,27 @@ EXPERIMENT=l96_const julia --project=. calibration_diagnostic_plots_l96.jl EXPERIMENT=l96_const julia --project=. emulate_sample_l96.jl EXPERIMENT=l96_const julia --project=. pushforward_from_posterior_l96.jl EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl +EXPERIMENT=l96_const julia --project=. exp_to_leaderboard.jl ``` You can also run a single cell by passing its 1-based task index for the array-capable scripts: ```bash -julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only -julia --project=. emulate_sample_l63.jl 5 # fifth cell only -julia --project=. pushforward_from_posterior_l63.jl 5 # fifth cell only +julia --project=. calibrate_l63.jl 1 # first (N_ens, rng_idx) cell only +julia --project=. emulate_sample_l63.jl 5 # fifth cell only +julia --project=. pushforward_from_posterior_l63.jl 5 # fifth cell only +julia --project=. posterior_diagnostic_plots_l63.jl 5 # fifth cell only EXPERIMENT=l96_const julia --project=. pushforward_from_posterior_l96.jl 3 EXPERIMENT=l96_const julia --project=. posterior_diagnostic_plots_l96.jl 3 ``` -Leaderboard conversion runs all cells serially in a single call (requires +`exp_to_leaderboard.jl` runs all cells serially in a single call (requires pushforward to have been run first): ```bash -julia --project=. l63_exp_to_leaderboard_utilities.jl -EXPERIMENT=l96_const julia --project=. l96_exp_to_leaderboard_utilities.jl +julia --project=. exp_to_leaderboard.jl +EXPERIMENT=l96_const julia --project=. exp_to_leaderboard.jl ``` ## HPC (Caltech Resnick cluster, SLURM) @@ -128,6 +134,10 @@ PUSHFWD_JID=$(sbatch --parsable -A esm \ pushforward_from_posterior.sbatch) sbatch -A esm \ --dependency=afterany:${PUSHFWD_JID} \ + posterior_diagnostic_plots_l63.sbatch +sbatch -A esm \ + --dependency=afterany:${PUSHFWD_JID} \ + --export=ALL,EXPERIMENT=l63 \ exp_to_leaderboard.sbatch # L96 @@ -163,7 +173,8 @@ sbatch -A esm \ | `calibrate_array.sbatch` | array (1–180) | One task per `(N_ens, rng_idx)` cell | | `emulate_sample_array.sbatch` | array (1–180) | One task per `(N_ens, rng_idx)` cell | | `pushforward_from_posterior.sbatch` | array (1–180) | Posterior pushforward (forcing + output), one task per cell; saves results into the posterior JLD2 | -| `calibration_diagnostic_plots_l96.sbatch` | single job | Calibration figures, all cells serially (L96) | +| `calibration_diagnostic_plots_l96.sbatch` | single job | Calibration figures, all cells serially (L96 only) | +| `posterior_diagnostic_plots_l63.sbatch` | array (1–180) | Posterior ribbon/scatter figures, one task per cell (L63); loads pushforward from JLD2 | | `posterior_diagnostic_plots_l96.sbatch` | array (1–180) | Posterior ribbon/scatter figures, one task per cell (L96); loads pushforward from JLD2 | | `exp_to_leaderboard.sbatch` | single job | NetCDF leaderboard file, all cells serially; loads pushforward from JLD2 | | `precompile.sbatch` | single job | `Pkg.instantiate()` + `Pkg.precompile()` | From dc05fc21c08e6f1624ce55af295340084451413f Mon Sep 17 00:00:00 2001 From: odunbar Date: Mon, 22 Jun 2026 15:37:28 -0700 Subject: [PATCH 84/86] updated local pipeline to reflect hpc-variant --- examples/EKIRace/README.md | 54 +- examples/EKIRace/calibrate_l63.jl | 17 + .../EKIRace/compute_leaderboard_metrics.jl | 492 ++++++++++++++++-- examples/EKIRace/exp_to_leaderboard.jl | 418 +++++++++++++++ examples/EKIRace/experiment_config.jl | 31 +- examples/EKIRace/minimal_leaderboard_nc.jl | 110 ++++ .../EKIRace/posterior_diagnostic_plots_l63.jl | 251 +++++++++ .../EKIRace/pushforward_from_posterior_l63.jl | 75 +++ .../EKIRace/pushforward_from_posterior_l96.jl | 99 ++++ 9 files changed, 1484 insertions(+), 63 deletions(-) create mode 100644 examples/EKIRace/exp_to_leaderboard.jl create mode 100644 examples/EKIRace/minimal_leaderboard_nc.jl create mode 100644 examples/EKIRace/posterior_diagnostic_plots_l63.jl create mode 100644 examples/EKIRace/pushforward_from_posterior_l63.jl create mode 100644 examples/EKIRace/pushforward_from_posterior_l96.jl diff --git a/examples/EKIRace/README.md b/examples/EKIRace/README.md index d7532b33d..0a67cabf9 100644 --- a/examples/EKIRace/README.md +++ b/examples/EKIRace/README.md @@ -18,7 +18,8 @@ In `experiment_config.jl`, set `EXPERIMENT` to the case you want to run and pin `calibrate_date` before starting a run, e.g. ```julia -EXPERIMENT = :l96_vec +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[3] # :l96_vec calibrate_date = Date("2026-06-04", "yyyy-mm-dd") ``` @@ -48,8 +49,11 @@ julia --project -e 'include("calibrate_l63.jl")' julia --project -e 'include("calibrate_l96.jl")' ``` -The L96 calibration also writes a `l96_computed_preliminaries_.jld2` -file to `output/` that the later stages require. +Each calibration script also writes a shared preliminaries file to `output/` +that all later stages require: + +- `calibrate_l63.jl` → `output/l63_computed_preliminaries.jld2` +- `calibrate_l96.jl` → `output/l96_computed_preliminaries_.jld2` ### 2. Calibration diagnostic plots (L96 only) @@ -75,17 +79,32 @@ Reads the calibration output for each `(N_ens, rng_idx)` cell, trains an emulator, runs MCMC, and writes per-cell posterior `.jld2` files to the calibration output directory. -### 4. Posterior diagnostic plots (L96 only) +### 4. Pushforward from posterior + +```bash +# L63 +julia --project -e 'include("pushforward_from_posterior_l63.jl")' + +# L96 +julia --project -e 'include("pushforward_from_posterior_l96.jl")' +``` + +For each posterior `.jld2`, draws samples from the fitted posterior and runs +them through the Lorenz forward map, saving the resulting forcing and output +sample arrays back into the same `.jld2` file. This step is required before +running stage 5 (posterior diagnostics) or stage 6 (leaderboard). + +### 5. Posterior diagnostic plots (L96 only) ```bash julia --project -e 'include("posterior_diagnostic_plots_l96.jl")' ``` -Reads each posterior `.jld2`, pushes 400 samples through the Lorenz 96 -forward map, and writes pushforward and posterior-ribbon diagnostic plots -into the calibration output directory. +Reads each posterior `.jld2` (including the pushforward samples written in +stage 4) and writes pushforward and posterior-ribbon diagnostic plots into +the calibration output directory. -### 5. Leaderboard +### 6. Leaderboard Convert the per-cell posterior files into a leaderboard NetCDF file: @@ -97,22 +116,28 @@ julia --project -e 'include("l63_exp_to_leaderboard_utilities.jl")' julia --project -e 'include("l96_exp_to_leaderboard_utilities.jl")' ``` -To compute summary metrics from the resulting netCDF file, edit the `filename` -variable at the top of `compute_leaderboard_metrics.jl` to point to the target -file, then run: +These scripts require the pushforward samples from stage 4 to be present in +each posterior `.jld2`. + +To compute summary metrics from the resulting NetCDF file, pair the `filename` +and `prelim_jld2_file` variables at the top of `compute_leaderboard_metrics.jl` +to point to the target file and its prelim JLD2, then run: ```bash julia --project -e 'include("compute_leaderboard_metrics.jl")' ``` -This prints per-ensemble-size Mahalanobis and log-posterior scores against -chi-squared reference quantiles. +This prints per-ensemble-size calibration scores (Mahalanobis, log-posterior +ratio, marginal coverage) against chi-squared reference quantiles, and saves +budget-for-coverage figures to `indir/`. The prelim JLD2 is needed for +R-whitened PCA coverage; without it, only raw coverage is reported. ## Full run example (L96 vector forcing) ```julia # In experiment_config.jl: -EXPERIMENT = :l96_vec +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[3] # :l96_vec calibrate_date = Date("2026-06-04", "yyyy-mm-dd") ``` @@ -120,6 +145,7 @@ calibrate_date = Date("2026-06-04", "yyyy-mm-dd") julia --project -e 'include("calibrate_l96.jl")' julia --project -e 'include("calibration_diagnostic_plots_l96.jl")' # optional julia --project -e 'include("emulate_sample_l96.jl")' +julia --project -e 'include("pushforward_from_posterior_l96.jl")' julia --project -e 'include("posterior_diagnostic_plots_l96.jl")' # optional julia --project -e 'include("l96_exp_to_leaderboard_utilities.jl")' ``` diff --git a/examples/EKIRace/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl index fe2e7bd1c..15deb01ec 100644 --- a/examples/EKIRace/calibrate_l63.jl +++ b/examples/EKIRace/calibrate_l63.jl @@ -113,6 +113,23 @@ cov_solve = lorenz_solve(truth_params, x0, LorenzConfig(t, covT)) ic_cov = 0.1 * cov(cov_solve, dims = 2) ic_cov_sqrt = sqrt(ic_cov) +prelim_file = joinpath(output_dir, "l63_computed_preliminaries.jld2") +if !isfile(prelim_file) + prelim_tmp = splitext(prelim_file)[1] * ".tmp.$(getpid()).jld2" + JLD2.save( + prelim_tmp, + "x0", x0, + "y", y, + "lorenz_config_settings", lorenz_config_settings, + "observation_config", observation_config, + "ic_cov_sqrt", ic_cov_sqrt, + "R", R, + "R_inv_var", R_inv_var, + ) + mv(prelim_tmp, prelim_file) + @info "Saved computed quantities to $(prelim_file)" +end + ######################################################################## ########################### Running EKI Race ########################### ######################################################################## diff --git a/examples/EKIRace/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl index cdd9404a5..214b3aee7 100644 --- a/examples/EKIRace/compute_leaderboard_metrics.jl +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -3,22 +3,31 @@ using Distributions using Statistics using Printf using DataFrames +using JLD2 +using LinearAlgebra indir = "." #joinpath("output","from-hpc_2026-05-31") #filename = "ces-eki-dmc_l63_ensemble_results_2026-05-29.nc" +#prelim_jld2_file = "l63_computed_preliminaries.jld2" #filename = "ces-eki-dmc_l96_ensemble_results_2026-05-29.nc" +#prelim_jld2_file = "l96_computed_preliminaries_const-force.jld2" #filename = "ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-05-29.nc" -filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-31.nc" +#prelim_jld2_file = "l96_computed_preliminaries_vec-force.jld2" +filename = "ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-05-31.nc" +prelim_jld2_file = "l96_computed_preliminaries_flux-force.jld2" -filename = joinpath(indir, filename) +filename = joinpath(indir, filename) +prelim_filename = joinpath(indir, prelim_jld2_file) @info "computing leaderboard metrics from $(filename)" ########################################################################### #################### Metric parameters ################################### ########################################################################### -calibration_check_quantiles = [0.1, 0.5, 0.9] # quantile levels for calibration check scores and coverage -n_lowrank_modes = 5 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +calibration_check_quantiles = [0.15, 0.5, 0.85] # quantile levels for calibration check scores and coverage +n_lowrank_modes = 2 # top EOF modes for perturbed-low-rank (aI+UDU') Mahalanobis +R_variance_retain = 0.99 # fraction of R variance to retain for R-whitened PCA coverage +budget_target_scalings = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5] # c values for α_c(q) = c·√(q(1−q)/N_y) ########################################################################### #################### Display filters #################################### @@ -33,15 +42,17 @@ n_lowrank_modes = 5 # top EOF modes for perturbed-low # :output — output space (pushforward posterior vs observations) # # display_metrics — which metric types to display: -# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) -# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) -# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) -# :coverage — marginal coverage fractions (forcing and output only) +# :mahalanobis — full Mahalanobis distance M = (m−truth)ᵀ C⁻¹ (m−truth) +# :logratio — log PDF ratio L = log p(truth|post) − log p(mode|post) +# :plr_mahalanobis — perturbed-low-rank M (top / residual / total; forcing and output only) +# :coverage — marginal coverage fractions (forcing and output only) +# :budget_for_coverage — per-repeat budget (N_ens·k_iter) to reach c·√(q(1-q)/N_y) tolerance +# for each (c, N_ens) pair; reports mean/median/IQR/min/max across repeats # # display_ens_sizes — which ensemble-size values (integers) to show in tables. display_spaces = [:output] # e.g. [:param, :forcing] -display_metrics = [:coverage] # e.g. [:mahalanobis, :plr_mahalanobis] +display_metrics = [:coverage, :budget_for_coverage] # e.g. [:mahalanobis, :plr_mahalanobis] display_ens_sizes = nothing # e.g. [10, 50] show_metric(m) = isnothing(display_metrics) || m in display_metrics @@ -52,6 +63,7 @@ show_ens(e) = isnothing(display_ens_sizes) || e in display_ens_sizes #################### Load file and report available options ############## ########################################################################### +@info "computing leaderboard metrics from $(filename)" ncd = NCDataset(filename) mh = ncd[:mahalanobis] lp = ncd[:posterior_logpdf_true_v_map] @@ -59,13 +71,33 @@ lp = ncd[:posterior_logpdf_true_v_map] ens_vals = try Int.(ncd["ensemble_size"][:]) catch; collect(1:n_ens_size) end kiter_vals = try Int.(ncd["k_iter"][:]) catch; collect(1:n_k_iter) end +# Load truth and observation-noise covariance from the prelim JLD2 (calibrate output). +# Falls back gracefully if the file is absent. +if isfile(prelim_filename) + @info "using precomputed quantities from $(prelim_filename)" + _pd = JLD2.load(prelim_filename) + y_truth = Float64.(_pd["y"]) + R_obs = Float64.(_pd["R"]) + @info "Loaded y and R from $(prelim_jld2_file)" +elseif haskey(ncd, "truth_output") + y_truth = Float64.(ncd["truth_output"][:]) + R_obs = nothing +else + @warn "Prelim file not found at $(prelim_filename) and truth_output absent from NC; R-whitened coverage disabled." + y_truth = nothing + R_obs = nothing +end + # Discover what is actually present in this file avail_spaces = Symbol[:param] avail_metrics = Symbol[:mahalanobis, :logratio] -(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage")) && push!(avail_spaces, :forcing) -(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage")) && push!(avail_spaces, :output) +(haskey(ncd, "forcing_mahalanobis") || haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "forcing_coverage") || (haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing"))) && push!(avail_spaces, :forcing) +(haskey(ncd, "output_mahalanobis") || haskey(ncd, "output_plr_mahalanobis_top") || haskey(ncd, "output_coverage") || (haskey(ncd, "output_samples") && !isnothing(y_truth))) && push!(avail_spaces, :output) (haskey(ncd, "forcing_plr_mahalanobis_top") || haskey(ncd, "output_plr_mahalanobis_top")) && push!(avail_metrics, :plr_mahalanobis) -(haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) +((haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing")) || (haskey(ncd, "output_samples") && !isnothing(y_truth)) || haskey(ncd, "forcing_coverage") || haskey(ncd, "output_coverage")) && push!(avail_metrics, :coverage) +(haskey(ncd, "output_coverage") || haskey(ncd, "forcing_coverage")) && push!(avail_metrics, :budget_for_coverage) + +avail_quantiles = haskey(ncd, "coverage_quantile") ? Float64.(ncd["coverage_quantile"][:]) : nothing println("="^72) println("File: $(filename)") @@ -74,11 +106,13 @@ println("All available options:") println(" display_spaces — $(avail_spaces)") println(" display_metrics — $(avail_metrics)") println(" display_ens_sizes — $(ens_vals)") +isnothing(avail_quantiles) || println(" coverage_quantile — $(avail_quantiles) ($(length(avail_quantiles)) levels)") println() println("Currently selected (edit display_* variables above to filter):") println(" display_spaces = $(isnothing(display_spaces) ? "nothing → $(avail_spaces)" : display_spaces)") println(" display_metrics = $(isnothing(display_metrics) ? "nothing → $(avail_metrics)" : display_metrics)") println(" display_ens_sizes = $(isnothing(display_ens_sizes) ? "nothing → $(ens_vals)" : display_ens_sizes)") +isnothing(avail_quantiles) || println(" calibration_check_quantiles = $(calibration_check_quantiles) (used to select quantile subset for budget metric)") println("="^72) mh_nonmiss = sum(.!ismissing.(mh), dims=1)[1,:,:] @@ -429,52 +463,143 @@ end ########################################################################### #################### Marginal coverage fractions ######################### ########################################################################### +# Coverage is computed on-the-fly from raw pushforward samples and truth, +# allowing any quantile level. When R_obs is loaded from the prelim JLD2, +# an additional R-whitened (PCA) coverage is computed: each coordinate is +# decorrelated by the eigenvectors of R and scaled by the inverse sqrt of +# the corresponding eigenvalue, making dimensions more independent. if show_metric(:coverage) - for (space_label, space_sym, cov_sym) in [ - ("Forcing", :forcing, :forcing_coverage), - ("Output", :output, :output_coverage), - ] - (show_space(space_sym) && haskey(ncd, cov_sym)) || continue - cov = ncd[cov_sym] # (n_rng, n_ens_size, n_k_iter, n_cov_q) - cov_qs_all = Float64.(ncd["coverage_quantile"][:]) - qi_keep = findall(qp -> any(isapprox(qp, p, atol=1e-8) for p in calibration_check_quantiles), cov_qs_all) - cov_qs = cov_qs_all[qi_keep] + # ── Output-space coverage ────────────────────────────────────────────── + if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) + + n_out = ncd.dim["output_dim"] + os_all = coalesce.(Array(ncd["output_samples"]), NaN) # (n_rng, n_ens_size, n_k_iter, n_ps, n_out) + actual_out_dim = size(os_all, 5) + actual_out_dim != n_out && @warn "output_samples trailing dim ($actual_out_dim) ≠ output_dim ($n_out); raw coverage uses actual dim, R-whitened coverage disabled" + + do_whiten = !isnothing(R_obs) + if do_whiten + eig_R = eigen(Symmetric(R_obs)) + ord = sortperm(eig_R.values; rev=true) # descending eigenvalue order + λ_all = eig_R.values[ord] + V_all = eig_R.vectors[:, ord] + cum_var = cumsum(λ_all) ./ sum(λ_all) + k_R = something(findfirst(>=(R_variance_retain), cum_var), length(λ_all)) + V_R = V_all[:, 1:k_R] + λ_R = λ_all[1:k_R] + yw = V_R' * y_truth ./ sqrt.(λ_R) + println(@sprintf("\nR eigenvalue truncation: %d / %d modes retained (threshold %.1f%% → %.4f%% variance)", + k_R, n_out, 100*R_variance_retain, 100*cum_var[k_R])) + end + + for (label, do_white) in [("raw", false), ("R-whitened PCA", true)] + (do_white && !do_whiten) && continue + + local cov_arr = fill(NaN, n_rng, n_ens_size, n_k_iter, length(calibration_check_quantiles)) + for ri in 1:n_rng, ei in 1:n_ens_size, ki in 1:n_k_iter + os = os_all[ri, ei, ki, :, :] # (n_ps, n_out) + all(isnan.(os)) && continue + if do_white + size(os, 2) != size(V_R, 1) && continue + sw = Matrix((V_R' * Matrix(os')) ./ sqrt.(λ_R))' # (n_ps, k_R) + truth_vec = yw + n_cov_dim = k_R + else + sw = os + truth_vec = y_truth + n_cov_dim = actual_out_dim + end + for (qi, qp) in enumerate(calibration_check_quantiles) + cov_arr[ri, ei, ki, qi] = mean(truth_vec[d] <= quantile(sw[:, d], qp) for d in 1:n_cov_dim) + end + end + + println("\n" * "="^72) + println("Output-space marginal coverage fractions [$(label)]") + if do_white + println(" x̃_d = (Vᵀx)_d / √λ_d (R = VΛVᵀ, top $(k_R)/$(n_out) modes, $(round(Int, 100*R_variance_retain))% var threshold)") + println(" Coverage evaluated over effective dimension = $(k_R) (full output dim = $(n_out))") + end + println(" (dims as trials; expected mean_coverage ≈ q_prob for calibrated posterior)") + println("="^72) + println("Pooled mean coverage across all experiments:") + for (qi, qp) in enumerate(calibration_check_quantiles) + pool = collect(filter(!isnan, vec(cov_arr[:, :, :, qi]))) + isempty(pool) && continue + println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) + end + + local df_cov = DataFrame( + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mean_coverage=Union{Float64,Missing}[], std_coverage=Union{Float64,Missing}[], n_valid=Int[], + ) + for (qi, qp) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = filter(!isnan, vec(cov_arr[:, ei, ki, qi])) + push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, + mean_coverage = isempty(vals) ? missing : mean(vals), + std_coverage = length(vals) > 1 ? std(vals) : missing, + n_valid = length(vals))) + end + + println("\nOutput-space coverage by ens_size and k_iter:") + for ens in filter(show_ens, sort(unique(df_cov.ens_size))) + println("\n=== ens_size = $(ens) ===") + for (qi, qp) in enumerate(calibration_check_quantiles) + sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) + println(" q = $(qp) (expected ≈ $(qp))") + show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) + println() + end + end + end + end + + # ── Forcing-space coverage ───────────────────────────────────────────── + if show_space(:forcing) && haskey(ncd, "forcing_samples") && haskey(ncd, "truth_forcing") + + n_for = ncd.dim["forcing_dim"] + fs_all = coalesce.(Array(ncd["forcing_samples"]), NaN) # (n_rng, n_ens_size, n_k_iter, n_ps, n_for) + tf_all = coalesce.(Array(ncd["truth_forcing"]), NaN) # (n_rng, n_ens_size, n_for) + + cov_arr = fill(NaN, n_rng, n_ens_size, n_k_iter, length(calibration_check_quantiles)) + for ri in 1:n_rng, ei in 1:n_ens_size, ki in 1:n_k_iter + fs = fs_all[ri, ei, ki, :, :] + truth_f = tf_all[ri, ei, :] + all(isnan.(fs)) && continue + for (qi, qp) in enumerate(calibration_check_quantiles) + cov_arr[ri, ei, ki, qi] = mean(truth_f[d] <= quantile(fs[:, d], qp) for d in 1:n_for) + end + end println("\n" * "="^72) - println("$(space_label)-space marginal coverage fractions") - println(" (Spatial dims as trials; spatially correlated → effective n < stated n)") + println("Forcing-space marginal coverage fractions [raw]") + println(" (dims as trials; expected mean_coverage ≈ q_prob for calibrated posterior)") println("="^72) println("Pooled mean coverage across all experiments:") - for (qi_local, qi_file) in enumerate(qi_keep) - qp = cov_qs[qi_local] - pool = collect(skipmissing(vec(Array(cov[:, :, :, qi_file])))) + for (qi, qp) in enumerate(calibration_check_quantiles) + pool = collect(filter(!isnan, vec(cov_arr[:, :, :, qi]))) + isempty(pool) && continue println(@sprintf(" q = %.2f mean = %.3f (expected %.3f)", qp, mean(pool), qp)) end df_cov = DataFrame( - q_prob = Float64[], - ens_size = Int[], - k_iter = Int[], - mean_coverage = Union{Float64,Missing}[], - std_coverage = Union{Float64,Missing}[], - n_valid = Int[], + q_prob=Float64[], ens_size=Int[], k_iter=Int[], + mean_coverage=Union{Float64,Missing}[], std_coverage=Union{Float64,Missing}[], n_valid=Int[], ) - for (qi_local, qi_file) in enumerate(qi_keep), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) - qp = cov_qs[qi_local] - vals = collect(skipmissing(cov[:, ei, ki, qi_file])) + for (qi, qp) in enumerate(calibration_check_quantiles), (ei, ev) in enumerate(ens_vals), (ki, kv) in enumerate(kiter_vals) + vals = filter(!isnan, vec(cov_arr[:, ei, ki, qi])) push!(df_cov, (q_prob=qp, ens_size=ev, k_iter=kv, mean_coverage = isempty(vals) ? missing : mean(vals), std_coverage = length(vals) > 1 ? std(vals) : missing, n_valid = length(vals))) end - println("\n$(space_label)-space coverage by ens_size and k_iter:") - println(" (Expected: mean_coverage ≈ q_prob for calibrated marginals)") + println("\nForcing-space coverage by ens_size and k_iter:") for ens in filter(show_ens, sort(unique(df_cov.ens_size))) println("\n=== ens_size = $(ens) ===") - for (qi, qp) in enumerate(cov_qs) + for (qi, qp) in enumerate(calibration_check_quantiles) sub = sort(filter(r -> r.ens_size == ens && r.q_prob == qp, df_cov), :k_iter) println(" q = $(qp) (expected ≈ $(qp))") show(stdout, select(sub, :k_iter, :mean_coverage, :std_coverage, :n_valid), allrows=true) @@ -483,3 +608,294 @@ if show_metric(:coverage) end end end + +########################################################################### +#################### Budget for coverage ################################# +########################################################################### +# For each (c, N_ens) pair: smallest k_iter (per repeat) such that +# |S(q) − q| ≤ α_c(q) = c·√(q(1−q)/N_y) for ALL quantiles q. +# Budget = N_ens·k_iter. NaN = target never reached within the k_iter range. +# Prints mean/median/IQR/min/max; reports 0 for all stats when n_conv = 0. +# Saves one figure per space with one panel per c, box-and-whisker per N_ens. + +if show_metric(:budget_for_coverage) + + ENV["GKSwstype"] = "100" + using Plots + using Plots.Measures + + bq_vals = calibration_check_quantiles + + # ── Table printer ────────────────────────────────────────────────────── + function _bfcov_print_table(budgets_mat, label, n_dim_label) + println("\n$(label) (N_y = $(n_dim_label))") + println(@sprintf(" %-10s %-6s %8s %8s %14s %8s %8s %s", + "N_ens", "c", "mean", "median", "IQR[25,75]", "min", "max", "conv")) + for (ei, ev) in enumerate(ens_vals) + show_ens(ev) || continue + for (si, c) in enumerate(budget_target_scalings) + b = filter(!isnan, vec(budgets_mat[:, ei, si])) + n_conv = length(b) + if n_conv == 0 + println(@sprintf(" %-10d %-6.2f %8.1f %8.1f [%5.0f,%5.0f] %8.1f %8.1f %d/%d", + ev, c, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0, n_rng)) + else + q25, q75 = quantile(b, [0.25, 0.75]) + println(@sprintf(" %-10d %-6.2f %8.1f %8.1f [%5.0f,%5.0f] %8.1f %8.1f %d/%d", + ev, c, mean(b), median(b), q25, q75, minimum(b), maximum(b), n_conv, n_rng)) + end + end + end + end + + # ── Figure builder ───────────────────────────────────────────────────── + function _bfcov_make_figure(budgets_mat, label, n_dim_label; + ylabel_str = "budget (N_ens · k_iter)", + title_prefix = "Budget for coverage") + n_c = length(budget_target_scalings) + n_cols = n_c <= 4 ? n_c : 4 + n_rows = ceil(Int, n_c / n_cols) + dx = 0.20 # half-width of whisker cap ticks in x-index space + + filt_ens = [(ei, ev) for (ei, ev) in enumerate(ens_vals) if show_ens(ev)] + n_x = length(filt_ens) + + p = plot(layout = (2 * n_rows, n_cols), + size = (max(560 * n_cols, 1120), max(640 * n_rows, 700)), + plot_title = "$(title_prefix) — $(label) (N_y = $(n_dim_label))", + plot_titlefontsize = 10, + margin = 6mm) + + # ── Top tier: budget whiskers ────────────────────────────────────── + for (si, c) in enumerate(budget_target_scalings) + b_all_panel = filter(!isnan, vec(budgets_mat[:, :, si])) + y_ref = isempty(b_all_panel) ? 1.0 : max(maximum(b_all_panel), 1.0) + + for (xi, (ei, ev)) in enumerate(filt_ens) + b_all = vec(budgets_mat[:, ei, si]) + b_conv = filter(!isnan, b_all) + n_conv = length(b_conv) + n_tot = n_rng + + if n_conv == 0 + scatter!(p, [xi], [0.0]; subplot = si, color = :red, + marker = :x, markersize = 8, markerstrokewidth = 2, label = false) + y_ann = 0.06 * y_ref + + elseif n_conv == 1 + scatter!(p, [xi], b_conv; subplot = si, color = :steelblue, + marker = :circle, markersize = 6, label = false) + y_ann = b_conv[1] + + elseif n_conv == 2 + mn, mx = minimum(b_conv), maximum(b_conv) + plot!(p, [xi, xi], [mn, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mn, mn]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mx, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + y_ann = mx + + else # n_conv ≥ 3 + mn, mx, med = minimum(b_conv), maximum(b_conv), median(b_conv) + plot!(p, [xi, xi], [mn, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mn, mn]; subplot = si, color = :steelblue, linewidth = 2, label = false) + plot!(p, [xi-dx, xi+dx], [mx, mx]; subplot = si, color = :steelblue, linewidth = 2, label = false) + scatter!(p, [xi], [med]; subplot = si, color = :steelblue, + marker = :circle, markersize = 7, markerstrokecolor = :white, label = false) + y_ann = mx + end + + ann_y = y_ann + 0.07 * y_ref + annotate!(p, [xi], [ann_y], + [Plots.text("$(n_conv)/$(n_tot)", 7, :center, :black)]; subplot = si) + end + + is_left = mod(si - 1, n_cols) == 0 + plot!(p; subplot = si, + title = @sprintf("c = %.2f", c), + titlefontsize = 9, + xlabel = "N_ens", + ylabel = is_left ? ylabel_str : "", + xticks = (collect(1:n_x), string.(last.(filt_ens))), + xlims = (0.5, n_x + 0.5), + xrotation = 30, + legend = false) + end + + # ── Bottom tier: failure-fraction bars ───────────────────────────── + for (si, c) in enumerate(budget_target_scalings) + si_frac = si + n_rows * n_cols + is_left = mod(si - 1, n_cols) == 0 + + frac_fail_vec = [(n_rng - count(!isnan, budgets_mat[:, ei, si])) / n_rng + for (ei, _ev) in filt_ens] + + bar!(p, collect(1:n_x), frac_fail_vec; subplot = si_frac, + bar_width = 0.6, + color = :tomato, + alpha = 0.75, + label = false, + xlabel = "N_ens", + ylabel = is_left ? "frac. failed" : "", + xticks = (collect(1:n_x), string.(last.(filt_ens))), + xlims = (0.5, n_x + 0.5), + ylims = (0.0, 1.05), + yticks = [0.0, 0.5, 1.0], + xrotation = 30, + legend = false) + end + + # ── Hide unused subplots in both tiers ──────────────────────────── + for si in (n_c + 1):(n_rows * n_cols) + plot!(p; subplot = si, axis = false, grid = false, showaxis = false, framestyle = :none) + plot!(p; subplot = si + n_rows * n_cols, axis = false, grid = false, showaxis = false, framestyle = :none) + end + + return p + end + + println("\n" * "="^72) + println("Budget for coverage") + println(" α_c(q) = c·√(q(1−q)/N_y), budget = N_ens·k_iter") + println(" Find min k_iter per repeat where |S(q) − q| ≤ α_c(q) for ALL q") + println(" Quantile grid: $(length(bq_vals)) levels [$(bq_vals[1]), …, $(bq_vals[end])]") + println(" target_scalings c = $(budget_target_scalings)") + println("="^72) + + budget_plot_data = [] # [(label, n_dim_label, budgets_mat), ...] + kiter_plot_data = [] # [(label, n_dim_label, kiters_mat), ...] + nc_tag = splitext(basename(filename))[1] + + # ── Raw coverage from stored NC variables ────────────────────────────── + for (space_label, space_sym, cov_key, dim_key) in [ + ("Output-space raw", :output, "output_coverage", "output_dim"), + ("Forcing-space raw", :forcing, "forcing_coverage", "forcing_dim"), + ] + (show_space(space_sym) && haskey(ncd, cov_key) && haskey(ncd, dim_key)) || continue + + n_dim_bud = ncd.dim[dim_key] + + # Subset coverage to the requested quantiles; fall back to full grid when + # coverage_quantile coordinate is absent (old files). + can_load = haskey(ncd, "coverage_quantile") + bq_idx = nothing + if can_load + bq_nc = Float64.(ncd["coverage_quantile"][:]) + bq_idx = [findfirst(q -> abs(q - qc) < 1e-10, bq_nc) for qc in bq_vals] + can_load = all(!isnothing, bq_idx) + end + + if can_load + @info "Computing $(space_label) budget from NC (quantile subset: $(bq_vals))" + cov_bq = Array(ncd[cov_key])[:, :, :, bq_idx] # (n_rng, n_ens, n_k, n_bq) + budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = coalesce.(cov_bq[ri, ei, ki, :], NaN) + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_raw[ri, ei, si] = ev * kv + kiters_raw[ri, ei, si] = kv + break + end + end + end + end + else + cov_nc = Array(ncd[cov_key]) # (n_rng, n_ens_size, n_k_iter, n_cov_q) + budgets_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_raw = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_bud) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = coalesce.(cov_nc[ri, ei, ki, :], NaN) + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_raw[ri, ei, si] = ev * kv + kiters_raw[ri, ei, si] = kv + break + end + end + end + end + end + _bfcov_print_table(budgets_raw, space_label, n_dim_bud) + push!(budget_plot_data, (space_label, string(n_dim_bud), budgets_raw)) + push!(kiter_plot_data, (space_label, string(n_dim_bud), kiters_raw)) + end + + # ── R-whitened output coverage (recomputed from raw samples) ────────── + if show_space(:output) && haskey(ncd, "output_samples") && !isnothing(y_truth) && !isnothing(R_obs) + n_out_bud = ncd.dim["output_dim"] + os_all_bud = coalesce.(Array(ncd["output_samples"]), NaN) + + eig_Rb = eigen(Symmetric(R_obs)) + ord_b = sortperm(eig_Rb.values; rev=true) + λ_all_b = eig_Rb.values[ord_b] + V_all_b = eig_Rb.vectors[:, ord_b] + cum_var_b = cumsum(λ_all_b) ./ sum(λ_all_b) + k_Rb = something(findfirst(>=(R_variance_retain), cum_var_b), length(λ_all_b)) + V_Rb = V_all_b[:, 1:k_Rb] + λ_Rb = λ_all_b[1:k_Rb] + yw_b = V_Rb' * y_truth ./ sqrt.(λ_Rb) + n_dim_wb = k_Rb + + n_bq = length(bq_vals) + wh_cov_b = fill(NaN, n_rng, n_ens_size, n_k_iter, n_bq) + for ri in 1:n_rng, (ei, _ev) in enumerate(ens_vals), ki in 1:n_k_iter + os = os_all_bud[ri, ei, ki, :, :] + all(isnan.(os)) && continue + size(os, 2) != size(V_Rb, 1) && continue + sw = Matrix((V_Rb' * Matrix(os')) ./ sqrt.(λ_Rb))' # (n_ps, k_Rb) + for (qi, qp) in enumerate(bq_vals) + wh_cov_b[ri, ei, ki, qi] = mean(yw_b[d] <= quantile(sw[:, d], qp) for d in 1:n_dim_wb) + end + end + + budgets_wh = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + kiters_wh = fill(NaN, n_rng, n_ens_size, length(budget_target_scalings)) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(bq_vals .* (1 .- bq_vals) ./ n_dim_wb) + for (ei, ev) in enumerate(ens_vals), ri in 1:n_rng + for (ki, kv) in enumerate(kiter_vals) + s_q = wh_cov_b[ri, ei, ki, :] + all(isnan.(s_q)) && continue + if all(abs.(s_q .- bq_vals) .<= tol) + budgets_wh[ri, ei, si] = ev * kv + kiters_wh[ri, ei, si] = kv + break + end + end + end + end + wh_label = "Output-space R-whitened" + wh_ndim = "$(n_dim_wb) eff. dims ($(round(Int, 100*R_variance_retain))% var of R)" + _bfcov_print_table(budgets_wh, wh_label, wh_ndim) + push!(budget_plot_data, (wh_label, wh_ndim, budgets_wh)) + push!(kiter_plot_data, (wh_label, wh_ndim, kiters_wh)) + end + + # ── Generate and save figures ────────────────────────────────────────── + for (label, n_dim_label, bmat) in budget_plot_data + p_fig = _bfcov_make_figure(bmat, label, n_dim_label) + tag = replace(lowercase(label), r"[^a-z0-9]+" => "_") + base = joinpath(indir, "budget_for_coverage_$(tag)_$(nc_tag)") + savefig(p_fig, base * ".pdf") + savefig(p_fig, base * ".png") + @info "Saved: $(base).pdf" + end + + for (label, n_dim_label, kmat) in kiter_plot_data + p_fig = _bfcov_make_figure(kmat, label, n_dim_label; + ylabel_str = "k_iter", + title_prefix = "Iterations for coverage") + tag = replace(lowercase(label), r"[^a-z0-9]+" => "_") + base = joinpath(indir, "iters_for_coverage_$(tag)_$(nc_tag)") + savefig(p_fig, base * ".pdf") + savefig(p_fig, base * ".png") + @info "Saved: $(base).pdf" + end +end diff --git a/examples/EKIRace/exp_to_leaderboard.jl b/examples/EKIRace/exp_to_leaderboard.jl new file mode 100644 index 000000000..b927d01af --- /dev/null +++ b/examples/EKIRace/exp_to_leaderboard.jl @@ -0,0 +1,418 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +n_lowrank_modes = 5 +marginal_coverage_quantiles = collect(0.05:0.05:0.95) +n_marginal_coverage_quantiles = length(marginal_coverage_quantiles) +budget_target_scalings = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5] +n_target_scalings = length(budget_target_scalings) + +########################################################################### +#################### Config setup ######################################## +########################################################################### + +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) +has_forcing = cfg.force_case !== nothing + +nc_save_filename = nc_filename(cfg, method) + +########################################################################### +#################### Locate valid posterior files ######################## +########################################################################### + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss, rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end +end + +@info "Converting data from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample first.") +end + +########################################################################### +#################### Determine dimensions ################################ +########################################################################### + +first_post_fn = joinpath(data_save_directory, posterior_filename(cfg, valid_file_items[1]...)) +first_loaded = JLD2.load(first_post_fn) +if !haskey(first_loaded, "pushforward_output_samples") + error("Pushforward data not found in $(first_post_fn). Run pushforward_from_posterior_$(cfg.model).jl first.") +end + +n_params = length(vec(mean(first_loaded["posteriors_by_k"][1]))) +n_pushforward_samples = first_loaded["pushforward_n_samples"] +n_output = size(first_loaded["pushforward_output_samples"], 2) +n_forcing = has_forcing ? size(first_loaded["pushforward_forcing_samples"], 2) : 0 + +n_k = maximum( + maximum(JLD2.load(joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)))["k_values"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +########################################################################### +#################### Load truth observation vector ####################### +########################################################################### + +prelim_file = if cfg.force_case === nothing + joinpath(homedir, "output", "$(cfg.model)_computed_preliminaries.jld2") +else + joinpath(homedir, "output", "$(cfg.model)_computed_preliminaries_$(cfg.force_case).jld2") +end +isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate first.") +y = JLD2.load(prelim_file, "y") + +########################################################################### +#################### Pre-allocate arrays ################################# +########################################################################### + +true_param_arr = fill(NaN, n_rng, n_ens, n_params) +n_evals_arr = fill(NaN, n_rng, n_ens) +post_mean_arr = fill(NaN, n_rng, n_ens, n_k, n_params) +post_cov_arr = fill(NaN, n_rng, n_ens, n_k, n_params, n_params) +mahal_arr = fill(NaN, n_rng, n_ens, n_k) +logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_output) +output_mahal_arr = fill(NaN, n_rng, n_ens, n_k) +output_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) +output_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) +output_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) + +if has_forcing + forcing_samples_arr = fill(NaN, n_rng, n_ens, n_k, n_pushforward_samples, n_forcing) + forcing_mahal_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_logpdf_true_v_map_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_plr_mahal_top_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_plr_mahal_residual_arr = fill(NaN, n_rng, n_ens, n_k) + forcing_coverage_arr = fill(NaN, n_rng, n_ens, n_k, n_marginal_coverage_quantiles) + truth_forcing_arr = fill(NaN, n_rng, n_ens, n_forcing) +end + +########################################################################### +#################### Main loop over posterior files ###################### +########################################################################### + +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded = JLD2.load(joinpath(data_save_directory, post_fn)) + + if !haskey(loaded, "pushforward_output_samples") + @warn "Pushforward data missing for $(post_fn); skipping. Run pushforward_from_posterior_$(cfg.model).jl first." + continue + end + + posteriors_by_k = loaded["posteriors_by_k"] + k_values = loaded["k_values"] + truth_params = loaded["truth_params"] + + pf_output = loaded["pushforward_output_samples"] # (n_samples, n_output, n_k_pos) + pf_k_values = loaded["pushforward_k_values"] + + if has_forcing + pf_forcing = loaded["pushforward_forcing_samples"] # (n_samples, n_forcing, n_k_pos) + truth_forcing_vec = loaded["truth_forcing"] + end + + ekp_loaded = JLD2.load(joinpath(data_save_directory, ekp_filename(cfg, N_ens, rng_idx))) + conv_alg_iters = length(get_g(ekp_loaded["ekpobj"])) + + i = findfirst(==(rng_idx), rng_idxs) + j = findfirst(==(N_ens), N_enss) + + true_param_arr[i, j, :] = truth_params + n_evals_arr[i, j] = conv_alg_iters * N_ens + if has_forcing + truth_forcing_arr[i, j, :] = truth_forcing_vec + end + + for k in k_values + post_dist = posteriors_by_k[k] + pm = vec(mean(post_dist)) + pc = cov(post_dist) + C_reg = Symmetric(pc + 1e-10 * I) + post_normal = MvNormal(pm, C_reg) + post_samples = reduce(vcat, [get_distribution(post_dist)[name] for name in get_name(post_dist)]) + pmode = post_samples[:, argmax(logpdf(post_normal, post_samples))] + diff = pm - truth_params + + num_samples = size(post_samples, 2) + r = rank(pc) + if r == num_samples - 1 && r < n_params - 1 + @warn "Posterior covariance rank $(r) = num_samples-1 = $(num_samples-1) < n_params-1 = $(n_params-1). Metric may be inaccurate; recommend num_samples > $(n_params)." + end + + post_mean_arr[i, j, k, :] = pm + post_cov_arr[i, j, k, :, :] = pc + mahal_arr[i, j, k] = diff' * (C_reg \ diff) + logpdf_true_v_map_arr[i, j, k] = logpdf(post_normal, truth_params) - logpdf(post_normal, pmode) + + ki = findfirst(==(k), pf_k_values) + + # ── Output-space metrics ────────────────────────────────────────── + output_samples_arr[i, j, k, :, :] = pf_output[:, :, ki] + os = output_samples_arr[i, j, k, :, :] + om = vec(mean(os, dims=1)) + oc = Symmetric(cov(os) + 1e-10 * I) + o_normal = MvNormal(om, oc) + o_cols = Matrix(os') + o_mode = o_cols[:, argmax(logpdf(o_normal, o_cols))] + o_diff = om - y + output_mahal_arr[i, j, k] = o_diff' * (oc \ o_diff) + output_logpdf_true_v_map_arr[i, j, k] = logpdf(o_normal, y) - logpdf(o_normal, o_mode) + Fo = eigen(oc) + a_o = mean(Fo.values[1:end-n_lowrank_modes]) + V_o = Fo.vectors[:, end-n_lowrank_modes+1:end] + λ_o = Fo.values[end-n_lowrank_modes+1:end] + if a_o < 1e-8 + @warn "Output-space PLR skipped: noise floor a_o=$(a_o) ≈ regularization level. Entries left as NaN." + else + proj_o = V_o' * o_diff + output_plr_mahal_top_arr[i, j, k] = sum(proj_o.^2 ./ λ_o) + output_plr_mahal_residual_arr[i, j, k] = (sum(o_diff.^2) - sum(proj_o.^2)) / a_o + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + output_coverage_arr[i, j, k, qi] = mean(y .<= [quantile(os[:, d], qp) for d in 1:n_output]) + end + + # ── Forcing-space metrics (L96 only) ───────────────────────────── + if has_forcing + forcing_samples_arr[i, j, k, :, :] = pf_forcing[:, :, ki] + fs = forcing_samples_arr[i, j, k, :, :] + fm = vec(mean(fs, dims=1)) + fc = Symmetric(cov(fs) + 1e-10 * I) + f_normal = MvNormal(fm, fc) + f_cols = Matrix(fs') + f_mode = f_cols[:, argmax(logpdf(f_normal, f_cols))] + f_diff = fm - truth_forcing_vec + forcing_mahal_arr[i, j, k] = f_diff' * (fc \ f_diff) + forcing_logpdf_true_v_map_arr[i, j, k] = logpdf(f_normal, truth_forcing_vec) - logpdf(f_normal, f_mode) + Ff = eigen(fc) + a_f = mean(Ff.values[1:end-n_lowrank_modes]) + V_f = Ff.vectors[:, end-n_lowrank_modes+1:end] + λ_f = Ff.values[end-n_lowrank_modes+1:end] + if a_f < 1e-8 + @warn "Forcing-space PLR skipped: noise floor a_f=$(a_f) ≈ regularization level (e.g. const-force). Entries left as NaN." + else + proj_f = V_f' * f_diff + forcing_plr_mahal_top_arr[i, j, k] = sum(proj_f.^2 ./ λ_f) + forcing_plr_mahal_residual_arr[i, j, k] = (sum(f_diff.^2) - sum(proj_f.^2)) / a_f + end + for (qi, qp) in enumerate(marginal_coverage_quantiles) + forcing_coverage_arr[i, j, k, qi] = mean(truth_forcing_vec .<= [quantile(fs[:, d], qp) for d in 1:n_forcing]) + end + end + end +end + +########################################################################### +#################### Budget for coverage ################################# +########################################################################### + +output_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) +output_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) +for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_output) + for ri in 1:n_rng, ei in 1:n_ens, (qi, qp) in enumerate(marginal_coverage_quantiles) + for k in 1:n_k + s = output_coverage_arr[ri, ei, k, qi] + isnan(s) && continue + if abs(s - qp) <= tol[qi] + output_budget_to_target[ri, ei, si, qi] = N_enss[ei] * k + output_iters_to_target[ri, ei, si, qi] = k + break + end + end + end +end + +if has_forcing + forcing_budget_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) + forcing_iters_to_target = fill(NaN, n_rng, n_ens, n_target_scalings, n_marginal_coverage_quantiles) + for (si, c) in enumerate(budget_target_scalings) + tol = c .* sqrt.(marginal_coverage_quantiles .* (1 .- marginal_coverage_quantiles) ./ n_forcing) + for ri in 1:n_rng, ei in 1:n_ens, (qi, qp) in enumerate(marginal_coverage_quantiles) + for k in 1:n_k + s = forcing_coverage_arr[ri, ei, k, qi] + isnan(s) && continue + if abs(s - qp) <= tol[qi] + forcing_budget_to_target[ri, ei, si, qi] = N_enss[ei] * k + forcing_iters_to_target[ri, ei, si, qi] = k + break + end + end + end + end +end + +########################################################################### +#################### Save to NetCDF ##################################### +########################################################################### + +ds = NCDataset(nc_save_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "param_dim", n_params) +defDim(ds, "param_dim_2", n_params) +defDim(ds, "pushforward_sample", n_pushforward_samples) +defDim(ds, "output_dim", n_output) +defDim(ds, "coverage_quantile", n_marginal_coverage_quantiles) +defDim(ds, "target_scaling", n_target_scalings) +if has_forcing + defDim(ds, "forcing_dim", n_forcing) +end + +# ── Coordinate variables ────────────────────────────────────────────── +rng_var = defVar(ds, "random_seed", Int64, ("random_seed",)) +rng_var[:] = rng_idxs + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = N_enss + +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = collect(1:n_k) + +cov_q_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cov_q_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cov_q_var[:] = marginal_coverage_quantiles + +ts_var = defVar(ds, "target_scaling", Float64, ("target_scaling",)) +ts_var.attrib["description"] = "Scaling c in α_c(q) = c·√(q(1−q)/N_y); tolerance for budget_to_target / iters_to_target" +ts_var[:] = budget_target_scalings + +# ── Calibration cost and truth ──────────────────────────────────────── +n_evals_v = defVar(ds, "n_evals_to_target", Float64, ("random_seed", "ensemble_size"); fillvalue=NaN) +n_evals_v.attrib["description"] = "Total calibration cost: conv_alg_iters * ensemble_size." +n_evals_v[:, :] = n_evals_arr + +true_param_v = defVar(ds, "true_param", Float64, ("random_seed", "ensemble_size", "param_dim"); fillvalue=NaN) +true_param_v.attrib["description"] = "truth_param" +true_param_v[:, :, :] = true_param_arr + +if has_forcing + truth_forcing_v = defVar(ds, "truth_forcing", Float64, ("random_seed", "ensemble_size", "forcing_dim"); fillvalue=NaN) + truth_forcing_v.attrib["description"] = "True forcing vector for each (random_seed, ensemble_size) pair." + truth_forcing_v[:, :, :] = truth_forcing_arr +end + +# ── Param-space metrics ─────────────────────────────────────────────── +post_mean_v = defVar(ds, "post_mean", Float64, ("random_seed", "ensemble_size", "k_iter", "param_dim"); fillvalue=NaN) +post_mean_v.attrib["description"] = "mean(posterior) for each emulator training size k" +post_mean_v[:, :, :, :] = post_mean_arr + +post_cov_v = defVar(ds, "post_cov", Float64, ("random_seed", "ensemble_size", "k_iter", "param_dim", "param_dim_2"); fillvalue=NaN) +post_cov_v.attrib["description"] = "cov(posterior) for each emulator training size k" +post_cov_v[:, :, :, :, :] = post_cov_arr + +mahalanobis_v = defVar(ds, "mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +mahalanobis_v.attrib["description"] = "M = (m-truth)' C^{-1} (m-truth) in parameter space." +mahalanobis_v[:, :, :] = mahal_arr + +posterior_logpdf_true_v_map_v = defVar(ds, "posterior_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +posterior_logpdf_true_v_map_v.attrib["description"] = "P = logpdf(posterior, truth_param) - logpdf(posterior, mode). For Gaussian, -2P = M." +posterior_logpdf_true_v_map_v[:, :, :] = logpdf_true_v_map_arr + +# ── Output-space metrics ────────────────────────────────────────────── +output_samples_v = defVar(ds, "output_samples", Float64, ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "output_dim"); fillvalue=NaN) +output_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples in output space" +output_samples_v[:, :, :, :, :] = output_samples_arr + +output_mahal_v = defVar(ds, "output_mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_mahal_v.attrib["description"] = "Mahalanobis distance in output space: (om-y)' Co^{-1} (om-y)." +output_mahal_v[:, :, :] = output_mahal_arr + +output_logpdf_true_v_map_v = defVar(ds, "output_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in output space: logpdf(N(om,Co), y) - logpdf(N(om,Co), mode)." +output_logpdf_true_v_map_v[:, :, :] = output_logpdf_true_v_map_arr + +output_plr_top_v = defVar(ds, "output_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_top_v.attrib["description"] = "PLR Mahalanobis top term in output space: sum((Vo'*(om-y))^2 / λo_i), top $(n_lowrank_modes) modes. Ref: Chisq($(n_lowrank_modes))." +output_plr_top_v[:, :, :] = output_plr_mahal_top_arr + +output_plr_res_v = defVar(ds, "output_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) +output_plr_res_v.attrib["description"] = "PLR Mahalanobis residual term in output space: (||om-y||^2 - ||Vo'*(om-y)||^2) / a_o. Ref: Chisq($(n_output - n_lowrank_modes))." +output_plr_res_v[:, :, :] = output_plr_mahal_residual_arr + +output_coverage_v = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +output_coverage_v.attrib["description"] = "Marginal coverage in output space: fraction of output dims where y[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." +output_coverage_v[:, :, :, :] = output_coverage_arr + +output_budget_v = defVar(ds, "output_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) +output_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) per quantile q in output space. NaN = not reached. Take max over quantiles for the all-quantile condition." +output_budget_v[:, :, :, :] = output_budget_to_target + +output_iters_v = defVar(ds, "output_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) +output_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target per quantile in output space. NaN = not reached." +output_iters_v[:, :, :, :] = output_iters_to_target + +# ── Forcing-space metrics (L96 only) ───────────────────────────────── +if has_forcing + forcing_samples_v = defVar(ds, "forcing_samples", Float64, ("random_seed", "ensemble_size", "k_iter", "pushforward_sample", "forcing_dim"); fillvalue=NaN) + forcing_samples_v.attrib["description"] = "$(n_pushforward_samples) posterior pushforward samples in forcing space" + forcing_samples_v[:, :, :, :, :] = forcing_samples_arr + + forcing_mahal_v = defVar(ds, "forcing_mahalanobis", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_mahal_v.attrib["description"] = "Mahalanobis distance in forcing space: (fm-truth_forcing)' Cf^{-1} (fm-truth_forcing)." + forcing_mahal_v[:, :, :] = forcing_mahal_arr + + forcing_logpdf_true_v_map_v = defVar(ds, "forcing_logpdf_true_v_map", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_logpdf_true_v_map_v.attrib["description"] = "Log PDF ratio in forcing space: logpdf(N(fm,Cf), truth_forcing) - logpdf(N(fm,Cf), mode)." + forcing_logpdf_true_v_map_v[:, :, :] = forcing_logpdf_true_v_map_arr + + forcing_plr_top_v = defVar(ds, "forcing_plr_mahalanobis_top", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_plr_top_v.attrib["description"] = "PLR Mahalanobis top term in forcing space: sum((Vf'*(fm-truth))^2 / λf_i), top $(n_lowrank_modes) modes. Ref: Chisq($(n_lowrank_modes))." + forcing_plr_top_v[:, :, :] = forcing_plr_mahal_top_arr + + forcing_plr_res_v = defVar(ds, "forcing_plr_mahalanobis_residual", Float64, ("random_seed", "ensemble_size", "k_iter"); fillvalue=NaN) + forcing_plr_res_v.attrib["description"] = "PLR Mahalanobis residual term in forcing space: (||fm-truth||^2 - ||Vf'*(fm-truth)||^2) / a_f. Ref: Chisq($(n_forcing - n_lowrank_modes))." + forcing_plr_res_v[:, :, :] = forcing_plr_mahal_residual_arr + + forcing_coverage_v = defVar(ds, "forcing_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) + forcing_coverage_v.attrib["description"] = "Marginal coverage in forcing space: fraction of forcing dims where truth_forcing[d] ≤ q_p of the $(n_pushforward_samples) pushforward samples." + forcing_coverage_v[:, :, :, :] = forcing_coverage_arr + + forcing_budget_v = defVar(ds, "forcing_budget_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) + forcing_budget_v.attrib["description"] = "Budget (N_ens·k_iter) to first reach |S(q)−q| ≤ c·√(q(1−q)/N_y) per quantile q in forcing space. NaN = not reached. Take max over quantiles for the all-quantile condition." + forcing_budget_v[:, :, :, :] = forcing_budget_to_target + + forcing_iters_v = defVar(ds, "forcing_iters_to_target", Float64, ("random_seed", "ensemble_size", "target_scaling", "coverage_quantile"); fillvalue=NaN) + forcing_iters_v.attrib["description"] = "Iterations k_iter to first reach coverage target per quantile in forcing space. NaN = not reached." + forcing_iters_v[:, :, :, :] = forcing_iters_to_target +end + +close(ds) +@info "Saved leaderboard data to $(nc_save_filename)" diff --git a/examples/EKIRace/experiment_config.jl b/examples/EKIRace/experiment_config.jl index cee4c77c4..350cd927e 100644 --- a/examples/EKIRace/experiment_config.jl +++ b/examples/EKIRace/experiment_config.jl @@ -3,12 +3,14 @@ using Dates ######################################################################## ############### USER TOGGLE ######################################### ######################################################################## -# Set EXPERIMENT to one of: :l63, :l96_const, :l96_vec, :l96_flux -EXPERIMENT = :l96_flux +# Set EXPERIMENT to one of: +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[4] # Date identifying this calibration run — written by calibrate, read by # emulate_sample and exp_to_leaderboard (so all stages stay in sync). -#calibrate_date = Date("2026-05-31", "yyyy-mm-dd") +# PIN this to reproduce a specific run directory. +#calibrate_date = Date("2026-06-03", "yyyy-mm-dd") calibrate_date = today() ######################################################################## @@ -46,14 +48,17 @@ forcing_cases_key = Dict( # n_repeats: number of rng seeds function experiment_config(case::Symbol) + n_ens_step = 8 + n_repeats = 20 if case == :l63 + ens_step = 2 return ( model = "l63", force_case = nothing, - N_ens_sizes = [10, 25, 40], + N_ens_sizes = collect(4:ens_step:4+n_ens_step*ens_step), N_iter = 20, terminate_at = 2.0, # DataMisfitController end time - n_repeats = 20, + n_repeats = n_repeats, max_iter = 10, retain_var = 0.99, n_features = 100, @@ -61,13 +66,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_const + ens_step_const = 2 return ( model = "l96", force_case = "const-force", - N_ens_sizes = [5, 15, 30], + N_ens_sizes = collect(4:ens_step_const:4+n_ens_step*ens_step_const), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, @@ -75,13 +81,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_vec + ens_step_vec = 5 return ( model = "l96", force_case = "vec-force", - N_ens_sizes = [50, 75, 100], + N_ens_sizes = collect(40:ens_step_vec:40+n_ens_step*ens_step_vec), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, @@ -89,13 +96,14 @@ function experiment_config(case::Symbol) calibrate_date = calibrate_date, ) elseif case == :l96_flux + ens_step_flux = 5 return ( model = "l96", force_case = "flux-force", - N_ens_sizes = [50, 75, 100], + N_ens_sizes = collect(30:ens_step_flux:30+n_ens_step*ens_step_flux), N_iter = 20, terminate_at = 2.0, - n_repeats = 20, + n_repeats = n_repeats, max_iter = 15, retain_var = 0.99, n_features = 200, @@ -148,3 +156,4 @@ function nc_filename(cfg, method) return "$(key)_$(cfg.model)_$(forcing_cases_key[cfg.force_case])_$(cfg.calibrate_date).nc" end end + diff --git a/examples/EKIRace/minimal_leaderboard_nc.jl b/examples/EKIRace/minimal_leaderboard_nc.jl new file mode 100644 index 000000000..93a5f3c66 --- /dev/null +++ b/examples/EKIRace/minimal_leaderboard_nc.jl @@ -0,0 +1,110 @@ +using NCDatasets +using JLD2 +using LinearAlgebra +using Statistics +using Dates + +include("experiment_config.jl") + +cfg = experiment_config(EXPERIMENT) +method = method_cases[1] + +in_filename = nc_filename(cfg, method) +out_filename = replace(in_filename, r"\.nc$" => "_minimal.nc") +prelim_filename = if cfg.force_case === nothing + joinpath(pwd(), "output", "$(cfg.model)_computed_preliminaries.jld2") +else + joinpath(pwd(), "output", "$(cfg.model)_computed_preliminaries_$(cfg.force_case).jld2") +end + +R_variance_retain = 0.99 # fraction of R eigenvalue variance retained for whitened PCA coverage + +########################################################################### +#################### Read required fields from input NC ################## +########################################################################### + +@info "Reading $(in_filename)" +ncd_in = NCDataset(in_filename) + +n_rng = ncd_in.dim["random_seed"] +n_ens = ncd_in.dim["ensemble_size"] +n_k = ncd_in.dim["k_iter"] +n_cq = ncd_in.dim["coverage_quantile"] +n_ts = ncd_in.dim["target_scaling"] +n_out = ncd_in.dim["output_dim"] + +ens_vals = Int.(ncd_in["ensemble_size"][:]) +k_vals = Int.(ncd_in["k_iter"][:]) +cq_vals = Float64.(ncd_in["coverage_quantile"][:]) +ts_vals = Float64.(ncd_in["target_scaling"][:]) +os_all = coalesce.(Array(ncd_in["output_samples"]), NaN) # (n_rng, n_ens, n_k, n_ps, n_out) + +close(ncd_in) + +########################################################################### +#################### Load preliminaries and compute whitened coverage #### +########################################################################### + +@info "Loading preliminaries from $(prelim_filename)" +_pd = JLD2.load(prelim_filename) +y_truth = Float64.(_pd["y"]) +R_obs = Float64.(_pd["R"]) + +eig_R = eigen(Symmetric(R_obs)) +ord = sortperm(eig_R.values; rev=true) +λ_all = eig_R.values[ord] +V_all = eig_R.vectors[:, ord] +cum_var = cumsum(λ_all) ./ sum(λ_all) +k_R = something(findfirst(>=(R_variance_retain), cum_var), length(λ_all)) +V_R = V_all[:, 1:k_R] +λ_R = λ_all[1:k_R] +yw = V_R' * y_truth ./ sqrt.(λ_R) + +@info "R-whitened PCA: retaining $(k_R)/$(n_out) modes ($(round(100*cum_var[k_R]; digits=2))% variance)" + +output_coverage = fill(NaN, n_rng, n_ens, n_k, n_cq) +for ri in 1:n_rng, ei in 1:n_ens, ki in 1:n_k + os = os_all[ri, ei, ki, :, :] + all(isnan.(os)) && continue + size(os, 2) != size(V_R, 1) && continue + sw = Matrix((V_R' * Matrix(os')) ./ sqrt.(λ_R))' # (n_ps, k_R) + for (qi, qp) in enumerate(cq_vals) + output_coverage[ri, ei, ki, qi] = mean(yw[d] <= quantile(sw[:, d], qp) for d in 1:k_R) + end +end + +########################################################################### +#################### Write minimal NC #################################### +########################################################################### + +ds = NCDataset(out_filename, "c") + +defDim(ds, "random_seed", n_rng) +defDim(ds, "ensemble_size", n_ens) +defDim(ds, "k_iter", n_k) +defDim(ds, "coverage_quantile", n_cq) +defDim(ds, "target_scaling", n_ts) +defDim(ds, "output_dim", k_R) # effective whitened dimension (not full output_dim = $(n_out)) + +ens_var = defVar(ds, "ensemble_size", Int64, ("ensemble_size",)) +ens_var.attrib["description"] = "Number of ensemble members" +ens_var[:] = ens_vals + +k_var = defVar(ds, "k_iter", Int64, ("k_iter",)) +k_var.attrib["description"] = "Number of EKP training iterations used to fit the emulator (1-indexed)" +k_var[:] = k_vals + +cq_var = defVar(ds, "coverage_quantile", Float64, ("coverage_quantile",)) +cq_var.attrib["description"] = "Quantile levels used for marginal coverage fraction metrics" +cq_var[:] = cq_vals + +ts_var = defVar(ds, "target_scaling", Float64, ("target_scaling",)) +ts_var.attrib["description"] = "Scaling c in α_c(q) = c·√(q(1−q)/N_y); tolerance for budget_to_target / iters_to_target" +ts_var[:] = ts_vals + +oc_var = defVar(ds, "output_coverage", Float64, ("random_seed", "ensemble_size", "k_iter", "coverage_quantile"); fillvalue=NaN) +oc_var.attrib["description"] = "R-whitened PCA marginal coverage: fraction of whitened output dims d where ỹ[d] ≤ q_p of whitened pushforward samples. Whitening: x̃_d = (Vᵀx)_d / √λ_d where R = VΛVᵀ. Retained $(k_R)/$(n_out) R-eigenmodes ($(round(100*cum_var[k_R]; digits=1))% variance, threshold $(R_variance_retain)). The output_dim dimension = $(k_R) is the effective whitened N_y used in α_c(q) = c·√(q(1−q)/N_y)." +oc_var[:, :, :, :] = output_coverage + +close(ds) +@info "Saved minimal leaderboard data to $(out_filename)" diff --git a/examples/EKIRace/posterior_diagnostic_plots_l63.jl b/examples/EKIRace/posterior_diagnostic_plots_l63.jl new file mode 100644 index 000000000..32d9f5d98 --- /dev/null +++ b/examples/EKIRace/posterior_diagnostic_plots_l63.jl @@ -0,0 +1,251 @@ +# Import modules +using Distributions +using LinearAlgebra +using Random +using JLD2 +using Statistics +using Plots +using Plots.Measures + +# CES +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("Lorenz63.jl") +include("experiment_config.jl") + +n_samples_pushforward = 1000 + +######################################################################## +################## file-certainty for loading ########################## +######################################################################## + +cfg = experiment_config(:l63) +method = method_cases[1] +calib_dir = calib_directory(method, cfg) +N_enss = cfg.N_ens_sizes +rng_idxs = collect(1:cfg.n_repeats) + +prelim_dir = joinpath(@__DIR__, "output") +if !isdir(prelim_dir) + mkdir(prelim_dir) +end +prelim_file = joinpath(prelim_dir, "l63_computed_preliminaries.jld2") +if isfile(prelim_file) + loaded_data = JLD2.load(prelim_file) + x0 = loaded_data["x0"] + nx = length(x0) + y = loaded_data["y"] + ic_cov_sqrt = loaded_data["ic_cov_sqrt"] + R = loaded_data["R"] + lorenz_config_settings = loaded_data["lorenz_config_settings"] + observation_config = loaded_data["observation_config"] + + @info "loaded precomputed preliminary quantities from $(prelim_file)" +else + throw(ErrorException("preliminaries file not found. \n First run: \n > julia --project calibrate_l63.jl")) +end + +homedir = joinpath(pwd()) +data_save_directory = joinpath(homedir, "output", calib_dir) + +valid_file_items = [] +valid_files = [] +for N_ens in N_enss + for rng_idx in rng_idxs + data_file = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + if isfile(data_file) + push!(valid_files, case_suffix(cfg, N_ens, rng_idx)) + push!(valid_file_items, (N_ens, rng_idx)) + end + end +end + +@info "Pushing forward posteriors through the forward map from valid files:" +display(valid_files) + +if isempty(valid_file_items) + error("No valid posterior files found in $(data_save_directory). Run emulate_sample_l63.jl first.") +end + +function pushforward_metrics(samples::AbstractMatrix, truth::AbstractVector) + m = vec(mean(samples, dims=2)) + C_raw = cov(samples, dims=2) + num_samples = size(samples, 2) + dim_samples = size(samples, 1) + r = rank(C_raw) + if r == num_samples - 1 && r < dim_samples - 1 + @warn "Covariance rank $(r) = num_samples-1 = $(num_samples-1) < dim-1 = $(dim_samples-1). Metric may be inaccurate due to insufficient samples; recommend num_samples > $(dim_samples)." + end + λ = max(1e-10, 1e-4 * mean(diag(C_raw))) + C = Symmetric(C_raw + λ * I) + dist = MvNormal(m, C) + pmode = samples[:, argmax(logpdf(dist, samples))] + diff = m - truth + mah = diff' * (C \ diff) + lp = logpdf(dist, truth) - logpdf(dist, pmode) + return mah, lp +end + +for (N_ens, rng_idx) in valid_file_items + post_fn = posterior_filename(cfg, N_ens, rng_idx) + @info "loading case $(post_fn)" + loaded_p = JLD2.load(joinpath(data_save_directory, post_fn)) + + posteriors_by_k = loaded_p["posteriors_by_k"] + priors = loaded_p["priors"] + k_values = loaded_p["k_values"] + truth_params = loaded_p["truth_params"] + truth_params_constrained = loaded_p["truth_params_constrained"] + + n_par = length(truth_params_constrained) + ny = length(y) + + # --- sample from prior for grey background --- + prior_samples_unconstrained = sample(priors, n_samples_pushforward) + prior_samples_constrained = transform_unconstrained_to_constrained(priors, prior_samples_unconstrained) + prior_param_diffs = reduce(hcat, [prior_samples_constrained[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + G_prior = hcat( + [ + lorenz_forward( + EnsembleMemberConfig(prior_samples_constrained[:, j]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) for j in 1:n_samples_pushforward + ]..., + ) + + if !haskey(loaded_p, "pushforward_output_samples") + error("Pushforward data not found in $(post_fn). Run pushforward_from_posterior_l63.jl first.") + end + pf_output = loaded_p["pushforward_output_samples"] # (n_samples, n_output, n_k) + pf_k_values = loaded_p["pushforward_k_values"] + + for k in k_values + post_dist = posteriors_by_k[k] + + # sample posterior for parameter-space analysis (no forward map needed) + push_ensemble = sample(post_dist, n_samples_pushforward) + constrained_push_ensemble = transform_unconstrained_to_constrained(post_dist, push_ensemble) + + # load precomputed posterior pushforward (produced by pushforward_from_posterior_l63.jl) + ki = findfirst(==(k), pf_k_values) + G_ens = pf_output[:, :, ki]' # (n_output, n_samples) + + param_diffs = reduce(hcat, [constrained_push_ensemble[:, j] - truth_params_constrained for j in 1:n_samples_pushforward]) + + param_mah, param_lp = pushforward_metrics(push_ensemble, truth_params) + output_mah, output_lp = pushforward_metrics(G_ens, y) + lowbd_par, upbd_par = round.(quantile(Chisq(n_par), [0.01, 0.99]), digits=2) + lowbd_out, upbd_out = round.(quantile(Chisq(ny), [0.01, 0.99]), digits=2) + @info "--- Posterior metrics (N_ens=$(N_ens), rng=$(rng_idx), k=$(k)) ---" + @info " param [d=$(n_par)]: target: [$(lowbd_par), $(upbd_par)] mahal=$(round(param_mah, digits=2)) -2*logpdf_ratio=$(round(-2*param_lp, digits=2))" + @info " output [d=$(ny)]: target: [$(lowbd_out), $(upbd_out)] mahal=$(round(output_mah, digits=2)) -2*logpdf_ratio=$(round(-2*output_lp, digits=2))" + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameters (differences from truth) + p2 = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Panel (ii): output + p4 = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Output index", + ylabel = "State statistics (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + + # Add prior samples (grey) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), prior_param_diffs[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 1], label = "prior samples", color = :grey, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_prior[:, 2:end], label = "", color = :grey, linewidth = 4, linealpha = 0.1) + + # Add posterior samples (green) + plot!(p2, collect(1:n_par), param_diffs[:, 1], label = "posterior samples", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p2, collect(1:n_par), param_diffs[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 1], label = "pushforward outputs", color = :lightgreen, linewidth = 4, linealpha = 0.1) + plot!(p4, 1:ny, G_ens[:, 2:end], label = "", color = :lightgreen, linewidth = 4, linealpha = 0.1) + + l = @layout [a b] + plt = plot(p2, p4, layout = l) + + suffix = case_suffix(cfg, N_ens, rng_idx) + figure_fn = "pushforward_from_posterior_$(suffix)_k$(k)_full_ens.png" + savefig(plt, joinpath(data_save_directory, figure_fn)) + savefig(plt, joinpath(data_save_directory, replace(figure_fn, ".png" => ".pdf"))) + + # --- posterior ribbons plot --- + post_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(param_diffs)])' + post_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_ens)])' + prior_param_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(prior_param_diffs)])' + prior_G_quantiles = reduce(hcat, [quantile(row, [0.05, 0.5, 0.95]) for row in eachrow(G_prior)])' + + gr(size = (2 * 1.6 * 600, 600), guidefontsize = 18, tickfontsize = 16, legendfontsize = 16) + + # Panel (i): parameter ribbons + pa = plot( + collect(1:n_par), + zeros(n_par), + label = "solution", + color = :black, + linewidth = 4, + xlabel = "Parameter index", + ylabel = "parameters (input)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + plot!(pa, collect(1:n_par), prior_param_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_param_quantiles[:, 2] - prior_param_quantiles[:, 1] prior_param_quantiles[:, 3] - prior_param_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pa, collect(1:n_par), post_param_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_param_quantiles[:, 2] - post_param_quantiles[:, 1] post_param_quantiles[:, 3] - post_param_quantiles[:, 2]], + fillalpha = 0.2) + + # Panel (ii): output ribbons + pc = plot( + 1:ny, + y, + ribbon = sqrt.(diag(R)), + label = "data", + color = :black, + linewidth = 4, + xlabel = "Output index", + ylabel = "State statistics (output)", + left_margin = 15mm, + bottom_margin = 15mm, + ) + plot!(pc, 1:ny, prior_G_quantiles[:, 2], + color = :grey, label = "prior", linewidth = 4, + ribbon = [prior_G_quantiles[:, 2] - prior_G_quantiles[:, 1] prior_G_quantiles[:, 3] - prior_G_quantiles[:, 2]], + fillalpha = 0.2) + plot!(pc, 1:ny, post_G_quantiles[:, 2], + color = :blue, label = "posterior", linewidth = 4, + ribbon = [post_G_quantiles[:, 2] - post_G_quantiles[:, 1] post_G_quantiles[:, 3] - post_G_quantiles[:, 2]], + fillalpha = 0.2) + + ribbons_plt = plot(pa, pc, layout = @layout [a b]) + ribbons_fn = "posterior_ribbons_$(suffix)_k$(k)" + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".png")) + savefig(ribbons_plt, joinpath(data_save_directory, ribbons_fn * ".pdf")) + end +end diff --git a/examples/EKIRace/pushforward_from_posterior_l63.jl b/examples/EKIRace/pushforward_from_posterior_l63.jl new file mode 100644 index 000000000..09b29991b --- /dev/null +++ b/examples/EKIRace/pushforward_from_posterior_l63.jl @@ -0,0 +1,75 @@ +using JLD2 +using Distributions +using LinearAlgebra +using Random +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz63.jl") + +const n_pushforward_samples = 1000 + +function pushforward_one(cfg, N_ens, rng_idx; method = method_cases[1]) + calib_dir = calib_directory(method, cfg) + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + post_fn = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + + if !isfile(post_fn) + @warn "No posterior file for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + loaded_p = JLD2.load(post_fn) + if haskey(loaded_p, "pushforward_output_samples") + @info "Pushforward already present in $(post_fn); skipping." + return + end + + prelim_file = joinpath(homedir, "output", "l63_computed_preliminaries.jld2") + isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l63.jl first.") + prelim = JLD2.load(prelim_file) + x0 = prelim["x0"] + nx = length(x0) + ic_cov_sqrt = prelim["ic_cov_sqrt"] + lorenz_config_settings = prelim["lorenz_config_settings"] + observation_config = prelim["observation_config"] + n_output = 9 # 3 means + 3 variances + 3 covariances + + posteriors_by_k = loaded_p["posteriors_by_k"] + k_values = loaded_p["k_values"] + n_k = length(k_values) + + output_arr = Array{Float64}(undef, n_pushforward_samples, n_output, n_k) + + for (ki, k) in enumerate(k_values) + post_dist = posteriors_by_k[k] + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz63 evals" + for s in 1:n_pushforward_samples + output_arr[s, :, ki] = lorenz_forward( + EnsembleMemberConfig(push_con[:, s]), + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + end + + JLD2.jldopen(post_fn, "r+") do f + f["pushforward_output_samples"] = output_arr # (n_samples, n_output, n_k) + f["pushforward_k_values"] = k_values + f["pushforward_n_samples"] = n_pushforward_samples + end + @info "Saved pushforward to $(post_fn)" +end + +cfg = experiment_config(:l63) +for rng_idx in 1:cfg.n_repeats + for N_ens in cfg.N_ens_sizes + pushforward_one(cfg, N_ens, rng_idx) + end +end diff --git a/examples/EKIRace/pushforward_from_posterior_l96.jl b/examples/EKIRace/pushforward_from_posterior_l96.jl new file mode 100644 index 000000000..3443d60a5 --- /dev/null +++ b/examples/EKIRace/pushforward_from_posterior_l96.jl @@ -0,0 +1,99 @@ +using JLD2 +using Distributions +using LinearAlgebra +using Random +using Flux +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") +include("Lorenz96.jl") + +const n_pushforward_samples = 1000 + +function pushforward_one(cfg, N_ens, rng_idx; method = method_cases[1]) + force_case = cfg.force_case + calib_dir = calib_directory(method, cfg) + homedir = joinpath(pwd()) + data_save_directory = joinpath(homedir, "output", calib_dir) + + post_fn = joinpath(data_save_directory, posterior_filename(cfg, N_ens, rng_idx)) + calib_fn = joinpath(data_save_directory, results_filename(cfg, N_ens, rng_idx)) + + if !isfile(post_fn) + @warn "No posterior file for $(case_suffix(cfg, N_ens, rng_idx)); skipping." + return + end + + loaded_p = JLD2.load(post_fn) + if haskey(loaded_p, "pushforward_output_samples") + @info "Pushforward already present in $(post_fn); skipping." + return + end + + prelim_file = joinpath(homedir, "output", "l96_computed_preliminaries_$(force_case).jld2") + isfile(prelim_file) || error("Prelim file not found at $(prelim_file). Run calibrate_l96.jl first.") + prelim = JLD2.load(prelim_file) + x0 = prelim["x0"] + nx = length(x0) + ic_cov_sqrt = prelim["ic_cov_sqrt"] + lorenz_config_settings = prelim["lorenz_config_settings"] + observation_config = prelim["observation_config"] + n_output = 2 * nx + + calib_loaded = JLD2.load(calib_fn) + (truth_phi, _) = calib_loaded["truth_params_structure"] + (phi_structure, sample_range) = if force_case == "const-force" + (nothing, nothing) + elseif force_case == "vec-force" + (nothing, nothing) + elseif force_case == "flux-force" + (truth_phi.model, truth_phi.sample_range) + end + n_forcing = force_case == "flux-force" ? length(sample_range) : nx + + truth_params_constrained = loaded_p["truth_params_constrained"] + truth_emc = build_forcing(truth_phi, truth_params_constrained, phi_structure, sample_range) + truth_forcing_vec = forcing(truth_emc, x0) + + posteriors_by_k = loaded_p["posteriors_by_k"] + k_values = loaded_p["k_values"] + n_k = length(k_values) + + forcing_arr = Array{Float64}(undef, n_pushforward_samples, n_forcing, n_k) + output_arr = Array{Float64}(undef, n_pushforward_samples, n_output, n_k) + + for (ki, k) in enumerate(k_values) + post_dist = posteriors_by_k[k] + push_unc = sample(post_dist, n_pushforward_samples) + push_con = transform_unconstrained_to_constrained(post_dist, push_unc) + @info "Pushforward k=$(k), N_ens=$(N_ens), rng_idx=$(rng_idx): $(n_pushforward_samples) Lorenz96 evals" + for s in 1:n_pushforward_samples + emc = build_forcing(truth_phi, push_con[:, s], phi_structure, sample_range) + forcing_arr[s, :, ki] = forcing(emc, x0) + output_arr[s, :, ki] = lorenz_forward( + emc, + x0 .+ ic_cov_sqrt * rand(Normal(0.0, 1.0), nx, 1), + lorenz_config_settings, + observation_config, + ) + end + end + + JLD2.jldopen(post_fn, "r+") do f + f["pushforward_forcing_samples"] = forcing_arr # (n_samples, n_forcing, n_k) + f["pushforward_output_samples"] = output_arr # (n_samples, n_output, n_k) + f["pushforward_k_values"] = k_values + f["pushforward_n_samples"] = n_pushforward_samples + f["truth_forcing"] = truth_forcing_vec + end + @info "Saved pushforward to $(post_fn)" +end + +@assert EXPERIMENT in (:l96_const, :l96_vec, :l96_flux) "EXPERIMENT must be :l96_const, :l96_vec, or :l96_flux (got $EXPERIMENT)" +cfg = experiment_config(EXPERIMENT) +for rng_idx in 1:cfg.n_repeats + for N_ens in cfg.N_ens_sizes + pushforward_one(cfg, N_ens, rng_idx) + end +end From 61545be90e8005c9e33d0643ebc0262b693a1d73 Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 25 Jun 2026 14:27:21 -0700 Subject: [PATCH 85/86] docs for minibatcher tools --- docs/src/calibrate.md | 19 +++++++++++++++++++ docs/src/emulate.md | 6 ++++++ 2 files changed, 25 insertions(+) diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md index 44fad7bf1..da653a142 100644 --- a/docs/src/calibrate.md +++ b/docs/src/calibrate.md @@ -19,3 +19,22 @@ Documentation on how to construct an EnsembleKalmanProcess from the computer mod One draw of our approach is that it does not require the forward map to be written in Julia. To aid construction of such a workflow, EnsembleKalmanProcesses.jl provides a documented example of a BASH workflow for the [sinusoid problem](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/examples/sinusoid_example_toml/), with source code [here](https://github.com/CliMA/EnsembleKalmanProcesses.jl/tree/main/examples/SinusoidInterface). The forward map interacts with the calibration tools (EKP) only though TOML file reading an writing, and thus can be written in any language; for example, to be used with [slurm HPC scripts](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/examples/ClimateMachine_example/), with source code [here](https://github.com/CliMA/EnsembleKalmanProcesses.jl/tree/main/examples/ClimateMachine). +## Extracting training data after calibration + +Two helper functions from `CalibrateEmulateSample.Utilities` simplify passing the calibration results to the emulator. + +**`get_training_points(ekp, n)`** collects the stored parameter ensembles and forward-model outputs into a [`PairedDataContainer`](https://github.com/CliMA/EnsembleKalmanProcesses.jl/blob/main/src/DataContainers.jl) ready for emulator training. The argument `n` can be an integer (uses iterations `1:n`) or an index vector. An optional keyword `g_final` accepts the forward-model outputs at the final, not-yet-stored parameter ensemble. + +```julia +using CalibrateEmulateSample.Utilities +input_output_pairs = get_training_points(ekp, 5) # first 5 iterations +input_output_pairs = get_training_points(ekp, 1:2:9) # odd iterations 1,3,5,7,9 +input_output_pairs = get_training_points(ekp, 5; g_final = G(get_ϕ_final(prior, ekp))) +``` + +**`encoder_kwargs_from(ekp, prior)`** collects additional information from the calibration — the prior covariance, the observational noise covariance, and the sequence of parameter and output sample distributions across EKP iterations — into a named tuple that can be passed directly as the `encoder_kwargs` argument of the `Emulator`. This enables data-adaptive preprocessing such as likelihood-informed subspace reduction. + +```julia +encoder_kwargs = encoder_kwargs_from(ekp, prior) +``` + diff --git a/docs/src/emulate.md b/docs/src/emulate.md index bfb2d7f42..6f51c61f0 100644 --- a/docs/src/emulate.md +++ b/docs/src/emulate.md @@ -11,6 +11,12 @@ First, obtain data in a `PairedDataContainer`, for example, get this from an `En using CalibrateEmulateSample.Utilities input_output_pairs = Utilities.get_training_points(ekpobj, 5) # use first 5 iterations as data ``` + +!!! note "Minibatched calibration" + When the `EnsembleKalmanProcess` is built with an `ObservationSeries` and a minibatcher, each iteration uses a batch ``B_1, \ldots, B_n`` (``n \leq N``) drawn as a disjoint partition of statistically similar observations ``y_1, \ldots, y_N \sim \rho``. Minibatching is a technique for accelerating the calibration stage; it does not change the inverse problem being solved. + + `get_training_points` and `encoder_kwargs_from` therefore break the stacked batch outputs back into their per-observation components automatically, so the emulator is trained on ``(\theta, \mathcal{G}(\theta))`` pairs representative of the full distribution ``\rho``. In the Sampling stage the full observation set ``y_1, \ldots, y_N`` is used and batch structure is ignored. + Wrapping a predefined machine learning tool, e.g. a Gaussian process `gauss_proc`, the `Emulator` can then be built: ```julia From 35b524a107db32d9d9956c2f7bb3e0aeea82ad0a Mon Sep 17 00:00:00 2001 From: odunbar Date: Thu, 25 Jun 2026 14:29:03 -0700 Subject: [PATCH 86/86] add minibatch helpers into encoder_kwargs and get_training_points --- src/Utilities.jl | 181 ++++++++++++++++++++++++++++++++----- test/Utilities/runtests.jl | 124 +++++++++++++++++++++++++ 2 files changed, 282 insertions(+), 23 deletions(-) diff --git a/src/Utilities.jl b/src/Utilities.jl index ee7260277..3e8b35dc0 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -17,6 +17,7 @@ using TSVD import LinearAlgebra: norm export get_training_points +export flatten_minibatch_pairs export create_compact_linear_map export PairedDataContainerProcessor, DataContainerProcessor export create_encoder_schedule, @@ -78,36 +79,37 @@ function get_training_points( ) end - u_tp = [] - g_tp = [] - + os = get_observation_series(ekp) g_len = length(get_g(ekp)) - g_full = [get_g(ekp, i) for i in iter_range[iter_range .<= g_len]] + paired_range = [ir for ir in iter_range if ir <= g_len] + + # flatten_minibatch_pairs handles any batch size, including 1 (no-op for batch_size=1) + u_tp, g_tp = flatten_minibatch_pairs(ekp, paired_range) + if !isnothing(g_final) - if (isa(g_final, AbstractMatrix)) # add the matrix - if !(size(g_full[1]) == size(g_final)) - throw(ArgumentError("Expected `g_final` size: $(size(g_full[1])), received $(size(g_final)).")) + isa(g_final, AbstractMatrix) || + throw(ArgumentError("Expected `g_final` to be type `<:AbstractMatrix`, received $(typeof(g_final)).")) + final_u = get_u(ekp, maximum(iter_range)) + gdim = size(g_tp, 1) # per-element output dimension + last_bs = length(get_minibatch(os, last(paired_range))) + + if size(g_final, 1) == gdim + # g_final is in per-element (unbatched) form + u_tp = hcat(u_tp, final_u) + g_tp = hcat(g_tp, g_final) + elseif last_bs > 1 && size(g_final, 1) == gdim * last_bs + # g_final is stacked over batch elements — split it + for b in 1:last_bs + idx = (gdim * (b - 1) + 1):(gdim * b) + u_tp = hcat(u_tp, final_u) + g_tp = hcat(g_tp, g_final[idx, :]) end - push!(g_full, g_final) else - throw( - ArgumentError( - "Expected `g_final` to be type `<:AbstractMatrix`, of size: $(size(g_full[1])), received $(typeof(g_final)).", - ), - ) + throw(ArgumentError("Expected `g_final` size: $((gdim, size(final_u, 2))), received $(size(g_final)).")) end end - for (idx, ir) in enumerate(iter_range) - push!(u_tp, get_u(ekp, ir)) #N_parameters x N_ens - push!(g_tp, g_full[idx]) #N_data x N_ens - end - u_tp = reduce(hcat, u_tp) # N_parameters x (N_ek_it x N_ensemble)] - g_tp = reduce(hcat, g_tp) # N_data x (N_ek_it x N_ensemble) - - training_points = PairedDataContainer(u_tp, g_tp, data_are_columns = true) - - return training_points + return PairedDataContainer(u_tp, g_tp, data_are_columns = true) end # Using Observation Objects: @@ -162,6 +164,94 @@ end """ $(TYPEDSIGNATURES) +Flatten the paired parameter and batched-output ensembles from an +`EnsembleKalmanProcess` with minibatching into individual +(parameter, single-observation) pairs. + +In a minibatched EKP run each stored output has row dimension `batch_size × gdim`, +where `gdim` is the dimension of a single observation. This function splits those +stacked outputs into `batch_size` blocks of `gdim` rows and replicates the +corresponding parameter ensemble, treating each batch element as an independent +(parameter, observation) sample. + +`iter_range` must be a subset of `1:length(get_g(ekp))`. An integer `n` is +interpreted as `1:n` (clamped to the number of stored outputs). + +Returns `(u_flat, g_flat)` where +- `u_flat` is `dim_par × (Σ_k bs_k × N_ens)` — parameters repeated per batch element. +- `g_flat` is `gdim × (Σ_k bs_k × N_ens)` — outputs split to single-observation blocks. + +When `include_dt = true`, returns `(u_flat, g_flat, dt_flat)` where `dt_flat` is a +`Vector{Float64}` of per-pair algorithm times: 0 for iteration 1 (the initial ensemble), +and `get_algorithm_time(ekp)[ir - 1]` for subsequent iterations, each replicated +`batch_size` times. + +Used internally by `get_training_points` and `encoder_kwargs_from` whenever a batch +size > 1 is detected, so most users will not need to call this directly. +""" +function flatten_minibatch_pairs( + ekp::EKP.EnsembleKalmanProcess{FT, IT, P}, + iter_range::Union{IT, AbstractVector{IT}}; + include_dt::Bool = false, +) where {FT, IT, P} + if !isa(iter_range, AbstractVector) + iter_range = 1:min(iter_range, length(get_g(ekp))) + end + + os = get_observation_series(ekp) + g_len = length(get_g(ekp)) + alg_time = include_dt ? get_algorithm_time(ekp) : nothing + + u_out = [] + g_out = [] + dt_out = include_dt ? Float64[] : nothing + + for ir in iter_range + if ir > g_len + throw( + ArgumentError( + "iter_range contains iteration $ir, but EKP only stores $g_len " * + "outputs. iter_range must be a subset of 1:length(get_g(ekp)).", + ), + ) + end + uu = get_u(ekp, ir) + gg = get_g(ekp, ir) + batch = get_minibatch(os, ir) + bs = length(batch) + + ndim = size(gg, 1) + ndim % bs == 0 || throw( + ArgumentError( + "At iteration $ir: output dimension ($ndim) is not divisible by batch " * + "size ($bs). Ensure outputs are stacked as [g_b1; g_b2; …] for each " * + "batch element.", + ), + ) + gdim = ndim ÷ bs + + for b in 1:bs + idx = (gdim * (b - 1) + 1):(gdim * b) + push!(u_out, uu) + push!(g_out, gg[idx, :]) + end + + if include_dt + di = (ir == 1) ? zero(eltype(alg_time)) : alg_time[ir - 1] + append!(dt_out, fill(di, bs)) + end + end + + if include_dt + return reduce(hcat, u_out), reduce(hcat, g_out), dt_out + else + return reduce(hcat, u_out), reduce(hcat, g_out) + end +end + +""" +$(TYPEDSIGNATURES) + Extracts the relevant encoder kwargs from a vector triple (samples_in, samples_out, dt). Samples describe an ordered sequence of distributions in input and output space, each indexed with a temperature, or algorithm time, `dt`. Contains @@ -216,6 +306,8 @@ function encoder_kwargs_from( final_samples_out = nothing, ) where {PD <: ParameterDistribution, EKP <: EnsembleKalmanProcess} observation_series = isnothing(observation_series) ? get_observation_series(ekp) : observation_series + # track whether samples were user-supplied (don't auto-flatten those) + auto_flatten_io = isnothing(samples_in) && isnothing(samples_out) samples_in = isnothing(samples_in) ? get_u(ekp) : samples_in samples_out = isnothing(samples_out) ? get_g(ekp) : samples_out dt = isnothing(dt) ? get_algorithm_time(ekp) : dt @@ -238,6 +330,49 @@ function encoder_kwargs_from( end end + # when samples were derived from ekp, flatten any minibatched outputs + orig_g_len = length(get_g(ekp)) + if auto_flatten_io && + orig_g_len > 0 && + any(length(get_minibatch(observation_series, ir)) > 1 for ir in 1:orig_g_len) + + # delegate u/g/dt splitting to the public helper; rebuild per-pair lists from the result + u_flat, g_flat, flat_dt = flatten_minibatch_pairs(ekp, 1:orig_g_len; include_dt = true) + n_ens_per = size(get_u(ekp, 1), 2) + n_pairs = size(u_flat, 2) ÷ n_ens_per # = Σ bs_k over stored iterations + flat_si = [u_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs] + flat_so = [g_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs] + + # handle final_samples_out that was appended beyond orig_g_len + if length(samples_out) > orig_g_len + final_so = samples_out[orig_g_len + 1] + gdim_ref = size(flat_so[1], 1) + last_bs = length(get_minibatch(observation_series, orig_g_len)) + final_si = samples_in[min(orig_g_len + 1, length(samples_in))] + final_di = Float64(dt[min(orig_g_len, length(dt))]) + if size(final_so, 1) == gdim_ref + push!(flat_si, final_si) + push!(flat_so, final_so) + push!(flat_dt, final_di) + elseif size(final_so, 1) == gdim_ref * last_bs + for b in 1:last_bs + idx = (gdim_ref * (b - 1) + 1):(gdim_ref * b) + push!(flat_si, final_si) + push!(flat_so, final_so[idx, :]) + push!(flat_dt, final_di) + end + else + push!(flat_si, final_si) + push!(flat_so, final_so) + push!(flat_dt, final_di) + end + end + + samples_in = flat_si + samples_out = flat_so + dt = flat_dt + end + prior_kwargs = encoder_kwargs_from(prior) obs_kwargs = encoder_kwargs_from(observation_series) io_kwargs = encoder_kwargs_from(samples_in, samples_out, dt) diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl index 2b267bcae..604e61dd3 100644 --- a/test/Utilities/runtests.jl +++ b/test/Utilities/runtests.jl @@ -689,6 +689,130 @@ end +@testset "Minibatch flattening" begin + rng_mb = Random.MersenneTwister(77) + + gdim = 3 # dimension of one observation + n_obs = 4 # total observations in the series + dim_par = 2 + n_ens = 5 + batch_size = 2 # two batch elements per EKP iteration + + obs_vec = [Observation(randn(rng_mb, gdim), Matrix{Float64}(I, gdim, gdim), "obs_$i") for i in 1:n_obs] + given_batches = [[1, 2], [3, 4]] + obs_series_mb = ObservationSeries(obs_vec, FixedMinibatcher(given_batches)) + + prior_mb = constrained_gaussian("test", 0.0, 1.0, -Inf, Inf, repeats = dim_par) + initial_ensemble_mb = construct_initial_ensemble(rng_mb, prior_mb, n_ens) + ekp_mb = + EnsembleKalmanProcesses.EnsembleKalmanProcess(initial_ensemble_mb, obs_series_mb, Inversion(); rng = rng_mb) + + # store the raw stacked outputs for 2 iterations (one per FixedMinibatcher batch) + g_stored = Matrix{Float64}[] + for _ in 1:2 + bs = length(get_current_minibatch(ekp_mb)) + g_batch = randn(rng_mb, gdim * bs, n_ens) + push!(g_stored, copy(g_batch)) + EnsembleKalmanProcesses.update_ensemble!(ekp_mb, g_batch) + end + n_iters = 2 + + # ---- flatten_minibatch_pairs: shape ---- + u_flat, g_flat = flatten_minibatch_pairs(ekp_mb, 1:n_iters) + total_cols = n_iters * batch_size * n_ens + @test size(u_flat) == (dim_par, total_cols) + @test size(g_flat) == (gdim, total_cols) + + # both batch-element copies of u at iteration 1 must equal get_u(ekp_mb, 1) + u1 = get_u(ekp_mb, 1) + @test all(isapprox.(u_flat[:, 1:n_ens], u1, atol = 1e-12)) + @test all(isapprox.(u_flat[:, (n_ens + 1):(2n_ens)], u1, atol = 1e-12)) + + # g_flat slices match the two gdim-row blocks of the stored batched output + @test all(isapprox.(g_flat[:, 1:n_ens], g_stored[1][1:gdim, :], atol = 1e-12)) + @test all(isapprox.(g_flat[:, (n_ens + 1):(2n_ens)], g_stored[1][(gdim + 1):(2gdim), :], atol = 1e-12)) + + # integer shorthand: flatten_minibatch_pairs(ekp, 2) == flatten_minibatch_pairs(ekp, 1:2) + u_flat2, g_flat2 = flatten_minibatch_pairs(ekp_mb, n_iters) + @test all(isapprox.(u_flat2, u_flat, atol = 1e-12)) + @test all(isapprox.(g_flat2, g_flat, atol = 1e-12)) + + # include_dt: length and first batch_size entries are 0 (initial ensemble) + u_flat3, g_flat3, dt_flat = flatten_minibatch_pairs(ekp_mb, 1:n_iters; include_dt = true) + @test all(isapprox.(u_flat3, u_flat, atol = 1e-12)) + @test all(isapprox.(g_flat3, g_flat, atol = 1e-12)) + @test length(dt_flat) == total_cols ÷ n_ens + @test all(dt_flat[1:batch_size] .== 0) + + # error: iter_range references an iteration with no stored output + @test_throws ArgumentError flatten_minibatch_pairs(ekp_mb, 1:5) + + # ---- get_training_points: minibatch path ---- + tp = get_training_points(ekp_mb, n_iters) + @test all(isapprox.(get_inputs(tp), u_flat, atol = 1e-12)) + @test all(isapprox.(get_outputs(tp), g_flat, atol = 1e-12)) + + # g_final in per-element (unbatched) form: one extra column group appended + g_final_flat = randn(rng_mb, gdim, n_ens) + tp_f = get_training_points(ekp_mb, n_iters, g_final = g_final_flat) + @test size(get_inputs(tp_f), 2) == total_cols + n_ens + @test all(isapprox.(get_outputs(tp_f)[:, (total_cols + 1):end], g_final_flat, atol = 1e-12)) + + # g_final in batched (stacked) form: split into batch_size groups + g_final_batched = randn(rng_mb, gdim * batch_size, n_ens) + tp_fb = get_training_points(ekp_mb, n_iters, g_final = g_final_batched) + @test size(get_inputs(tp_fb), 2) == total_cols + batch_size * n_ens + @test all( + isapprox.( + get_outputs(tp_fb)[:, (total_cols + 1):(total_cols + n_ens)], + g_final_batched[1:gdim, :], + atol = 1e-12, + ), + ) + @test all( + isapprox.( + get_outputs(tp_fb)[:, (total_cols + n_ens + 1):end], + g_final_batched[(gdim + 1):(2gdim), :], + atol = 1e-12, + ), + ) + + # wrong g_final size + @test_throws ArgumentError get_training_points(ekp_mb, n_iters, g_final = randn(rng_mb, gdim + 1, n_ens)) + + # ---- encoder_kwargs_from: minibatch path ---- + ekp_kw = encoder_kwargs_from(ekp_mb, prior_mb) + + # prior and obs kwargs are structurally unchanged + @test ekp_kw[:obs_noise_cov] == encoder_kwargs_from(obs_series_mb)[:obs_noise_cov] + @test all(isapprox.(ekp_kw[:prior_cov], encoder_kwargs_from(prior_mb)[:prior_cov], atol = 1e-12)) + + # samples_in/out lists reflect per-element (flattened) structure + flat_si_list = ekp_kw[:input_structure_vecs][:samples_in] + flat_so_list = ekp_kw[:output_structure_vecs][:samples_out] + @test length(flat_si_list) == n_iters * batch_size + @test length(flat_so_list) == n_iters * batch_size + @test size(flat_si_list[1]) == (dim_par, n_ens) + @test size(flat_so_list[1]) == (gdim, n_ens) + + # dt: first batch_size entries correspond to the initial distribution (time = 0) + flat_dt_kw = ekp_kw[:input_structure_vecs][:dt] + @test all(flat_dt_kw[1:batch_size] .== 0) + + # final_samples_out in per-element form → one extra pair appended + ekp_kw_f = encoder_kwargs_from(ekp_mb, prior_mb, final_samples_out = g_final_flat) + @test length(ekp_kw_f[:output_structure_vecs][:samples_out]) == n_iters * batch_size + 1 + + # final_samples_out in batched form → split into batch_size extra pairs + ekp_kw_fb = encoder_kwargs_from(ekp_mb, prior_mb, final_samples_out = g_final_batched) + @test length(ekp_kw_fb[:output_structure_vecs][:samples_out]) == n_iters * batch_size + batch_size + + # user-supplied samples_in/out are NOT auto-flattened + ekp_kw_custom = encoder_kwargs_from(ekp_mb, prior_mb; samples_in = get_u(ekp_mb), samples_out = get_g(ekp_mb)) + @test size(ekp_kw_custom[:output_structure_vecs][:samples_out][1], 1) == gdim * batch_size + +end + @testset "Decorrelator: Large observational covariance" begin # loop over output dim