From 6bbd58405bf04fbbdfb9ef27753ceb6330bcf8ed Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 25 Jun 2026 14:27:21 -0700
Subject: [PATCH 01/12] docs for minibatcher tools

---
 docs/src/calibrate.md | 19 +++++++++++++++++++
 docs/src/emulate.md   |  6 ++++++
 2 files changed, 25 insertions(+)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index 44fad7bf1..da653a142 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -19,3 +19,22 @@ Documentation on how to construct an EnsembleKalmanProcess from the computer mod
 
 One draw of our approach is that it does not require the forward map to be written in Julia. To aid construction of such a workflow, EnsembleKalmanProcesses.jl provides a documented example of a BASH workflow for the [sinusoid problem](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/examples/sinusoid_example_toml/), with source code [here](https://github.com/CliMA/EnsembleKalmanProcesses.jl/tree/main/examples/SinusoidInterface). The forward map interacts with the calibration tools (EKP) only though TOML file reading an writing, and thus can be written in any language; for example, to be used with [slurm HPC scripts](https://clima.github.io/EnsembleKalmanProcesses.jl/dev/examples/ClimateMachine_example/), with source code [here](https://github.com/CliMA/EnsembleKalmanProcesses.jl/tree/main/examples/ClimateMachine).
 
+## Extracting training data after calibration
+
+Two helper functions from `CalibrateEmulateSample.Utilities` simplify passing the calibration results to the emulator.
+
+**`get_training_points(ekp, n)`** collects the stored parameter ensembles and forward-model outputs into a [`PairedDataContainer`](https://github.com/CliMA/EnsembleKalmanProcesses.jl/blob/main/src/DataContainers.jl) ready for emulator training. The argument `n` can be an integer (uses iterations `1:n`) or an index vector.  An optional keyword `g_final` accepts the forward-model outputs at the final, not-yet-stored parameter ensemble.
+
+```julia
+using CalibrateEmulateSample.Utilities
+input_output_pairs = get_training_points(ekp, 5)              # first 5 iterations
+input_output_pairs = get_training_points(ekp, 1:2:9)          # odd iterations 1,3,5,7,9
+input_output_pairs = get_training_points(ekp, 5; g_final = G(get_ϕ_final(prior, ekp)))
+```
+
+**`encoder_kwargs_from(ekp, prior)`** collects additional information from the calibration — the prior covariance, the observational noise covariance, and the sequence of parameter and output sample distributions across EKP iterations — into a named tuple that can be passed directly as the `encoder_kwargs` argument of the `Emulator`. This enables data-adaptive preprocessing such as likelihood-informed subspace reduction.
+
+```julia
+encoder_kwargs = encoder_kwargs_from(ekp, prior)
+```
+
diff --git a/docs/src/emulate.md b/docs/src/emulate.md
index bfb2d7f42..6f51c61f0 100644
--- a/docs/src/emulate.md
+++ b/docs/src/emulate.md
@@ -11,6 +11,12 @@ First, obtain data in a `PairedDataContainer`, for example, get this from an `En
 using CalibrateEmulateSample.Utilities
 input_output_pairs = Utilities.get_training_points(ekpobj, 5) # use first 5 iterations as data
 ```
+
+!!! note "Minibatched calibration"
+    When the `EnsembleKalmanProcess` is built with an `ObservationSeries` and a minibatcher, each iteration uses a batch ``B_1, \ldots, B_n`` (``n \leq N``) drawn as a disjoint partition of statistically similar observations ``y_1, \ldots, y_N \sim \rho``. Minibatching is a technique for accelerating the calibration stage; it does not change the inverse problem being solved.
+
+    `get_training_points` and `encoder_kwargs_from` therefore break the stacked batch outputs back into their per-observation components automatically, so the emulator is trained on ``(\theta, \mathcal{G}(\theta))`` pairs representative of the full distribution ``\rho``. In the Sampling stage the full observation set ``y_1, \ldots, y_N`` is used and batch structure is ignored.
+
 Wrapping a predefined machine learning tool, e.g. a Gaussian process `gauss_proc`, the `Emulator` can then be built:
 
 ```julia

From 3586fb88fc5855415989ebeddd3ca1e299fadd79 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 25 Jun 2026 14:29:03 -0700
Subject: [PATCH 02/12] add minibatch helpers into encoder_kwargs and
 get_training_points

---
 src/Utilities.jl           | 181 ++++++++++++++++++++++++++++++++-----
 test/Utilities/runtests.jl | 124 +++++++++++++++++++++++++
 2 files changed, 282 insertions(+), 23 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index ee7260277..3e8b35dc0 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -17,6 +17,7 @@ using TSVD
 import LinearAlgebra: norm
 
 export get_training_points
+export flatten_minibatch_pairs
 export create_compact_linear_map
 export PairedDataContainerProcessor, DataContainerProcessor
 export create_encoder_schedule,
@@ -78,36 +79,37 @@ function get_training_points(
         )
     end
 
-    u_tp = []
-    g_tp = []
-
+    os = get_observation_series(ekp)
     g_len = length(get_g(ekp))
-    g_full = [get_g(ekp, i) for i in iter_range[iter_range .<= g_len]]
+    paired_range = [ir for ir in iter_range if ir <= g_len]
+
+    # flatten_minibatch_pairs handles any batch size, including 1 (no-op for batch_size=1)
+    u_tp, g_tp = flatten_minibatch_pairs(ekp, paired_range)
+
     if !isnothing(g_final)
-        if (isa(g_final, AbstractMatrix)) # add the matrix
-            if !(size(g_full[1]) == size(g_final))
-                throw(ArgumentError("Expected `g_final` size: $(size(g_full[1])), received $(size(g_final))."))
+        isa(g_final, AbstractMatrix) ||
+            throw(ArgumentError("Expected `g_final` to be type `<:AbstractMatrix`, received $(typeof(g_final))."))
+        final_u = get_u(ekp, maximum(iter_range))
+        gdim = size(g_tp, 1)              # per-element output dimension
+        last_bs = length(get_minibatch(os, last(paired_range)))
+
+        if size(g_final, 1) == gdim
+            # g_final is in per-element (unbatched) form
+            u_tp = hcat(u_tp, final_u)
+            g_tp = hcat(g_tp, g_final)
+        elseif last_bs > 1 && size(g_final, 1) == gdim * last_bs
+            # g_final is stacked over batch elements — split it
+            for b in 1:last_bs
+                idx = (gdim * (b - 1) + 1):(gdim * b)
+                u_tp = hcat(u_tp, final_u)
+                g_tp = hcat(g_tp, g_final[idx, :])
             end
-            push!(g_full, g_final)
         else
-            throw(
-                ArgumentError(
-                    "Expected `g_final` to be type `<:AbstractMatrix`, of size: $(size(g_full[1])), received $(typeof(g_final)).",
-                ),
-            )
+            throw(ArgumentError("Expected `g_final` size: $((gdim, size(final_u, 2))), received $(size(g_final))."))
         end
     end
 
-    for (idx, ir) in enumerate(iter_range)
-        push!(u_tp, get_u(ekp, ir)) #N_parameters x N_ens
-        push!(g_tp, g_full[idx]) #N_data x N_ens
-    end
-    u_tp = reduce(hcat, u_tp) # N_parameters x (N_ek_it x N_ensemble)]
-    g_tp = reduce(hcat, g_tp) # N_data x (N_ek_it x N_ensemble)
-
-    training_points = PairedDataContainer(u_tp, g_tp, data_are_columns = true)
-
-    return training_points
+    return PairedDataContainer(u_tp, g_tp, data_are_columns = true)
 end
 
 # Using Observation Objects:
@@ -162,6 +164,94 @@ end
 """
 $(TYPEDSIGNATURES)
 
+Flatten the paired parameter and batched-output ensembles from an
+`EnsembleKalmanProcess` with minibatching into individual
+(parameter, single-observation) pairs.
+
+In a minibatched EKP run each stored output has row dimension `batch_size × gdim`,
+where `gdim` is the dimension of a single observation.  This function splits those
+stacked outputs into `batch_size` blocks of `gdim` rows and replicates the
+corresponding parameter ensemble, treating each batch element as an independent
+(parameter, observation) sample.
+
+`iter_range` must be a subset of `1:length(get_g(ekp))`.  An integer `n` is
+interpreted as `1:n` (clamped to the number of stored outputs).
+
+Returns `(u_flat, g_flat)` where
+- `u_flat` is `dim_par × (Σ_k bs_k × N_ens)` — parameters repeated per batch element.
+- `g_flat` is `gdim × (Σ_k bs_k × N_ens)` — outputs split to single-observation blocks.
+
+When `include_dt = true`, returns `(u_flat, g_flat, dt_flat)` where `dt_flat` is a
+`Vector{Float64}` of per-pair algorithm times: 0 for iteration 1 (the initial ensemble),
+and `get_algorithm_time(ekp)[ir - 1]` for subsequent iterations, each replicated
+`batch_size` times.
+
+Used internally by `get_training_points` and `encoder_kwargs_from` whenever a batch
+size > 1 is detected, so most users will not need to call this directly.
+"""
+function flatten_minibatch_pairs(
+    ekp::EKP.EnsembleKalmanProcess{FT, IT, P},
+    iter_range::Union{IT, AbstractVector{IT}};
+    include_dt::Bool = false,
+) where {FT, IT, P}
+    if !isa(iter_range, AbstractVector)
+        iter_range = 1:min(iter_range, length(get_g(ekp)))
+    end
+
+    os = get_observation_series(ekp)
+    g_len = length(get_g(ekp))
+    alg_time = include_dt ? get_algorithm_time(ekp) : nothing
+
+    u_out = []
+    g_out = []
+    dt_out = include_dt ? Float64[] : nothing
+
+    for ir in iter_range
+        if ir > g_len
+            throw(
+                ArgumentError(
+                    "iter_range contains iteration $ir, but EKP only stores $g_len " *
+                    "outputs. iter_range must be a subset of 1:length(get_g(ekp)).",
+                ),
+            )
+        end
+        uu = get_u(ekp, ir)
+        gg = get_g(ekp, ir)
+        batch = get_minibatch(os, ir)
+        bs = length(batch)
+
+        ndim = size(gg, 1)
+        ndim % bs == 0 || throw(
+            ArgumentError(
+                "At iteration $ir: output dimension ($ndim) is not divisible by batch " *
+                "size ($bs). Ensure outputs are stacked as [g_b1; g_b2; …] for each " *
+                "batch element.",
+            ),
+        )
+        gdim = ndim ÷ bs
+
+        for b in 1:bs
+            idx = (gdim * (b - 1) + 1):(gdim * b)
+            push!(u_out, uu)
+            push!(g_out, gg[idx, :])
+        end
+
+        if include_dt
+            di = (ir == 1) ? zero(eltype(alg_time)) : alg_time[ir - 1]
+            append!(dt_out, fill(di, bs))
+        end
+    end
+
+    if include_dt
+        return reduce(hcat, u_out), reduce(hcat, g_out), dt_out
+    else
+        return reduce(hcat, u_out), reduce(hcat, g_out)
+    end
+end
+
+"""
+$(TYPEDSIGNATURES)
+
 Extracts the relevant encoder kwargs from a vector triple (samples_in, samples_out, dt). Samples describe an ordered sequence of distributions in input and output space, each indexed with a temperature, or algorithm time, `dt`.
 
 Contains
@@ -216,6 +306,8 @@ function encoder_kwargs_from(
     final_samples_out = nothing,
 ) where {PD <: ParameterDistribution, EKP <: EnsembleKalmanProcess}
     observation_series = isnothing(observation_series) ? get_observation_series(ekp) : observation_series
+    # track whether samples were user-supplied (don't auto-flatten those)
+    auto_flatten_io = isnothing(samples_in) && isnothing(samples_out)
     samples_in = isnothing(samples_in) ? get_u(ekp) : samples_in
     samples_out = isnothing(samples_out) ? get_g(ekp) : samples_out
     dt = isnothing(dt) ? get_algorithm_time(ekp) : dt
@@ -238,6 +330,49 @@ function encoder_kwargs_from(
         end
     end
 
+    # when samples were derived from ekp, flatten any minibatched outputs
+    orig_g_len = length(get_g(ekp))
+    if auto_flatten_io &&
+       orig_g_len > 0 &&
+       any(length(get_minibatch(observation_series, ir)) > 1 for ir in 1:orig_g_len)
+
+        # delegate u/g/dt splitting to the public helper; rebuild per-pair lists from the result
+        u_flat, g_flat, flat_dt = flatten_minibatch_pairs(ekp, 1:orig_g_len; include_dt = true)
+        n_ens_per = size(get_u(ekp, 1), 2)
+        n_pairs = size(u_flat, 2) ÷ n_ens_per   # = Σ bs_k over stored iterations
+        flat_si = [u_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs]
+        flat_so = [g_flat[:, ((k - 1) * n_ens_per + 1):(k * n_ens_per)] for k in 1:n_pairs]
+
+        # handle final_samples_out that was appended beyond orig_g_len
+        if length(samples_out) > orig_g_len
+            final_so = samples_out[orig_g_len + 1]
+            gdim_ref = size(flat_so[1], 1)
+            last_bs = length(get_minibatch(observation_series, orig_g_len))
+            final_si = samples_in[min(orig_g_len + 1, length(samples_in))]
+            final_di = Float64(dt[min(orig_g_len, length(dt))])
+            if size(final_so, 1) == gdim_ref
+                push!(flat_si, final_si)
+                push!(flat_so, final_so)
+                push!(flat_dt, final_di)
+            elseif size(final_so, 1) == gdim_ref * last_bs
+                for b in 1:last_bs
+                    idx = (gdim_ref * (b - 1) + 1):(gdim_ref * b)
+                    push!(flat_si, final_si)
+                    push!(flat_so, final_so[idx, :])
+                    push!(flat_dt, final_di)
+                end
+            else
+                push!(flat_si, final_si)
+                push!(flat_so, final_so)
+                push!(flat_dt, final_di)
+            end
+        end
+
+        samples_in = flat_si
+        samples_out = flat_so
+        dt = flat_dt
+    end
+
     prior_kwargs = encoder_kwargs_from(prior)
     obs_kwargs = encoder_kwargs_from(observation_series)
     io_kwargs = encoder_kwargs_from(samples_in, samples_out, dt)
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 2b267bcae..604e61dd3 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -689,6 +689,130 @@ end
 
 
 
+@testset "Minibatch flattening" begin
+    rng_mb = Random.MersenneTwister(77)
+
+    gdim = 3        # dimension of one observation
+    n_obs = 4       # total observations in the series
+    dim_par = 2
+    n_ens = 5
+    batch_size = 2  # two batch elements per EKP iteration
+
+    obs_vec = [Observation(randn(rng_mb, gdim), Matrix{Float64}(I, gdim, gdim), "obs_$i") for i in 1:n_obs]
+    given_batches = [[1, 2], [3, 4]]
+    obs_series_mb = ObservationSeries(obs_vec, FixedMinibatcher(given_batches))
+
+    prior_mb = constrained_gaussian("test", 0.0, 1.0, -Inf, Inf, repeats = dim_par)
+    initial_ensemble_mb = construct_initial_ensemble(rng_mb, prior_mb, n_ens)
+    ekp_mb =
+        EnsembleKalmanProcesses.EnsembleKalmanProcess(initial_ensemble_mb, obs_series_mb, Inversion(); rng = rng_mb)
+
+    # store the raw stacked outputs for 2 iterations (one per FixedMinibatcher batch)
+    g_stored = Matrix{Float64}[]
+    for _ in 1:2
+        bs = length(get_current_minibatch(ekp_mb))
+        g_batch = randn(rng_mb, gdim * bs, n_ens)
+        push!(g_stored, copy(g_batch))
+        EnsembleKalmanProcesses.update_ensemble!(ekp_mb, g_batch)
+    end
+    n_iters = 2
+
+    # ---- flatten_minibatch_pairs: shape ----
+    u_flat, g_flat = flatten_minibatch_pairs(ekp_mb, 1:n_iters)
+    total_cols = n_iters * batch_size * n_ens
+    @test size(u_flat) == (dim_par, total_cols)
+    @test size(g_flat) == (gdim, total_cols)
+
+    # both batch-element copies of u at iteration 1 must equal get_u(ekp_mb, 1)
+    u1 = get_u(ekp_mb, 1)
+    @test all(isapprox.(u_flat[:, 1:n_ens], u1, atol = 1e-12))
+    @test all(isapprox.(u_flat[:, (n_ens + 1):(2n_ens)], u1, atol = 1e-12))
+
+    # g_flat slices match the two gdim-row blocks of the stored batched output
+    @test all(isapprox.(g_flat[:, 1:n_ens], g_stored[1][1:gdim, :], atol = 1e-12))
+    @test all(isapprox.(g_flat[:, (n_ens + 1):(2n_ens)], g_stored[1][(gdim + 1):(2gdim), :], atol = 1e-12))
+
+    # integer shorthand: flatten_minibatch_pairs(ekp, 2) == flatten_minibatch_pairs(ekp, 1:2)
+    u_flat2, g_flat2 = flatten_minibatch_pairs(ekp_mb, n_iters)
+    @test all(isapprox.(u_flat2, u_flat, atol = 1e-12))
+    @test all(isapprox.(g_flat2, g_flat, atol = 1e-12))
+
+    # include_dt: length and first batch_size entries are 0 (initial ensemble)
+    u_flat3, g_flat3, dt_flat = flatten_minibatch_pairs(ekp_mb, 1:n_iters; include_dt = true)
+    @test all(isapprox.(u_flat3, u_flat, atol = 1e-12))
+    @test all(isapprox.(g_flat3, g_flat, atol = 1e-12))
+    @test length(dt_flat) == total_cols ÷ n_ens
+    @test all(dt_flat[1:batch_size] .== 0)
+
+    # error: iter_range references an iteration with no stored output
+    @test_throws ArgumentError flatten_minibatch_pairs(ekp_mb, 1:5)
+
+    # ---- get_training_points: minibatch path ----
+    tp = get_training_points(ekp_mb, n_iters)
+    @test all(isapprox.(get_inputs(tp), u_flat, atol = 1e-12))
+    @test all(isapprox.(get_outputs(tp), g_flat, atol = 1e-12))
+
+    # g_final in per-element (unbatched) form: one extra column group appended
+    g_final_flat = randn(rng_mb, gdim, n_ens)
+    tp_f = get_training_points(ekp_mb, n_iters, g_final = g_final_flat)
+    @test size(get_inputs(tp_f), 2) == total_cols + n_ens
+    @test all(isapprox.(get_outputs(tp_f)[:, (total_cols + 1):end], g_final_flat, atol = 1e-12))
+
+    # g_final in batched (stacked) form: split into batch_size groups
+    g_final_batched = randn(rng_mb, gdim * batch_size, n_ens)
+    tp_fb = get_training_points(ekp_mb, n_iters, g_final = g_final_batched)
+    @test size(get_inputs(tp_fb), 2) == total_cols + batch_size * n_ens
+    @test all(
+        isapprox.(
+            get_outputs(tp_fb)[:, (total_cols + 1):(total_cols + n_ens)],
+            g_final_batched[1:gdim, :],
+            atol = 1e-12,
+        ),
+    )
+    @test all(
+        isapprox.(
+            get_outputs(tp_fb)[:, (total_cols + n_ens + 1):end],
+            g_final_batched[(gdim + 1):(2gdim), :],
+            atol = 1e-12,
+        ),
+    )
+
+    # wrong g_final size
+    @test_throws ArgumentError get_training_points(ekp_mb, n_iters, g_final = randn(rng_mb, gdim + 1, n_ens))
+
+    # ---- encoder_kwargs_from: minibatch path ----
+    ekp_kw = encoder_kwargs_from(ekp_mb, prior_mb)
+
+    # prior and obs kwargs are structurally unchanged
+    @test ekp_kw[:obs_noise_cov] == encoder_kwargs_from(obs_series_mb)[:obs_noise_cov]
+    @test all(isapprox.(ekp_kw[:prior_cov], encoder_kwargs_from(prior_mb)[:prior_cov], atol = 1e-12))
+
+    # samples_in/out lists reflect per-element (flattened) structure
+    flat_si_list = ekp_kw[:input_structure_vecs][:samples_in]
+    flat_so_list = ekp_kw[:output_structure_vecs][:samples_out]
+    @test length(flat_si_list) == n_iters * batch_size
+    @test length(flat_so_list) == n_iters * batch_size
+    @test size(flat_si_list[1]) == (dim_par, n_ens)
+    @test size(flat_so_list[1]) == (gdim, n_ens)
+
+    # dt: first batch_size entries correspond to the initial distribution (time = 0)
+    flat_dt_kw = ekp_kw[:input_structure_vecs][:dt]
+    @test all(flat_dt_kw[1:batch_size] .== 0)
+
+    # final_samples_out in per-element form → one extra pair appended
+    ekp_kw_f = encoder_kwargs_from(ekp_mb, prior_mb, final_samples_out = g_final_flat)
+    @test length(ekp_kw_f[:output_structure_vecs][:samples_out]) == n_iters * batch_size + 1
+
+    # final_samples_out in batched form → split into batch_size extra pairs
+    ekp_kw_fb = encoder_kwargs_from(ekp_mb, prior_mb, final_samples_out = g_final_batched)
+    @test length(ekp_kw_fb[:output_structure_vecs][:samples_out]) == n_iters * batch_size + batch_size
+
+    # user-supplied samples_in/out are NOT auto-flattened
+    ekp_kw_custom = encoder_kwargs_from(ekp_mb, prior_mb; samples_in = get_u(ekp_mb), samples_out = get_g(ekp_mb))
+    @test size(ekp_kw_custom[:output_structure_vecs][:samples_out][1], 1) == gdim * batch_size
+
+end
+
 @testset "Decorrelator: Large observational covariance" begin
 
     # loop over output dim

From b7a20e56ecca7b5e9361719bff18820630050a41 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 25 Jun 2026 14:59:06 -0700
Subject: [PATCH 03/12] improve name

---
 src/Utilities.jl           | 10 +++++-----
 test/Utilities/runtests.jl | 12 ++++++------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 3e8b35dc0..3c68d2ca6 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -17,7 +17,7 @@ using TSVD
 import LinearAlgebra: norm
 
 export get_training_points
-export flatten_minibatch_pairs
+export get_flat_pairs
 export create_compact_linear_map
 export PairedDataContainerProcessor, DataContainerProcessor
 export create_encoder_schedule,
@@ -83,8 +83,8 @@ function get_training_points(
     g_len = length(get_g(ekp))
     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)
+    # get_flat_pairs handles any batch size, including 1 (no-op for batch_size=1)
+    u_tp, g_tp = get_flat_pairs(ekp, paired_range)
 
     if !isnothing(g_final)
         isa(g_final, AbstractMatrix) ||
@@ -189,7 +189,7 @@ and `get_algorithm_time(ekp)[ir - 1]` for subsequent iterations, each replicated
 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(
+function get_flat_pairs(
     ekp::EKP.EnsembleKalmanProcess{FT, IT, P},
     iter_range::Union{IT, AbstractVector{IT}};
     include_dt::Bool = false,
@@ -337,7 +337,7 @@ function encoder_kwargs_from(
        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)
+        u_flat, g_flat, flat_dt = get_flat_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]
diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl
index 604e61dd3..2d92bf980 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -717,8 +717,8 @@ end
     end
     n_iters = 2
 
-    # ---- flatten_minibatch_pairs: shape ----
-    u_flat, g_flat = flatten_minibatch_pairs(ekp_mb, 1:n_iters)
+    # ---- 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)
@@ -732,20 +732,20 @@ end
     @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)
+    # 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 = flatten_minibatch_pairs(ekp_mb, 1:n_iters; include_dt = true)
+    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 flatten_minibatch_pairs(ekp_mb, 1:5)
+    @test_throws ArgumentError get_flat_pairs(ekp_mb, 1:5)
 
     # ---- get_training_points: minibatch path ----
     tp = get_training_points(ekp_mb, n_iters)

From 3cb20db278e61053524d9a2775541300ac81816c Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Thu, 25 Jun 2026 19:26:36 -0700
Subject: [PATCH 04/12] fix test failures, and improve layout. Add dosctring

---
 src/Utilities.jl           | 169 ++++++++++++++++---------------------
 test/Utilities/runtests.jl |  18 ++--
 2 files changed, 79 insertions(+), 108 deletions(-)

diff --git a/src/Utilities.jl b/src/Utilities.jl
index 3c68d2ca6..ffa04b85b 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -84,30 +84,8 @@ function get_training_points(
     paired_range = [ir for ir in iter_range if ir <= g_len]
 
     # get_flat_pairs handles any batch size, including 1 (no-op for batch_size=1)
-    u_tp, g_tp = get_flat_pairs(ekp, paired_range)
-
-    if !isnothing(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
-        else
-            throw(ArgumentError("Expected `g_final` size: $((gdim, size(final_u, 2))), received $(size(g_final))."))
-        end
-    end
+    # 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 PairedDataContainer(u_tp, g_tp, data_are_columns = true)
 end
@@ -164,39 +142,52 @@ 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.
+Flatten the paired parameter and forward-model output ensembles from a minibatched
+`EnsembleKalmanProcess` into individual `(parameter, single-observation)` column pairs.
 
-`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).
+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.
 
-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.
+# Arguments
 
-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.
+- `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))
@@ -206,17 +197,23 @@ function get_flat_pairs(
     g_out = []
     dt_out = include_dt ? Float64[] : nothing
 
-    for ir in iter_range
+    # 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
-            throw(
+            # 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)
-        gg = get_g(ekp, ir)
         batch = get_minibatch(os, ir)
         bs = length(batch)
 
@@ -229,6 +226,14 @@ function get_flat_pairs(
             ),
         )
         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)
@@ -306,21 +311,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
-    # 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
-    # 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",
@@ -330,49 +350,6 @@ 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 = get_flat_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 2d92bf980..524c1608f 100644
--- a/test/Utilities/runtests.jl
+++ b/test/Utilities/runtests.jl
@@ -752,12 +752,6 @@ end
     @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)
@@ -777,8 +771,12 @@ end
         ),
     )
 
-    # wrong g_final size
-    @test_throws ArgumentError get_training_points(ekp_mb, n_iters, g_final = randn(rng_mb, gdim + 1, n_ens))
+    # 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)
@@ -799,10 +797,6 @@ end
     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

From fb911b18a293bf255025c637997747f9aabc1d6d Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Fri, 26 Jun 2026 11:20:31 -0700
Subject: [PATCH 05/12] add note about noise sensitivity

---
 docs/src/emulate.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docs/src/emulate.md b/docs/src/emulate.md
index 6f51c61f0..a563bce1f 100644
--- a/docs/src/emulate.md
+++ b/docs/src/emulate.md
@@ -13,9 +13,11 @@ input_output_pairs = Utilities.get_training_points(ekpobj, 5) # use first 5 iter
 ```
 
 !!! 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.
+    When the `EnsembleKalmanProcess` is built with an `ObservationSeries` and a minibatcher, each iteration uses a batch ``B_1, \ldots, B_n`` (``n \leq N``): a disjoint partition of statistically similar observations ``y_1, \ldots, y_N \sim \rho``. The batching can aid the calibration stage but does not effect the emulate-sample stages.
 
-    `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.
+    Internally `get_training_points` and `encoder_kwargs_from` reduce outputs/noise from the `EnsembleKalmanProcess` to thei 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`` 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:
 

From f72f57ae9a404127320c33381ee00ff8f4269fa8 Mon Sep 17 00:00:00 2001
From: Oliver Dunbar <47412152+odunbar@users.noreply.github.com>
Date: Tue, 30 Jun 2026 15:09:50 -0700
Subject: [PATCH 06/12] Update docs/src/calibrate.md

Co-authored-by: Kevin Phan  <98072684+ph-kev@users.noreply.github.com>
---
 docs/src/calibrate.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index da653a142..04e468ba3 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -23,7 +23,7 @@ One draw of our approach is that it does not require the forward map to be writt
 
 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.
+**`get_training_points(ekp, n; g_final = nothing)`** 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 (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

From 224e08b84d373e27e8eebc4780deb18b970f8e76 Mon Sep 17 00:00:00 2001
From: Oliver Dunbar <47412152+odunbar@users.noreply.github.com>
Date: Tue, 30 Jun 2026 15:11:14 -0700
Subject: [PATCH 07/12] Update docs/src/emulate.md

Co-authored-by: Kevin Phan  <98072684+ph-kev@users.noreply.github.com>
---
 docs/src/emulate.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/emulate.md b/docs/src/emulate.md
index a563bce1f..88e1abf54 100644
--- a/docs/src/emulate.md
+++ b/docs/src/emulate.md
@@ -13,7 +13,7 @@ input_output_pairs = Utilities.get_training_points(ekpobj, 5) # use first 5 iter
 ```
 
 !!! 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``): a disjoint partition of statistically similar observations ``y_1, \ldots, y_N \sim \rho``. The batching can aid the calibration stage but does not effect the emulate-sample stages.
+    When the `EnsembleKalmanProcess` is built with an `ObservationSeries` and a minibatcher, each iteration uses a batch ``B_1, \ldots, B_n`` (``n \leq N``): 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 thei 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`` is used and batch structure is ignored.
 

From b2ce37b47d1e4590aaf66d6689d7fb3644e25adb Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 30 Jun 2026 16:02:09 -0700
Subject: [PATCH 08/12] resolve comments

---
 docs/src/calibrate.md |  4 ++--
 src/Utilities.jl      | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index 04e468ba3..4e3fb2fe7 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -23,7 +23,7 @@ One draw of our approach is that it does not require the forward map to be writt
 
 Two helper functions from `CalibrateEmulateSample.Utilities` simplify passing the calibration results to the emulator.
 
-**`get_training_points(ekp, n; g_final = nothing)`** 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 (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.
+[`get_training_points(ekp, n; g_final = nothing)`](@ref get_training_points) 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 (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
@@ -32,7 +32,7 @@ input_output_pairs = get_training_points(ekp, 1:2:9)          # odd iterations 1
 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.
+[`encoder_kwargs_from(ekp, prior)`](@ref 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/src/Utilities.jl b/src/Utilities.jl
index ffa04b85b..91e9bb860 100644
--- a/src/Utilities.jl
+++ b/src/Utilities.jl
@@ -55,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},

From 6587b31c370bf4dfcdc8bc38fcd6932aa7968078 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 30 Jun 2026 16:09:03 -0700
Subject: [PATCH 09/12] clarify n,N etc.

---
 docs/src/emulate.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/emulate.md b/docs/src/emulate.md
index 88e1abf54..bfa0cdb4d 100644
--- a/docs/src/emulate.md
+++ b/docs/src/emulate.md
@@ -13,9 +13,9 @@ input_output_pairs = Utilities.get_training_points(ekpobj, 5) # use first 5 iter
 ```
 
 !!! 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``): 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.
+    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 thei 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`` is used and batch structure is ignored.
+    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.
 

From a4c9eb9564679e79484a6443d61071f7c5857fb1 Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 30 Jun 2026 16:31:15 -0700
Subject: [PATCH 10/12] try adding Module

---
 docs/src/calibrate.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index 4e3fb2fe7..6eacb0270 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -23,7 +23,7 @@ One draw of our approach is that it does not require the forward map to be writt
 
 Two helper functions from `CalibrateEmulateSample.Utilities` simplify passing the calibration results to the emulator.
 
-[`get_training_points(ekp, n; g_final = nothing)`](@ref get_training_points) 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 (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.
+[`get_training_points(ekp, n; g_final = nothing)`](@ref Utilities.get_training_points) 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 (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
@@ -32,7 +32,7 @@ input_output_pairs = get_training_points(ekp, 1:2:9)          # odd iterations 1
 input_output_pairs = get_training_points(ekp, 5; g_final = G(get_ϕ_final(prior, ekp)))
 ```
 
-[`encoder_kwargs_from(ekp, prior)`](@ref 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.
+[`encoder_kwargs_from(ekp, prior)`](@ref 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)

From 79f136d20fd8a9c695dd15a6fe0c76d3a30e9f5e Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 30 Jun 2026 17:45:59 -0700
Subject: [PATCH 11/12] qualify module even more?

---
 docs/src/calibrate.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index 6eacb0270..f72f836c8 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -23,7 +23,7 @@ One draw of our approach is that it does not require the forward map to be writt
 
 Two helper functions from `CalibrateEmulateSample.Utilities` simplify passing the calibration results to the emulator.
 
-[`get_training_points(ekp, n; g_final = nothing)`](@ref Utilities.get_training_points) 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 (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.
+[`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://github.com/CliMA/EnsembleKalmanProcesses.jl/blob/main/src/DataContainers.jl) 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
@@ -32,7 +32,7 @@ input_output_pairs = get_training_points(ekp, 1:2:9)          # odd iterations 1
 input_output_pairs = get_training_points(ekp, 5; g_final = G(get_ϕ_final(prior, ekp)))
 ```
 
-[`encoder_kwargs_from(ekp, prior)`](@ref 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.
+[`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)

From 4f7d4d3c3941c86429b325cb32ab81084b6165fc Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Wed, 1 Jul 2026 10:57:48 -0700
Subject: [PATCH 12/12] pdc

---
 docs/src/calibrate.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/src/calibrate.md b/docs/src/calibrate.md
index f72f836c8..fa3231e02 100644
--- a/docs/src/calibrate.md
+++ b/docs/src/calibrate.md
@@ -23,13 +23,13 @@ One draw of our approach is that it does not require the forward map to be writt
 
 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://github.com/CliMA/EnsembleKalmanProcesses.jl/blob/main/src/DataContainers.jl) 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.
+[`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 1,3,5,7,9
-input_output_pairs = get_training_points(ekp, 5; g_final = G(get_ϕ_final(prior, ekp)))
+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.