Minibatched data-handling

.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*)",

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)
+```
+

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

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

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

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"