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/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 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..ae607c4cf --- /dev/null +++ b/examples/EKIRace/Project.toml @@ -0,0 +1,14 @@ +[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" +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/README.md b/examples/EKIRace/README.md new file mode 100644 index 000000000..0a67cabf9 --- /dev/null +++ b/examples/EKIRace/README.md @@ -0,0 +1,156 @@ +# 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` before starting a run, e.g. + +```julia +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[3] # :l96_vec +calibrate_date = Date("2026-06-04", "yyyy-mm-dd") +``` + +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 + +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")' +``` + +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) + +```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 +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 to the +calibration output directory. + +### 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` (including the pushforward samples written in +stage 4) and writes pushforward and posterior-ribbon diagnostic plots into +the calibration output directory. + +### 6. Leaderboard + +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")' +``` + +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 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: +experiments = [:l63, :l96_const, :l96_vec, :l96_flux] +EXPERIMENT = experiments[3] # :l96_vec +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("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")' +``` + +## 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/calibrate_l63.jl b/examples/EKIRace/calibrate_l63.jl new file mode 100644 index 000000000..15deb01ec --- /dev/null +++ b/examples/EKIRace/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") # Contains Lorenz 96 source code +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true +######################################################################## +############### Choose problem type and structure ###################### +######################################################################## + +cfg = experiment_config(:l63) +N_ens_sizes = cfg.N_ens_sizes +N_iter = cfg.N_iter +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" +configuration = + 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 +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 +] +distribution = Parameterized(MvNormal(prior_mean, prior_cov)) +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 + + + +######################################################################## +############################ 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) + +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 ########################### +######################################################################## + +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) + +# method_names defined in experiment_config.jl + + +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(), +# TransformInversion(), +# GaussNewtonInversion(prior), +# Unscented(prior), + ] + + @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, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = terminate_at), + ) + else + ekpobj = EKP.EnsembleKalmanProcess( + initial_params, + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = terminate_at), + ) + end + Ne = get_N_ens(ekpobj) + + 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 (diagnostic; not used as stopping criterion) + ens_mean = mean(params_i, dims = 2)[:] # in constrained_space + G_ens_mean = lorenz_forward( + EnsembleMemberConfig(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)" + + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] + + G_ens = hcat( + [ + lorenz_forward( + EnsembleMemberConfig(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 + ]..., + ) + 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 + 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 + JLD2.save(joinpath(per_method_dir, prior_filename(cfg)), "prior", prior) + end + if save_all_ekp + # JLD2 + JLD2.save( + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rr)), + "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, rr)), + "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, summary_filename(cfg)) +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 new file mode 100644 index 000000000..79d2ae795 --- /dev/null +++ b/examples/EKIRace/calibrate_l96.jl @@ -0,0 +1,372 @@ +# 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 +include("experiment_config.jl") + +verbose_flag = false +save_all_ekp = true +######################################################################## +############### Choose problem type and structure ###################### +######################################################################## + +@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 +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" +# saved and loaded for plotting etc. +configuration = Dict( + "case" => case, + "N_iter" => N_iter, + "N_ens_sizes" => N_ens_sizes, + "rng_seeds" => rng_seeds, + "terminate_at" => terminate_at +) + +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) + 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(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(deepcopy(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)) # original cov + 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" + 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 ########################### +######################################################################## + +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) + +# method_names defined in experiment_config.jl + +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(), +# TransformInversion(), +# GaussNewtonInversion(prior), +# Unscented(prior), + ] + + @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, + R, + 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( + initial_params, + y, + R, + deepcopy(method); + rng = copy(rng), + verbose = verbose_flag, + # accelerator = DefaultAccelerator(), + localization_method = NoLocalization(), + scheduler = DataMisfitController(terminate_at = terminate_at), +# scheduler = DefaultScheduler(0.1), + ) + end + Ne = get_N_ens(ekpobj) + + 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 (diagnostic; not used as stopping criterion) + ens_mean = mean(params_i, dims = 2)[:] + 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)" + ens_mean_final = ens_mean + G_ens_mean_final = G_ens_mean[:] + + 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 + ]..., + ) + 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 + 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 + JLD2.save(joinpath(per_method_dir, prior_filename(cfg)), "prior", prior) + end + if save_all_ekp + # JLD2 + JLD2.save( + joinpath(per_method_dir, ekp_filename(cfg, N_ens, rr)), + "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, rr)), + "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, summary_filename(cfg)) +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/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/compute_leaderboard_metrics.jl b/examples/EKIRace/compute_leaderboard_metrics.jl new file mode 100644 index 000000000..214b3aee7 --- /dev/null +++ b/examples/EKIRace/compute_leaderboard_metrics.jl @@ -0,0 +1,901 @@ +using NCDatasets +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" +#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) +prelim_filename = joinpath(indir, prelim_jld2_file) +@info "computing leaderboard metrics from $(filename)" + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +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 #################################### +########################################################################### +# +# 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) +# :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, :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 +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 ############## +########################################################################### + +@info "computing leaderboard metrics from $(filename)" +ncd = NCDataset(filename) +mh = ncd[:mahalanobis] +lp = ncd[:posterior_logpdf_true_v_map] +(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 + +# 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") || (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_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() +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,:,:] +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) + +########################################################################### +#################### Param-space: Mahalanobis and log-ratio ############## +########################################################################### + +if show_space(:param) && (show_metric(:mahalanobis) || show_metric(:logratio)) + + println("\n" * "="^72) + println("Param-space (chi-sq ref dim = $(n_par))") + println("="^72) + + 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) + + 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 + + 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 + + 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 + + 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 (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 + +########################################################################### +#################### 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") + + 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 (chi-sq ref dim = $(n_for))") + println("="^72) + 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) + 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_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 + 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 + 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 + + 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}[], 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], 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 + + 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, score_cols), allrows=true) + println() + end + end +end + +########################################################################### +#################### 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] + 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 (chi-sq ref dim = $(n_out))") + println("="^72) + 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) + 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_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 + 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 + 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 + + 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}[], 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], 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 + + 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, score_cols), allrows=true) + println() + end + end +end + +########################################################################### +#################### 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 noise-floor directions +# total ~ Chisq(n) — full-space (= top + residual) + +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"), + ] + (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) + 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 + + 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 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 + +########################################################################### +#################### 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) + + # ── 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("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, 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[], + ) + 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("\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(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 + +########################################################################### +#################### 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/emulate_sample_l63.jl b/examples/EKIRace/emulate_sample_l63.jl new file mode 100644 index 000000000..622499a87 --- /dev/null +++ b/examples/EKIRace/emulate_sample_l63.jl @@ -0,0 +1,157 @@ +# 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 +include("experiment_config.jl") + + +function main() + + #### 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 = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L63: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + # loading relevant data + 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)...") + continue + 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 + + # one file per EKI experiment, containing all k posteriors + 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 + end +end + + +main() diff --git a/examples/EKIRace/emulate_sample_l96.jl b/examples/EKIRace/emulate_sample_l96.jl new file mode 100644 index 000000000..45b48d545 --- /dev/null +++ b/examples/EKIRace/emulate_sample_l96.jl @@ -0,0 +1,177 @@ +# 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 +include("experiment_config.jl") + + +function main() + + #### 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 = case_suffix(cfg, N_ens, rng_idx) + @info("Perform emulate-sample for L96: \n method: $(calib_dir) \n experiment: $(calib_filename_suffix)") + + # loading relevant data + 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"] + + 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)...") + continue + 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")] + + 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 = deepcopy(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(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) + 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 + + # one file per EKI experiment, containing all k posteriors + 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 + end + end +end + +main() 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 new file mode 100644 index 000000000..350cd927e --- /dev/null +++ b/examples/EKIRace/experiment_config.jl @@ -0,0 +1,159 @@ +using Dates + +######################################################################## +############### USER TOGGLE ######################################### +######################################################################## +# 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). +# PIN this to reproduce a specific run directory. +#calibrate_date = Date("2026-06-03", "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) + n_ens_step = 8 + n_repeats = 20 + if case == :l63 + ens_step = 2 + return ( + model = "l63", + force_case = nothing, + 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 = n_repeats, + max_iter = 10, + retain_var = 0.99, + n_features = 100, + n_features_opt = 60, + calibrate_date = calibrate_date, + ) + elseif case == :l96_const + ens_step_const = 2 + return ( + model = "l96", + force_case = "const-force", + N_ens_sizes = collect(4:ens_step_const:4+n_ens_step*ens_step_const), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_vec + ens_step_vec = 5 + return ( + model = "l96", + force_case = "vec-force", + N_ens_sizes = collect(40:ens_step_vec:40+n_ens_step*ens_step_vec), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_flux + ens_step_flux = 5 + return ( + model = "l96", + force_case = "flux-force", + N_ens_sizes = collect(30:ens_step_flux:30+n_ens_step*ens_step_flux), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + 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/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..ae607c4cf --- /dev/null +++ b/examples/EKIRace/hpc-variant/Project.toml @@ -0,0 +1,14 @@ +[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" +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..7f28f5393 --- /dev/null +++ b/examples/EKIRace/hpc-variant/README.md @@ -0,0 +1,197 @@ +# 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 before starting a run, e.g. + +```julia +calibrate_date = Date("2026-06-04", "yyyy-mm-dd") +``` + +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. + +## Pipeline + +### L63 + +``` + ┌──afterany──► posterior_diagnostic_plots_l63 +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──► pushforward_from_posterior ──afterany──► posterior_diagnostic_plots_l96 + ──afterany──► exp_to_leaderboard +``` + +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) + +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 +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 +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=. 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 +``` + +`exp_to_leaderboard.jl` runs all cells serially in a single call (requires +pushforward to have been run first): + +```bash +julia --project=. exp_to_leaderboard.jl +EXPERIMENT=l96_const julia --project=. exp_to_leaderboard.jl +``` + +## HPC (Caltech Resnick cluster, SLURM) + +### Submission scripts (recommended) + +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 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: + +```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 +wait +``` + +### Manual submission + +Precompile via `submit_precompile.sh` (or directly), then submit each stage: + +```bash +# L63 +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) +PUSHFWD_JID=$(sbatch --parsable -A esm \ + --dependency=afterany:${EMU_JID} \ + 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 +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) +PUSHFWD_JID=$(sbatch --parsable -A esm \ + --dependency=afterany:${EMU_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + pushforward_from_posterior.sbatch) +sbatch -A esm \ + --dependency=afterany:${PUSHFWD_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + posterior_diagnostic_plots_l96.sbatch +sbatch -A esm \ + --dependency=afterany:${PUSHFWD_JID} \ + --export=ALL,EXPERIMENT=l96_const \ + exp_to_leaderboard.sbatch +``` + +### Sbatch files reference + +| File | Type | Description | +|------|------|-------------| +| `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 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()` | + +### Adjusting array size + +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` 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 + +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..6ed860afe --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_array.sbatch @@ -0,0 +1,44 @@ +#!/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-180%100 +#SBATCH --time=02:00:00 +#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 + +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 "${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/calibrate_l63.jl b/examples/EKIRace/hpc-variant/calibrate_l63.jl new file mode 100644 index 000000000..9edcc1b41 --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_l63.jl @@ -0,0 +1,300 @@ +# 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) + pf_tmp = splitext(pf)[1] * ".tmp.$(getpid()).jld2" + JLD2.save(pf_tmp, "prior", setup.prior) + try + mv(pf_tmp, pf) + catch + rm(pf_tmp; force = true) + end + 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) + + 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() + + 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..b10cb724b --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibrate_l96.jl @@ -0,0 +1,366 @@ +# 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) + 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(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(deepcopy(phi_structure), x_train, prior_train) + #prior_cov = (0.1^2) * I(length(prior_mean)) # original cov + 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") + 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) + 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, + ) + # 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( + "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) + pf_tmp = splitext(pf)[1] * ".tmp.$(getpid()).jld2" + JLD2.save(pf_tmp, "prior", setup.prior) + try + mv(pf_tmp, pf) + catch + rm(pf_tmp; force = true) + end + 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/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..d0bafc31a --- /dev/null +++ b/examples/EKIRace/hpc-variant/calibration_diagnostic_plots_l96.sbatch @@ -0,0 +1,31 @@ +#!/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=00:30:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=1 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +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/compute_leaderboard_metrics.jl b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl new file mode 100644 index 000000000..cb8ba37a4 --- /dev/null +++ b/examples/EKIRace/hpc-variant/compute_leaderboard_metrics.jl @@ -0,0 +1,932 @@ +using NCDatasets +using Distributions +using Statistics +using Printf +using DataFrames +using Dates +using JLD2 +using LinearAlgebra + +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[2] +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) +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 ################################### +########################################################################### + +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 #################################### +########################################################################### +# +# 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) +# :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, :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 +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 ############## +########################################################################### + +@info "computing leaderboard metrics from $(filename)" +ncd = NCDataset(filename) +mh = ncd[:mahalanobis] +lp = ncd[:posterior_logpdf_true_v_map] +(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 + +# 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_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 + @error "precomputed quantities NOT FOUND at $(prelim_filename)" + 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") || (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_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() +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,:,:] +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) + +########################################################################### +#################### Param-space: Mahalanobis and log-ratio ############## +########################################################################### + +if show_space(:param) && (show_metric(:mahalanobis) || show_metric(:logratio)) + + println("\n" * "="^72) + println("Param-space (chi-sq ref dim = $(n_par))") + println("="^72) + + 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) + + 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 + + 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 + + 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 + + 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 (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 + +########################################################################### +#################### 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") + + 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 (chi-sq ref dim = $(n_for))") + println("="^72) + 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) + 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_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 + 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 + 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 + + 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}[], 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], 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 + + 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, score_cols), allrows=true) + println() + end + end +end + +########################################################################### +#################### 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] + 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 (chi-sq ref dim = $(n_out))") + println("="^72) + 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) + 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_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 + 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 + 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 + + 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}[], 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], 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 + + 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, score_cols), allrows=true) + println() + end + end +end + +########################################################################### +#################### 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 noise-floor directions +# total ~ Chisq(n) — full-space (= top + residual) + +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"), + ] + (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) + 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 + + 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 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 + +########################################################################### +#################### 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) + + # ── 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("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, 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[], + ) + 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("\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(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 + +########################################################################### +#################### 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, 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] + + # 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)_$(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/emulate_sample_array.sbatch b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch new file mode 100644 index 000000000..b1caba570 --- /dev/null +++ b/examples/EKIRace/hpc-variant/emulate_sample_array.sbatch @@ -0,0 +1,41 @@ +#!/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-180%100 +#SBATCH --time=12:00:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +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} + +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}" 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..0d5938968 --- /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 = deepcopy(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..6759f9e4a --- /dev/null +++ b/examples/EKIRace/hpc-variant/emulate_sample_l96.jl @@ -0,0 +1,188 @@ +# 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")] + + 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 = deepcopy(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(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) + 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/exp_to_leaderboard.jl b/examples/EKIRace/hpc-variant/exp_to_leaderboard.jl new file mode 100644 index 000000000..3918c4bc6 --- /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, 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/hpc-variant/exp_to_leaderboard.sbatch b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch new file mode 100644 index 000000000..96dafcebc --- /dev/null +++ b/examples/EKIRace/hpc-variant/exp_to_leaderboard.sbatch @@ -0,0 +1,31 @@ +#!/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=00:30:00 +#SBATCH --mem=8G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +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=. exp_to_leaderboard.jl diff --git a/examples/EKIRace/hpc-variant/experiment_config.jl b/examples/EKIRace/hpc-variant/experiment_config.jl new file mode 100644 index 000000000..388b8674e --- /dev/null +++ b/examples/EKIRace/hpc-variant/experiment_config.jl @@ -0,0 +1,190 @@ +using Dates + +######################################################################## +############### USER TOGGLE ######################################### +######################################################################## +# 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). +# PIN this before submitting an array job so all tasks use the same directory. +#calibrate_date = Date("2026-06-03", "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) + 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 = collect(4:ens_step:4+n_ens_step*ens_step), + N_iter = 20, + terminate_at = 2.0, # DataMisfitController end time + n_repeats = n_repeats, + max_iter = 10, + retain_var = 0.99, + n_features = 100, + n_features_opt = 60, + calibrate_date = calibrate_date, + ) + elseif case == :l96_const + ens_step_const = 2 + return ( + model = "l96", + force_case = "const-force", + N_ens_sizes = collect(4:ens_step_const:4+n_ens_step*ens_step_const), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_vec + ens_step_vec = 5 + return ( + model = "l96", + force_case = "vec-force", + N_ens_sizes = collect(40:ens_step_vec:40+n_ens_step*ens_step_vec), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + n_features = 200, + n_features_opt = 160, + calibrate_date = calibrate_date, + ) + elseif case == :l96_flux + ens_step_flux = 5 + return ( + model = "l96", + force_case = "flux-force", + N_ens_sizes = collect(30:ens_step_flux:30+n_ens_step*ens_step_flux), + N_iter = 20, + terminate_at = 2.0, + n_repeats = n_repeats, + max_iter = 15, + retain_var = 0.99, + 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/l63_exp_to_leaderboard_utilities.jl b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl new file mode 100644 index 000000000..db204bd06 --- /dev/null +++ b/examples/EKIRace/hpc-variant/l63_exp_to_leaderboard_utilities.jl @@ -0,0 +1,293 @@ +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 # 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 + +#### 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, 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"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# 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) +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) + +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 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"])) + + 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 + + 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) + + # 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) + 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 + + +### 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) +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",)) +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 + +# 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 + +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 new file mode 100644 index 000000000..c916040ac --- /dev/null +++ b/examples/EKIRace/hpc-variant/l96_exp_to_leaderboard_utilities.jl @@ -0,0 +1,375 @@ +using NCDatasets +using Dates +using JLD2 +using Distributions +using LinearAlgebra +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.EnsembleKalmanProcesses + +include("experiment_config.jl") + +EXPERIMENT = l96_experiment() # respects EXPERIMENT env var / CLI arg + +########################################################################### +#################### Metric parameters ################################### +########################################################################### + +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 + +#### 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, 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"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# 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.") +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) +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) +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) +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) +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) + @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_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))) + 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 + truth_forcing_arr[i, j, :] = truth_forcing_vec + + 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 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) + + # 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) + 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 (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 + + # 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 + + +### 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) +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",)) +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 + +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 + +# 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 + +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 + +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 - 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 + +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_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..eb55e318b --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l63.sbatch @@ -0,0 +1,30 @@ +#!/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 +#SBATCH --constraint=cascadelake + +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/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl new file mode 100644 index 000000000..eb0a1a6d0 --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.jl @@ -0,0 +1,379 @@ +# 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 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 + + 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 = 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" + 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 + ]..., + ) + + 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] + + # 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_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) + 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(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) + 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)]: 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) + + # 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..e280e6e19 --- /dev/null +++ b/examples/EKIRace/hpc-variant/posterior_diagnostic_plots_l96.sbatch @@ -0,0 +1,32 @@ +#!/bin/bash +# 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: +# sbatch --dependency=afterok: \ +# --export=ALL,EXPERIMENT=l96_vec posterior_diagnostic_plots_l96.sbatch + +#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-180%100 +#SBATCH --time=00:30:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +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 "${SLURM_ARRAY_TASK_ID}" diff --git a/examples/EKIRace/hpc-variant/precompile.sbatch b/examples/EKIRace/hpc-variant/precompile.sbatch new file mode 100644 index 000000000..67c6906a8 --- /dev/null +++ b/examples/EKIRace/hpc-variant/precompile.sbatch @@ -0,0 +1,37 @@ +#!/bin/bash +# Single SLURM job that instantiates and precompiles the project. +# 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=4 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +set -euo pipefail + +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 + +# 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 new file mode 100644 index 000000000..e58a74fcc --- /dev/null +++ b/examples/EKIRace/hpc-variant/pushforward_from_posterior.sbatch @@ -0,0 +1,40 @@ +#!/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=00:30:00 +#SBATCH --mem=16G +#SBATCH --cpus-per-task=4 +#SBATCH --ntasks=1 +#SBATCH --constraint=cascadelake + +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 new file mode 100755 index 000000000..db5ba198d --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l63.sh @@ -0,0 +1,70 @@ +#!/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 "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) ===" +CALIB_JID=$(sbatch --parsable \ + -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 \ + -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}" + +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 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 \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${PUSHFWD_JID} \ + --export=ALL,EXPERIMENT=l63 \ + 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="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 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..e9a18d56f --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_const.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# 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 +# 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 "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) ===" +CALIB_JID=$(sbatch --parsable \ + -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 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 \ + --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}" + +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:${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 ${PUSHFWD_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${PUSHFWD_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 --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 new file mode 100755 index 000000000..58edc2910 --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_flux.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# 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 +# 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 "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) ===" +CALIB_JID=$(sbatch --parsable \ + -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 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 \ + --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}" + +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:${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 ${PUSHFWD_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${PUSHFWD_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="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 new file mode 100755 index 000000000..9ee7e2d8f --- /dev/null +++ b/examples/EKIRace/hpc-variant/submit_l96_vec.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# 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 +# 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 "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) ===" +CALIB_JID=$(sbatch --parsable \ + -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 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 \ + --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}" + +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:${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 ${PUSHFWD_JID}) ===" +LB_JID=$(sbatch --parsable \ + -A esm \ + --job-name="leaderboard_${LABEL}" \ + --dependency=afterany:${PUSHFWD_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 --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 diff --git a/examples/EKIRace/hpc-variant/submit_precompile.sh b/examples/EKIRace/hpc-variant/submit_precompile.sh new file mode 100755 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 ===" 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..c2bef088e --- /dev/null +++ b/examples/EKIRace/l63_exp_to_leaderboard_utilities.jl @@ -0,0 +1,293 @@ +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 # 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 + +#### 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, 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"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# 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) +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) + +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_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"])) + + 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 + + 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) + + # 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) + 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 + + +### 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) +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",)) +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 + +# 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 + +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 new file mode 100644 index 000000000..bcca90bb1 --- /dev/null +++ b/examples/EKIRace/l96_exp_to_leaderboard_utilities.jl @@ -0,0 +1,373 @@ +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 # 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 + +#### 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, 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"]) + for (N_ens, rng_idx) in valid_file_items +) + +n_rng = length(rng_idxs) +n_ens = length(N_enss) + +# 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.") +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) +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) +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) +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) +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) + @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"] + + 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))) + 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 + truth_forcing_arr[i, j, :] = truth_forcing_vec + + 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 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) + + # 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) + 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 (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 + + # 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 + + +### 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) +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",)) +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 + +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 + +# 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 + +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 + +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 - 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 + +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/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/posterior_diagnostic_plots_l96.jl b/examples/EKIRace/posterior_diagnostic_plots_l96.jl new file mode 100644 index 000000000..bb64e2c58 --- /dev/null +++ b/examples/EKIRace/posterior_diagnostic_plots_l96.jl @@ -0,0 +1,403 @@ +# 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 = 400 # 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) + +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) + 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) + @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)) + + 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" + ([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 + + 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] + + # 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_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) + 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(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) + 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)]: 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) + + # 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 + + 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 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