diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md index 44fad7bf1..fa3231e02 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; g_final = nothing)`](@ref CalibrateEmulateSample.Utilities.get_training_points) collects the stored parameter ensembles and forward-model outputs into a [`PairedDataContainer`](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/internal_data_representation/) ready for emulator training. The argument `n` can be an integer (which 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 +input_output_pairs = get_training_points(ekp, 5; g_final = G.(get_ϕ_final(prior, ekp))) # an additional iteration pairing (u_final, g_final) +``` + +[`encoder_kwargs_from(ekp, prior)`](@ref CalibrateEmulateSample.Utilities.encoder_kwargs_from) 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..bfa0cdb4d 100644 --- a/docs/src/emulate.md +++ b/docs/src/emulate.md @@ -11,6 +11,14 @@ 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" + The `EnsembleKalmanProcess` can be built with an `ObservationSeries` of `N` samples and a minibatcher with `n` batches. In this case, each iteration fits the ensemble to one of the batches ``B_1, \ldots, B_n`` (``n \leq N``) that form a disjoint partition of statistically similar observations ``y_1, \ldots, y_N \sim \rho``. The batching can aid the calibration stage but does not affect the emulate-sample stages. + + Internally `get_training_points` and `encoder_kwargs_from` reduce outputs/noise from the `EnsembleKalmanProcess` to their per-observation components, thus the emulator is trained on ``(\theta, \mathcal{G}(\theta))`` pairs representative of the distribution ``\rho``. In the Sampling stage the full observation set ``y_1, \ldots, y_N \sim \rho`` is used and batch structure is ignored. + + Warning: using training data with several `(identical-input)->(different-output)` pairs, will induce more sensitivity on the `obs_noise_cov`. It is recommended to ensure it is well estimated (or well-regularized with white noise) to help training; one can also use the `noise_learn=true` flag but this is known to cause slow-downs in training. + Wrapping a predefined machine learning tool, e.g. a Gaussian process `gauss_proc`, the `Emulator` can then be built: ```julia diff --git a/src/Utilities.jl b/src/Utilities.jl index ee7260277..91e9bb860 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -17,6 +17,7 @@ using TSVD import LinearAlgebra: norm export get_training_points +export get_flat_pairs export create_compact_linear_map export PairedDataContainerProcessor, DataContainerProcessor export create_encoder_schedule, @@ -54,6 +55,26 @@ Extract and flatten the training data from an `EnsembleKalmanProcess` into a - `train_iterations`: integer `n` (uses iterations `1:n`) or an index vector such as `3:2:9`. - `g_final` (keyword, default `nothing`): optional `AbstractMatrix` of forward-model outputs for the final parameter ensemble (not yet stored in `ekp`), sized as `get_g(ekp, 1)`. + +To ensure consistency in `g_final`, we recommend that, given the `ekp::EnsembleKalmanProcess` +and `prior::ParameterDistribution` used to generate it, the user evaluates the forward map `G` +on the parameters returned by `get_ϕ_final`: + +```julia +params = get_ϕ_final(prior, ekp) +g_final = reduce(hcat, [G(param) for param in eachcol(params)]) +``` + +If `ekp` is using minibatching, `G` must be evaluated on the same minibatch that `ekp` is +currently on: + +```julia +params = get_ϕ_final(prior, ekp) +batch_final = get_current_minibatch(ekp) +g_final = reduce(hcat, [G(param, batch_final) for param in eachcol(params)]) +``` + +See `EnsembleKalmanProcesses.Observations` for more on minibatching and `ObservationSeries`. """ function get_training_points( ekp::EKP.EnsembleKalmanProcess{FT, IT, P}, @@ -78,36 +99,15 @@ 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]] - 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)).")) - 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)).", - ), - ) - 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) + paired_range = [ir for ir in iter_range if ir <= g_len] - training_points = PairedDataContainer(u_tp, g_tp, data_are_columns = true) + # get_flat_pairs handles any batch size, including 1 (no-op for batch_size=1) + # final_samples_out appends the final ensemble output in per-element or batched form + u_tp, g_tp = get_flat_pairs(ekp, paired_range; final_samples_out = g_final) - return training_points + return PairedDataContainer(u_tp, g_tp, data_are_columns = true) end # Using Observation Objects: @@ -162,6 +162,121 @@ end """ $(TYPEDSIGNATURES) +Flatten the paired parameter and forward-model output ensembles from a minibatched +`EnsembleKalmanProcess` into individual `(parameter, single-observation)` column pairs. + +In each stored EKP iteration the output matrix `g` has `batch_size × gdim` rows, where +`gdim` is the single-observation dimension. This function splits each `g` into +`batch_size` blocks of `gdim` rows and replicates the corresponding parameter ensemble +so that each batch element becomes an independent `(u, g)` training pair. For +non-minibatched runs (`batch_size = 1`) the function is a no-op. + +# Arguments + +- `ekp`: `EnsembleKalmanProcess` holding the parameter ensembles and forward-model outputs. +- `iter_range`: iterations to include — an integer `n` is interpreted as `1:n` + (clamped to the number of stored outputs); a vector selects specific iterations, + which must be a subset of `1:length(get_g(ekp))`. +- `include_dt` (keyword, default `false`): when `true`, also return `dt_flat`, a + `Vector{Float64}` of per-pair algorithm times (`0` for iteration 1, then + `get_algorithm_time(ekp)[ir - 1]` repeated `batch_size` times per iteration). +- `final_samples_out` (keyword, default `nothing`): optional `AbstractMatrix` of + forward-model outputs for the final parameter ensemble `get_u_final(ekp)` (not yet + stored in `ekp`). Must be stacked as `[g_b1; g_b2; …]` matching the minibatch + structure of the next observation-series cycle. + +# Returns + +- `(u_flat, g_flat)` — matrices of size `dim_par × (Σ_k bs_k × N_ens)` and + `gdim × (Σ_k bs_k × N_ens)`, with parameters repeated once per batch element. +- `(u_flat, g_flat, dt_flat)` — as above, plus the algorithm-time vector, when + `include_dt = true`. +""" +function get_flat_pairs( + ekp::EKP.EnsembleKalmanProcess{FT, IT, P}, + iter_range::Union{IT, AbstractVector{IT}}; + include_dt::Bool = false, + final_samples_out = nothing, +) where {FT, IT, P} + if !isa(iter_range, AbstractVector) + iter_range = 1:min(iter_range, length(get_g(ekp))) + end + !isnothing(final_samples_out) && + !isa(final_samples_out, AbstractMatrix) && + throw( + ArgumentError( + "Expected `final_samples_out` to be an `AbstractMatrix`, received $(typeof(final_samples_out)).", + ), + ) + + 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 + + # if a final output is provided, treat it as one extra iteration beyond g_len + full_range = isnothing(final_samples_out) ? iter_range : [collect(iter_range); g_len + 1] + + for ir in full_range + if ir > g_len + # extra iteration: use the caller-supplied output instead of get_g + isnothing(final_samples_out) && 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)).", + ), + ) + gg = final_samples_out + else + gg = get_g(ekp, ir) + end + uu = get_u(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 + if ir > g_len && !isempty(g_out) && gdim != size(g_out[1], 1) + throw( + ArgumentError( + "final_samples_out has per-element output dimension $gdim, " * + "but existing outputs have dimension $(size(g_out[1], 1)).", + ), + ) + end + + 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,19 +331,36 @@ function encoder_kwargs_from( final_samples_out = nothing, ) where {PD <: ParameterDistribution, EKP <: EnsembleKalmanProcess} observation_series = isnothing(observation_series) ? get_observation_series(ekp) : observation_series + 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 - # a common setting is you have one-less samples_out than (samples_in, dt) so we extend g - if !isnothing(final_samples_out) - if (isa(final_samples_out, AbstractMatrix)) # add one matrix + + orig_g_len = length(get_g(ekp)) + is_minibatched = + auto_flatten_io && + orig_g_len > 0 && + any(length(get_minibatch(observation_series, ir)) > 1 for ir in 1:orig_g_len) + + if is_minibatched + # delegate all splitting (including final_samples_out) to get_flat_pairs + u_flat, g_flat, flat_dt = + get_flat_pairs(ekp, 1:orig_g_len; include_dt = true, final_samples_out = final_samples_out) + n_ens_per = size(get_u(ekp, 1), 2) + n_pairs = size(u_flat, 2) ÷ n_ens_per + samples_in = [u_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs] + samples_out = [g_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs] + dt = flat_dt + elseif !isnothing(final_samples_out) + # non-minibatch path: append final_samples_out directly + if isa(final_samples_out, AbstractMatrix) if length(samples_out) > 0 size(samples_out[1]) == size(final_samples_out) || error( "size of final_samples_out ($(size(final_samples_out))) does not match existing samples_out entries ($(size(samples_out[1])))", ) end push!(samples_out, final_samples_out) - else # add a vector of matrices + else if length(samples_out) > 0 all(size(samples_out[1]) == size(so) for so in final_samples_out) || error( "not all final_samples_out entries have size $(size(samples_out[1])) to match existing samples_out", diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl index 2b267bcae..524c1608f 100644 --- a/test/Utilities/runtests.jl +++ b/test/Utilities/runtests.jl @@ -689,6 +689,124 @@ 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 + + # ---- get_flat_pairs: shape ---- + u_flat, g_flat = get_flat_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: get_flat_pairs(ekp, 2) == get_flat_pairs(ekp, 1:2) + u_flat2, g_flat2 = get_flat_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 = get_flat_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 get_flat_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 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: not divisible by batch_size + @test_throws ArgumentError get_training_points( + ekp_mb, + n_iters, + g_final = randn(rng_mb, gdim * batch_size + 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 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