diff --git a/examples/Lorenz/GModel.jl b/examples/Lorenz/GModel.jl index c9e00a6f3..fd2820f00 100644 --- a/examples/Lorenz/GModel.jl +++ b/examples/Lorenz/GModel.jl @@ -1,437 +1,18 @@ module GModel -using DocStringExtensions - -using Random -using Distributions -using LinearAlgebra -#using DifferentialEquations -#using OrdinaryDiffEq -using FFTW - export run_G export run_G_ensemble export lorenz_forward -""" -$(DocStringExtensions.TYPEDEF) - -Structure to hold all information to run the forward model *G*. - -# Fields -$(DocStringExtensions.TYPEDFIELDS) -""" -struct GSettings{FT <: AbstractFloat, KT, D} - "Model output." - out::Array{FT, 1} - "Time period over which to run the model, e.g., `(0, 1)`." - tspan::Tuple{FT, FT} - "X domain." - x::Array{FT, 1} - "Y domain." - y::Array{FT, 1} -end - -struct LSettings - # Stationary or transient dynamics - dynamics::Int32 - # G model statistics type - stats_type::Int32 - # Integration time start - t_start::Float64 - # Data collection length - T::Float64 - # Integration length - Ts::Float64 - # Duration of polynomial fit, number of samples of Ts to fit over - Tfit::Float64 - # Initial perturbation - Fp::Array{Float64} - # Number of longitude steps - N::Int32 - # Timestep - dt::Float64 - # Simulation end time - tend::Float64 - # For stats_type=2, number of frequencies to consider - kmax::Int32 - # Should ω be learned using CES? - ω_fixed::Bool - # Truth ω - ω_true::Float64 -end - -struct LParams - # Mean forcing - F::Float64 - # Forcing frequency for transient terms - ω::Float64 - # Forcing amplitude for transient terms - A::Float64 -end - -""" -$(DocStringExtensions.TYPEDSIGNATURES) - -Run the forward model *G* for an array of parameters by iteratively -calling `run_G` for each of the *N\\_ensemble* parameter values. +include("GModel_common.jl") -- `params` - array of size (*N\\_ensemble* × *N\\_parameters*) containing the parameters for - which *G* will be run. -- `settings` - a [GSetttings](@ref) struct. - -Returns `g_ens`, an array of size (*N\\_ensemble* × *N\\_data*), where g_ens[j,:] = G(params[j,:]). -""" -function run_G_ensemble(params::Array{FT, 2}, settings::LSettings, rng_seed = 42) where {FT <: AbstractFloat} - # Initialize ensemble - N_ens = size(params, 2) # params is N_params x N_ens - if settings.stats_type == 1 - nd = 2 - elseif settings.stats_type == 2 - nd = 1 + (2 * settings.kmax) - elseif settings.stats_type == 3 - nd = 3 - elseif settings.stats_type == 4 - nd = Int64(2 * (settings.T / settings.Ts / settings.Tfit)) - elseif settings.stats_type == 5 - nd = Int64(settings.T / settings.Ts / settings.Tfit) - end +function run_ensembles(settings, lorenz_params, nd, N_ens) g_ens = zeros(nd, N_ens) - # Lorenz parameters - #Fp = rand(Normal(0.0, 0.01), N); # Initial perturbation - F = params[1, :] # Forcing - if settings.dynamics == 2 - A = params[2, :] # Transience amplitude - else - A = F .* 0.0 # Set to zero, it is not used - end - if settings.ω_fixed == false - ω = params[3, :] # Transience frequency - end - - Random.seed!(rng_seed) for i in 1:N_ens # run the model with the current parameters, i.e., map θ to G(θ) - if settings.ω_fixed == false - lorenz_params = GModel.LParams(F[i], ω[i], A[i]) - elseif settings.ω_fixed == true - lorenz_params = GModel.LParams(F[i], settings.ω_true, A[i]) - end - g_ens[:, i] = lorenz_forward(settings, lorenz_params) + g_ens[:, i] = lorenz_forward(settings, lorenz_params[i]) end - return g_ens end -# Forward pass of the Lorenz 96 model -# Inputs: settings: structure with stats_type, t_start, T -# F: scalar forcing Float64(1), Fp: initial perturbation Float64(N) -# N: number of longitude steps Int64(1), dt: time step Float64(1) -# tend: End of simulation Float64(1), nstep: -function lorenz_forward(settings::LSettings, params::LParams) - # run the Lorenz simulation - xn, t = lorenz_solve(settings, params) - # Get statistics - gt = stats(settings, xn, t) - return gt -end - - -# Solve the Lorenz 96 system -# Inputs: F: scalar forcing Float64(1), Fp: initial perturbation Float64(N) -# N: number of longitude steps Int64(1), dt: time step Float64(1) -# tend: End of simulation Float64(1), nstep: -function lorenz_solve(settings::LSettings, params::LParams) - # Initialize - nstep = Int32(ceil(settings.tend / settings.dt)) - xn = zeros(Float64, settings.N, nstep) - t = zeros(Float64, nstep) - # Initial perturbation - X = fill(Float64(0.0), settings.N) - X = X + settings.Fp - # March forward in time - for j in 1:nstep - t[j] = settings.dt * j - if settings.dynamics == 1 - X = RK4(X, settings.dt, settings.N, params.F) - elseif settings.dynamics == 2 - F_local = params.F + params.A * sin(params.ω * t[j]) - X = RK4(X, settings.dt, settings.N, F_local) - end - xn[:, j] = X - end - # Output - return xn, t - -end - -# Lorenz 96 system -# f = dx/dt -# Inputs: x: state, N: longitude steps, F: forcing -function f(x, N, F) - f = zeros(Float64, 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] + F - end - # Periodic boundary conditions - f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + F - f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + F - f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + F - # Output - return f -end - -# RK4 solve -function RK4(xold, dt, N, F) - # Predictor steps - k1 = f(xold, N, F) - k2 = f(xold + k1 * dt / 2.0, N, F) - k3 = f(xold + k2 * dt / 2.0, N, F) - k4 = f(xold + k3 * dt, N, F) - # Step - xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) - # Output - return xnew -end - -function regression(X, y) - beta_tilde = [ones(length(X), 1) X] \ y - v = beta_tilde[1] - beta = beta_tilde[2] - return beta, v -end - -function spectra(signal, Ts, t) - # Init - Nt = length(t) - if mod(Nt, 2) != 0 - t = t[1:(Nt - 1)] - signal = signal[1:(Nt - 1)] - Nt = Nt - 1 - end - # Remove mean - u = signal .- mean(signal, dims = 1) - # FFT - F = fft(u) |> fftshift - freqs = fftfreq(length(t), 1.0 / Ts) |> fftshift - #k = gen_k(Nt,Ts) |> fftshift - T = (Nt - 1) * Ts - A = Ts^2 / (2 * pi * T) - uhat = A * abs.(F .^ 2) - # Output - mid = Int32(Nt / 2) + 1 - f = freqs[mid:Nt] - Fp = uhat[mid:Nt] - return f, Fp -end - -############################### -## Extract statistics -############################### -function stats(settings, xn, t) - # Define averaging indices range - indices = findall(x -> (x > settings.t_start) && (x < settings.t_start + settings.T), t) - # Define statistics of interest - if settings.stats_type == 1 # Mean - # Average in time and over longitude - gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) - # Variance - gtc = mean(mean((xn[:, indices] .- gtm) .^ 2, dims = 2), dims = 1) - # Combine statistics - gt = vcat(gtm, gtc) - elseif settings.stats_type == 2 #What about an integral under parts of the spectrum - # Average in time and over longitude - gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) - # Power spectra - # Get size - f, Fp = spectra(xn[1, indices]', settings.dt, t[indices]) - # Compute temporal spectrum for each longitude - Fp = zeros(size(Fp, 1), settings.N) - for i in 1:(settings.N) - f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices]) - end - # Average spectra over periodic directions - Fp = dropdims(mean(Fp, dims = 2), dims = 2) - ys = partialsortperm(Fp, 1:(settings.kmax); rev = true) - gt = vcat(gtm, Fp[ys]..., f[ys]...) - elseif settings.stats_type == 3 # Structure function - # Average in time and over longitude - gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) - # Maximum - mxval, mxind = findmax(xn[:, indices]; dims = 2) - mxval_out = mean(mxval, dims = 1) - # Power spectra - # Get size - f, Fp = spectra(xn[1, indices]', settings.dt, t[indices]) - # Compute temporal spectrum for each longitude - Fp = zeros(size(Fp, 1), settings.N) - for i in 1:(settings.N) - f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices]) - end - # Average spectra over periodic directions - Fp = dropdims(mean(Fp, dims = 2), dims = 2) - ys = partialsortperm(Fp, 1; rev = true) - # Period - T = 1.0 / f[ys] - mxval, r = findmin(abs.(t .- T); dims = 1) - r = r[1] - # Structure function - xp = xn .- gtm - st = zeros(settings.N) - for i in 1:(settings.N) - st[i] = mean( - ( - xp[1, indices[1]:(indices[length(indices)] - r)] .- - xp[1, (indices[1] + r):indices[length(indices)]] - ) .^ 2, - ) - end - # Combine - gt = vcat(gtm, mean(st)..., mxval_out...) - elseif settings.stats_type == 4 # Variance sub-samples, linear fit over window - T = settings.T - # Calculate variance on subsamples - # Variances over sliding windows - Ts = settings.Ts - N = Int64(T / Ts) # number of sliding windows - nt = Int64(Ts / settings.dt) - 1 # length of sliding windows - var_slide = zeros(N) - t_slide = zeros(nt, N) - for i in 1:N - var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...)) - t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)] - end - # Polynomial fit over a batch of sliding windows - Tfit = Int64(settings.Tfit) - NT = Int64(N / Tfit) - a1 = zeros(NT) - a2 = zeros(NT) - t_start = zeros(NT) - y = zeros(Tfit, NT) - ty = zeros(Tfit, NT) - ym = zeros(Tfit, NT) - for i in 1:NT - ind_low = (i - 1) * Tfit + 1 - ind_high = i * Tfit - t_start[i] = t_slide[1, ind_low] - ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i] - a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high]) - y[:, i] = a1[i] .* ty[:, i] .+ a2[i] - ym[i] = mean(var_slide[ind_low:ind_high]) - end - # Combine - gt = vcat(a1..., a2...) - elseif settings.stats_type == 5 # Variance sub-samples, mean over window - T = settings.T - # Calculate variance on subsamples - # Variances over sliding windows - Ts = settings.Ts - N = Int64(T / Ts) # number of sliding windows - nt = Int64(Ts / settings.dt) - 1 # length of sliding windows - var_slide = zeros(N) - t_slide = zeros(nt, N) - for i in 1:N - var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...)) - t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)] - end - # Polynomial fit over a batch of sliding windows - Tfit = Int64(settings.Tfit) - NT = Int64(N / Tfit) - a1 = zeros(NT) - a2 = zeros(NT) - t_start = zeros(NT) - y = zeros(Tfit, NT) - ty = zeros(Tfit, NT) - ym = zeros(NT) - for i in 1:NT - ind_low = (i - 1) * Tfit + 1 - ind_high = i * Tfit - t_start[i] = t_slide[1, ind_low] - ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i] - a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high]) - y[:, i] = a1[i] .* ty[:, i] .+ a2[i] - ym[i] = mean(var_slide[ind_low:ind_high]) - end - # Combine - gt = vcat(ym...) - else - ArgumentError("Setting " * string(settings.stats_type) * " not implemented.") - end - return gt -end - - - - -## Forward pass of the Lorenz 96 model -## Inputs: settings: structure with stats_type, t_start, T -## F: scalar forcing Float64(1), Fp: initial perturbation Float64(N) -## N: number of longitude steps Int64(1), dt: time step Float64(1) -## tend: End of simulation Float64(1), nstep: -#function lorenz_forward(settings::LSettings, F) -# # run the Lorenz simulation -# xn, t = lorenz_solve(settings, F) -# # Define averaging indices range -# indices = findall(x -> (x>settings.t_start) -# && (x fftshift + freqs = fftfreq(length(t), 1.0 / Ts) |> fftshift + #k = gen_k(Nt,Ts) |> fftshift + T = (Nt - 1) * Ts + A = Ts^2 / (2 * pi * T) + uhat = A * abs.(F .^ 2) + # Output + mid = Int32(Nt / 2) + 1 + f = freqs[mid:Nt] + Fp = uhat[mid:Nt] + return f, Fp +end + +############################### +## Extract statistics +############################### +function stats(settings, xn, t) + # Define averaging indices range + indices = findall(x -> (x > settings.t_start) && (x < settings.t_start + settings.T), t) + # Define statistics of interest + if settings.stats_type == 1 # Mean + # Average in time and over longitude + gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) + # Variance + gtc = mean(mean((xn[:, indices] .- gtm) .^ 2, dims = 2), dims = 1) + # Combine statistics + gt = vcat(gtm, gtc) + elseif settings.stats_type == 2 #What about an integral under parts of the spectrum + # Average in time and over longitude + gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) + # Power spectra + # Get size + f, Fp = spectra(xn[1, indices]', settings.dt, t[indices]) + # Compute temporal spectrum for each longitude + Fp = zeros(size(Fp, 1), settings.N) + for i in 1:(settings.N) + f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices]) + end + # Average spectra over periodic directions + Fp = dropdims(mean(Fp, dims = 2), dims = 2) + ys = partialsortperm(Fp, 1:(settings.kmax); rev = true) + gt = vcat(gtm, Fp[ys]..., f[ys]...) + elseif settings.stats_type == 3 # Structure function + # Average in time and over longitude + gtm = mean(mean(xn[:, indices], dims = 2), dims = 1) + # Maximum + mxval, mxind = findmax(xn[:, indices]; dims = 2) + mxval_out = mean(mxval, dims = 1) + # Power spectra + # Get size + f, Fp = spectra(xn[1, indices]', settings.dt, t[indices]) + # Compute temporal spectrum for each longitude + Fp = zeros(size(Fp, 1), settings.N) + for i in 1:(settings.N) + f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices]) + end + # Average spectra over periodic directions + Fp = dropdims(mean(Fp, dims = 2), dims = 2) + ys = partialsortperm(Fp, 1; rev = true) + # Period + T = 1.0 / f[ys] + mxval, r = findmin(abs.(t .- T); dims = 1) + r = r[1] + # Structure function + xp = xn .- gtm + st = zeros(settings.N) + for i in 1:(settings.N) + st[i] = mean( + ( + xp[1, indices[1]:(indices[length(indices)] - r)] .- + xp[1, (indices[1] + r):indices[length(indices)]] + ) .^ 2, + ) + end + # Combine + gt = vcat(gtm, mean(st)..., mxval_out...) + elseif settings.stats_type == 4 # Variance sub-samples, linear fit over window + T = settings.T + # Calculate variance on subsamples + # Variances over sliding windows + Ts = settings.Ts + N = Int64(T / Ts) # number of sliding windows + nt = Int64(Ts / settings.dt) - 1 # length of sliding windows + var_slide = zeros(N) + t_slide = zeros(nt, N) + for i in 1:N + var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...)) + t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)] + end + # Polynomial fit over a batch of sliding windows + Tfit = Int64(settings.Tfit) + NT = Int64(N / Tfit) + a1 = zeros(NT) + a2 = zeros(NT) + t_start = zeros(NT) + y = zeros(Tfit, NT) + ty = zeros(Tfit, NT) + ym = zeros(Tfit, NT) + for i in 1:NT + ind_low = (i - 1) * Tfit + 1 + ind_high = i * Tfit + t_start[i] = t_slide[1, ind_low] + ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i] + a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high]) + y[:, i] = a1[i] .* ty[:, i] .+ a2[i] + ym[i] = mean(var_slide[ind_low:ind_high]) + end + # Combine + gt = vcat(a1..., a2...) + elseif settings.stats_type == 5 # Variance sub-samples, mean over window + T = settings.T + # Calculate variance on subsamples + # Variances over sliding windows + Ts = settings.Ts + N = Int64(T / Ts) # number of sliding windows + nt = Int64(Ts / settings.dt) - 1 # length of sliding windows + var_slide = zeros(N) + t_slide = zeros(nt, N) + for i in 1:N + var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...)) + t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)] + end + # Polynomial fit over a batch of sliding windows + Tfit = Int64(settings.Tfit) + NT = Int64(N / Tfit) + a1 = zeros(NT) + a2 = zeros(NT) + t_start = zeros(NT) + y = zeros(Tfit, NT) + ty = zeros(Tfit, NT) + ym = zeros(NT) + for i in 1:NT + ind_low = (i - 1) * Tfit + 1 + ind_high = i * Tfit + t_start[i] = t_slide[1, ind_low] + ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i] + a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high]) + y[:, i] = a1[i] .* ty[:, i] .+ a2[i] + ym[i] = mean(var_slide[ind_low:ind_high]) + end + # Combine + gt = vcat(ym...) + else + ArgumentError("Setting " * string(settings.stats_type) * " not implemented.") + end + return gt +end diff --git a/examples/Lorenz/Lorenz_example.jl b/examples/Lorenz/Lorenz_example.jl deleted file mode 100644 index de638c23b..000000000 --- a/examples/Lorenz/Lorenz_example.jl +++ /dev/null @@ -1,428 +0,0 @@ -# Reference the in-tree version of CalibrateEmulateSample on Julias load path - -# Import modules -include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) -include(joinpath(@__DIR__, "GModel.jl")) # Contains Lorenz 96 source code - -# Import modules -using Distributions # probability distributions and associated functions -using LinearAlgebra -using StatsPlots -using GaussianProcesses -using Plots -using Random -using JLD2 - -# CES -using CalibrateEmulateSample.Emulators -using CalibrateEmulateSample.MarkovChainMonteCarlo -using CalibrateEmulateSample.Utilities -using CalibrateEmulateSample.EnsembleKalmanProcesses -using CalibrateEmulateSample.ParameterDistributions -using CalibrateEmulateSample.DataContainers -using CalibrateEmulateSample.Observations - -#rng_seed = 4137 -#rng_seed = 413798 -rng_seed = 2413798 -Random.seed!(rng_seed) - - -function get_standardizing_factors(data::Array{FT, 2}) where {FT} - # Input: data size: N_data x N_ensembles - # Ensemble median of the data - norm_factor = median(data, dims = 2) - return norm_factor -end - -function get_standardizing_factors(data::Array{FT, 1}) where {FT} - # Input: data size: N_data*N_ensembles (splatted) - # Ensemble median of the data - norm_factor = median(data) - return norm_factor -end - - -# Output figure save directory -example_directory = @__DIR__ -println(example_directory) -figure_save_directory = joinpath(example_directory, "output") -data_save_directory = joinpath(example_directory, "output") -if !isdir(figure_save_directory) - mkdir(figure_save_directory) -end -if !isdir(data_save_directory) - mkdir(data_save_directory) -end - -# Governing settings -# Characteristic time scale -τc = 5.0 # days, prescribed by the L96 problem -# Stationary or transient dynamics -dynamics = 2 # Transient is 2 -# Statistics integration length -# This has to be less than 360 and 360 must be divisible by Ts_days -Ts_days = 30.0 # Integration length in days -# Stats type, which statistics to construct from the L96 system -# 4 is a linear fit over a batch of length Ts_days -# 5 is the mean over a batch of length Ts_days -stats_type = 5 - - -### -### Define the (true) parameters -### -# Define the parameters that we want to learn -F_true = 8.0 # Mean F -A_true = 2.5 # Transient F amplitude -ω_true = 2.0 * π / (360.0 / τc) # Frequency of the transient F - -if dynamics == 2 - params_true = [F_true, A_true] - param_names = ["F", "A"] -else - params_true = [F_true] - param_names = ["F"] -end -n_param = length(param_names) -params_true = reshape(params_true, (n_param, 1)) - -println(n_param) -println(params_true) - - -### -### Define the parameter priors -### -# Lognormal prior or normal prior? -log_normal = false # THIS ISN't CURRENTLY IMPLEMENTED - -function logmean_and_logstd(μ, σ) - σ_log = sqrt(log(1.0 + σ^2 / μ^2)) - μ_log = log(μ / (sqrt(1.0 + σ^2 / μ^2))) - return μ_log, σ_log -end -#logmean_F, logstd_F = logmean_and_logstd(F_true, 5) -#logmean_A, logstd_A = logmean_and_logstd(A_true, 0.2*A_true) - -if dynamics == 2 - #prior_means = [F_true+0.5, A_true+0.5] - prior_means = [F_true, A_true] - prior_stds = [3.0, 0.5 * A_true] - prior_F = Dict( - "distribution" => Parameterized(Normal(prior_means[1], prior_stds[1])), - "constraint" => no_constraint(), - "name" => param_names[1], - ) - prior_A = Dict( - "distribution" => Parameterized(Normal(prior_means[2], prior_stds[2])), - "constraint" => no_constraint(), - "name" => param_names[2], - ) - - priors = ParameterDistribution([prior_F, prior_A]) - -else - prior_F = Dict("distribution" => Parameterized(Normal(F_true, 1)), "constraint" => no_constraint(), "name" => "F") - priors = ParameterDistribution(prior_F) -end - - -### -### Define the data from which we want to learn the parameters -### -data_names = ["y0", "y1"] - - -### -### L96 model settings -### - -# Lorenz 96 model parameters -# Behavior changes depending on the size of F, N -N = 36 -dt = 1 / 64.0 -# Start of integration -t_start = 800.0 -# Data collection length -#if dynamics==1 -# T = 2. -#else -# T = 360. / τc -#end -T = 360.0 / τc -# Batch length -Ts = 5.0 / τc # Nondimensionalize by L96 timescale -# Integration length -Tfit = Ts_days / τc -# Initial perturbation -Fp = rand(Normal(0.0, 0.01), N); -kmax = 1 -# Prescribe variance or use a number of forward passes to define true interval variability -var_prescribe = false -# Use CES to learn ω? -ω_fixed = true - -# Settings -# Constructs an LSettings structure, see GModel.jl for the descriptions -lorenz_settings = - GModel.LSettings(dynamics, stats_type, t_start, T, Ts, Tfit, Fp, N, dt, t_start + T, kmax, ω_fixed, ω_true); -lorenz_params = GModel.LParams(F_true, ω_true, A_true) - -### -### Generate (artificial) truth samples -### Note: The observables y are related to the parameters θ by: -### y = G(θ) + η -### - -# Lorenz forward -# Input: params: [N_params, N_ens] -# Output: gt: [N_data, N_ens] -# Dropdims of the output since the forward model is only being run with N_ens=1 -# corresponding to the truth construction -gt = dropdims(GModel.run_G_ensemble(params_true, lorenz_settings), dims = 2) - -# Compute internal variability covariance -n_samples = 50 -if var_prescribe == true - n_samples = 100 - yt = zeros(length(gt), n_samples) - noise_level = 0.05 - Γy = noise_level * convert(Array, Diagonal(gt)) - μ = zeros(length(gt)) - # Add noise - for i in 1:n_samples - yt[:, i] = gt .+ rand(MvNormal(μ, Γy)) - end -else - println("Using truth values to compute covariance") - yt = zeros(length(gt), n_samples) - for i in 1:n_samples - lorenz_settings_local = GModel.LSettings( - dynamics, - stats_type, - t_start + T * (i - 1), - T, - Ts, - Tfit, - Fp, - N, - dt, - t_start + T * (i - 1) + T, - kmax, - ω_fixed, - ω_true, - ) - yt[:, i] = GModel.run_G_ensemble(params_true, lorenz_settings_local) - end - # Variance of truth data - #Γy = convert(Array, Diagonal(dropdims(mean((yt.-mean(yt,dims=1)).^2,dims=1),dims=1))) - # Covariance of truth data - Γy = cov(yt, dims = 2) - - # println(Γy) - println(size(Γy), " ", rank(Γy)) -end - - -# Construct observation object -truth = Observations.Observation(yt, Γy, data_names) -# Truth sample for EKP -truth_sample = truth.mean -#sample_ind=randperm!(collect(1:n_samples))[1] -#truth_sample = yt[:,sample_ind] -#println("Truth sample:") -#println(sample_ind) -### -### Calibrate: Ensemble Kalman Inversion -### - -# L96 settings for the forward model in the EKP -# Here, the forward model for the EKP settings can be set distinctly from the truth runs -lorenz_settings_G = lorenz_settings; # initialize to truth settings - -# EKP parameters -log_transform(a::AbstractArray) = log.(a) -exp_transform(a::AbstractArray) = exp.(a) - -N_ens = 50 # number of ensemble members -N_iter = 5 # number of EKI iterations -# initial parameters: N_params x N_ens -initial_params = construct_initial_ensemble(priors, N_ens; rng_seed = rng_seed) - -ekiobj = EnsembleKalmanProcesses.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, Inversion()) - -# EKI iterations -println("EKP inversion error:") -err = zeros(N_iter) -err_params = zeros(N_iter) -for i in 1:N_iter - if log_normal == false - params_i = get_u_final(ekiobj) - else - params_i = exp_transform(get_u_final(ekiobj)) - end - g_ens = GModel.run_G_ensemble(params_i, lorenz_settings_G) - EnsembleKalmanProcesses.update_ensemble!(ekiobj, g_ens) - err[i] = get_error(ekiobj)[end] - err_params[i] = mean((params_true - mean(params_i, dims = 2)) .^ 2) - println("Iteration: " * string(i) * ", Error (data): " * string(err[i])) - println("Iteration: " * string(i) * ", Error (params): " * string(err_params[i])) -end - -# EKI results: Has the ensemble collapsed toward the truth? -println("True parameters: ") -println(params_true) - -println("\nEKI results:") -if log_normal == false - println(mean(get_u_final(ekiobj), dims = 2)) -else - println(mean(exp_transform(get_u_final(ekiobj)), dims = 2)) -end - -### -### Emulate: Gaussian Process Regression -### - -# Emulate-sample settings -standardize = true -retained_svd_frac = 0.95 - -gppackage = Emulators.GPJL() -pred_type = Emulators.YType() -gauss_proc = GaussianProcess( - gppackage; - kernel = nothing, # use default squared exponential kernel - prediction_type = pred_type, - noise_learn = false, -) - -# Standardize the output data -# Use median over all data since all data are the same type -yt_norm = vcat(yt...) -norm_factor = get_standardizing_factors(yt_norm) -println(size(norm_factor)) -#norm_factor = vcat(norm_factor...) -norm_factor = fill(norm_factor, size(truth_sample)) -println("Standardization factors") -println(norm_factor) - -# Get training points from the EKP iteration number in the second input term -N_iter = 5; -input_output_pairs = Utilities.get_training_points(ekiobj, N_iter) -normalized = true -emulator = Emulator( - gauss_proc, - input_output_pairs; - obs_noise_cov = Γy, - normalize_inputs = normalized, - standardize_outputs = standardize, - standardize_outputs_factors = norm_factor, - retained_svd_frac = retained_svd_frac, -) -optimize_hyperparameters!(emulator) - -# Check how well the Gaussian Process regression predicts on the -# true parameters -#if retained_svd_frac==1.0 -if log_normal == false - y_mean, y_var = Emulators.predict(emulator, reshape(params_true, :, 1), transform_to_real = true) -else - y_mean, y_var = Emulators.predict(emulator, reshape(log.(params_true), :, 1), transform_to_real = true) -end - -println("GP prediction on true parameters: ") -println(vec(y_mean)) -println(" GP variance") -println(diag(y_var[1], 0)) -println("true data: ") -println(truth.mean) # same, regardless of norm_factor -println("GP MSE: ") -println(mean((truth.mean - vec(y_mean)) .^ 2)) -#end -### -### Sample: Markov Chain Monte Carlo -### - -# initial values -u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) -println("initial parameters: ", u0) - -# First let's run a short chain to determine a good step size -yt_sample = truth_sample -mcmc = MCMCWrapper(RWMHSampling(), yt_sample, priors, emulator; init_params = u0) -new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0) - -# Now begin the actual MCMC -println("Begin MCMC - with step size ", new_step) -chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) -posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) - -post_mean = mean(posterior) -post_cov = cov(posterior) -println("post_mean") -println(post_mean) -println("post_cov") -println(post_cov) -println("D util") -println(det(inv(post_cov))) -println(" ") - -# Plot the posteriors together with the priors and the true parameter values -if log_normal == false - true_values = params_true -else - true_values = log.(params_true) -end -n_params = length(true_values) -y_folder = "qoi_" * string(lorenz_settings.stats_type) * "_" -save_directory = figure_save_directory * y_folder -for idx in 1:n_params - if idx == 1 - param = "F" - xbounds = [F_true - 1.0, F_true + 1.0] - xs = collect(xbounds[1]:((xbounds[2] - xbounds[1]) / 100):xbounds[2]) - elseif idx == 2 - param = "A" - xbounds = [A_true - 1.0, A_true + 1.0] - xs = collect(xbounds[1]:((xbounds[2] - xbounds[1]) / 100):xbounds[2]) - elseif idx == 3 - param = "ω" - xs = collect((ω_true - 0.2):0.01:(ω_true + 0.2)) - xbounds = [xs[1], xs[end]] - else - throw("not implemented") - end - - if log_normal == true - xs = log.(xs) - xbounds = log.(xbounds) - end - - label = "true " * param - posterior_distributions = get_distribution(posterior) - - histogram( - posterior_distributions[param_names[idx]][1, :], - bins = 100, - normed = true, - fill = :slategray, - lab = "posterior", - ) - prior_plot = get_distribution(mcmc.prior) - # This requires StatsPlots - plot!(xs, prior_plot[param_names[idx]], w = 2.6, color = :blue, lab = "prior") - #plot!(xs, mcmc.prior[idx].dist, w=2.6, color=:blue, lab="prior") - plot!([true_values[idx]], seriestype = "vline", w = 2.6, lab = label) - plot!(xlims = xbounds) - - title!(param) - - figpath = joinpath(figure_save_directory, "posterior_$(param)_T_$(T)_w_$(ω_true).png") - savefig(figpath) - linkfig(figpath) -end - -# Save data -@save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs diff --git a/examples/Lorenz/calibrate.jl b/examples/Lorenz/calibrate.jl new file mode 100644 index 000000000..79038276a --- /dev/null +++ b/examples/Lorenz/calibrate.jl @@ -0,0 +1,239 @@ +### This script is copied from Lorenz 96 example in EnsembleKalmanProcesses.jl v0.14.2 + +include("GModel.jl") # Contains Lorenz 96 source code + +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using StatsPlots +using Plots +using Random +using JLD2 + +# CES +using CalibrateEmulateSample +const EKP = CalibrateEmulateSample.EnsembleKalmanProcesses +const PD = EKP.ParameterDistributions + +function main() + + rng_seed = 4137 + Random.seed!(rng_seed) + + # Output figure save directory + homedir = pwd() + println(homedir) + figure_save_directory = homedir * "/output/" + data_save_directory = homedir * "/output/" + if !isdir(figure_save_directory) + mkdir(figure_save_directory) + end + if !isdir(data_save_directory) + mkdir(data_save_directory) + end + + # Governing settings + # Characteristic time scale + τc = 5.0 # days, prescribed by the L96 problem + # Stationary or transient dynamics + dynamics = 2 # Transient is 2 + # Statistics integration length + # This has to be less than 360 and 360 must be divisible by Ts_days + Ts_days = 90.0 # Integration length in days + # Stats type, which statistics to construct from the L96 system + # 4 is a linear fit over a batch of length Ts_days + # 5 is the mean over a batch of length Ts_days + stats_type = 5 + + + ### + ### Define the (true) parameters + ### + # Define the parameters that we want to learn + F_true = 8.0 # Mean F + A_true = 2.5 # Transient F amplitude + ω_true = 2.0 * π / (360.0 / τc) # Frequency of the transient F + + if dynamics == 2 + params_true = [F_true, A_true] + param_names = ["F", "A"] + else + params_true = [F_true] + param_names = ["F"] + end + n_param = length(param_names) + params_true = reshape(params_true, (n_param, 1)) + + println(n_param) + println(params_true) + + + ### + ### Define the parameter priors + ### + + if dynamics == 2 + prior_means = [F_true + 1.0, A_true + 0.5] + prior_stds = [2.0, 0.5 * A_true] + prior_F = PD.constrained_gaussian(param_names[1], prior_means[1], prior_stds[1], 0, Inf) + prior_A = PD.constrained_gaussian(param_names[2], prior_means[2], prior_stds[2], 0, Inf) + priors = PD.combine_distributions([prior_F, prior_A]) + else + priors = PD.constrained_gaussian("F", F_true, 1.0, 0, Inf) + end + + + ### + ### Define the data from which we want to learn the parameters + ### + data_names = ["y0", "y1"] + + + ### + ### L96 model settings + ### + + # Lorenz 96 model parameters + # Behavior changes depending on the size of F, N + N = 36 + dt = 1 / 64.0 + # Start of integration + t_start = 800.0 + # Data collection length + if dynamics == 1 + T = 2.0 + else + T = 360.0 / τc + end + # Batch length + Ts = 5.0 / τc # Nondimensionalize by L96 timescale + # Integration length + Tfit = Ts_days / τc + # Initial perturbation + Fp = rand(Normal(0.0, 0.01), N) + kmax = 1 + # Prescribe variance or use a number of forward passes to define true interval variability + var_prescribe = false + # Use CES to learn ω? + ω_fixed = true + + # Settings + # Constructs an LSettings structure, see GModel.jl for the descriptions + lorenz_settings = + GModel.LSettings(dynamics, stats_type, t_start, T, Ts, Tfit, Fp, N, dt, t_start + T, kmax, ω_fixed, ω_true) + lorenz_params = GModel.LParams(F_true, ω_true, A_true) + + ### + ### Generate (artificial) truth samples + ### Note: The observables y are related to the parameters θ by: + ### y = G(θ) + η + ### + + # Lorenz forward + # Input: params: [N_params, N_ens] + # Output: gt: [N_data, N_ens] + # Dropdims of the output since the forward model is only being run with N_ens=1 + # corresponding to the truth construction + gt = dropdims(GModel.run_G_ensemble(params_true, lorenz_settings), dims = 2) + + # Compute internal variability covariance + if var_prescribe == true + n_samples = 100 + yt = zeros(length(gt), n_samples) + noise_level = 0.05 + Γy = noise_level * convert(Array, Diagonal(gt)) + μ = zeros(length(gt)) + # Add noise + for i in 1:n_samples + yt[:, i] = gt .+ rand(MvNormal(μ, Γy)) + end + else + println("Using truth values to compute covariance") + n_samples = 20 + yt = zeros(length(gt), n_samples) + for i in 1:n_samples + lorenz_settings_local = GModel.LSettings( + dynamics, + stats_type, + t_start + T * (i - 1), + T, + Ts, + Tfit, + Fp, + N, + dt, + t_start + T * (i - 1) + T, + kmax, + ω_fixed, + ω_true, + ) + yt[:, i] = GModel.run_G_ensemble(params_true, lorenz_settings_local) + end + # Variance of truth data + #Γy = convert(Array, Diagonal(dropdims(mean((yt.-mean(yt,dims=1)).^2,dims=1),dims=1))) + # Covariance of truth data + Γy = cov(yt, dims = 2) + + println(Γy) + end + + + # Construct observation object + truth = Observations.Observation(yt, Γy, data_names) + truth_sample = truth.mean + ### + ### Calibrate: Ensemble Kalman Inversion + ### + + # L96 settings for the forward model in the EKP. Here, the forward model for the EKP settings can be set distinctly from the truth runs + lorenz_settings_G = lorenz_settings # initialize to truth settings + + # EKP parameters + N_ens = 20 # number of ensemble members + N_iter = 5 # number of EKI iterations + # initial parameters: N_params x N_ens + initial_params = EKP.construct_initial_ensemble(priors, N_ens; rng_seed = rng_seed) + + ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth_sample, truth.obs_noise_cov, EKP.Inversion()) + + # EKI iterations + println("EKP inversion error:") + err = zeros(N_iter) + for i in 1:N_iter + params_i = EKP.get_ϕ_final(priors, ekiobj) # the `ϕ` indicates that the `params_i` are in the constrained space + g_ens = GModel.run_G_ensemble(params_i, lorenz_settings_G) + EKP.update_ensemble!(ekiobj, g_ens) + err[i] = EKP.get_error(ekiobj)[end] #mean((params_true - mean(params_i,dims=2)).^2) + println("Iteration: " * string(i) * ", Error: " * string(err[i])) + end + + # EKI results: Has the ensemble collapsed toward the truth? + println("True parameters: ") + println(params_true) + + println("\nEKI results:") + println(EKP.get_ϕ_mean_final(priors, ekiobj)) + + u_stored = EKP.get_u(ekiobj, return_array = false) + g_stored = EKP.get_g(ekiobj, return_array = false) + + save( + joinpath(data_save_directory, "calibrate_results.jld2"), + "inputs", + u_stored, + "outputs", + g_stored, + "priors", + priors, + "eki", + ekiobj, + "truth_sample", + truth_sample, + "truth_sample_mean", + truth.mean, + "truth_input_constrained", + params_true, #constrained here, as these are in a physically constrained space (unlike the u inputs), + ) +end + +main() diff --git a/examples/Lorenz/emulate_sample.jl b/examples/Lorenz/emulate_sample.jl new file mode 100644 index 000000000..46e2286ba --- /dev/null +++ b/examples/Lorenz/emulate_sample.jl @@ -0,0 +1,214 @@ +# Import modules +include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) +include(joinpath(@__DIR__, "GModel.jl")) # Contains Lorenz 96 source code + +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +using StatsPlots +using Plots +using Random +using JLD2 + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers +using CalibrateEmulateSample.Observations + +function get_standardizing_factors(data::Array{FT, 2}) where {FT} + # Input: data size: N_data x N_ensembles + # Ensemble median of the data + norm_factor = median(data, dims = 2) + return norm_factor +end + +function get_standardizing_factors(data::Array{FT, 1}) where {FT} + # Input: data size: N_data*N_ensembles (splatted) + # Ensemble median of the data + norm_factor = median(data) + return norm_factor +end + + +function main() + ##### + # Should be loaded: + τc = 5.0 + F_true = 8.0 # Mean F + A_true = 2.5 # Transient F amplitude + ω_true = 2.0 * π / (360.0 / τc) # Frequency of the transient F + #### + + exp_name = "Lorenz_histogram_F$(F_true)_A$(A_true)-w$(ω_true)" + rng_seed = 44009 + Random.seed!(rng_seed) + # loading relevant data + homedir = pwd() + println(homedir) + figure_save_directory = joinpath(homedir, "output/") + data_save_directory = joinpath(homedir, "output/") + data_save_file = joinpath(data_save_directory, "calibrate_results.jld2") + + if !isfile(data_save_file) + throw( + ErrorException( + "data file $data_save_file not found. \n First run: \n > julia --project calibrate.jl \n and store results $data_save_file", + ), + ) + end + + ekiobj = load(data_save_file)["eki"] + priors = load(data_save_file)["priors"] + truth_sample = load(data_save_file)["truth_sample"] + truth_sample_mean = load(data_save_file)["truth_sample_mean"] + truth_params_constrained = load(data_save_file)["truth_input_constrained"] #true parameters in constrained space + truth_params = transform_constrained_to_unconstrained(priors, truth_params_constrained) + Γy = ekiobj.obs_noise_cov + ### + ### Emulate: Gaussian Process Regression + ### + + # Emulate-sample settings + standardize = true + retained_svd_frac = 0.95 + + gppackage = Emulators.GPJL() + pred_type = Emulators.YType() + gauss_proc = GaussianProcess( + gppackage; + kernel = nothing, # use default squared exponential kernel + prediction_type = pred_type, + noise_learn = false, + ) + + # Standardize the output data + # Use median over all data since all data are the same type + truth_sample_norm = vcat(truth_sample...) + norm_factor = get_standardizing_factors(truth_sample_norm) + println(size(norm_factor)) + #norm_factor = vcat(norm_factor...) + norm_factor = fill(norm_factor, size(truth_sample)) + println("Standardization factors") + println(norm_factor) + + # Get training points from the EKP iteration number in the second input term + N_iter = 5 + input_output_pairs = Utilities.get_training_points(ekiobj, N_iter) + # Save data + @save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs + + normalized = true + emulator = Emulator( + gauss_proc, + input_output_pairs; + obs_noise_cov = Γy, + normalize_inputs = normalized, + standardize_outputs = standardize, + standardize_outputs_factors = norm_factor, + retained_svd_frac = retained_svd_frac, + ) + optimize_hyperparameters!(emulator) + + # Check how well the Gaussian Process regression predicts on the + # true parameters + #if retained_svd_frac==1.0 + y_mean, y_var = Emulators.predict(emulator, reshape(truth_params, :, 1), transform_to_real = true) + + println("GP prediction on true parameters: ") + println(vec(y_mean)) + println(" GP variance") + println(diag(y_var[1], 0)) + println("true data: ") + println(truth_sample_mean) # same, regardless of norm_factor + println("GP MSE: ") + println(mean((truth_sample_mean - vec(y_mean)) .^ 2)) + + #end + ### + ### Sample: Markov Chain Monte Carlo + ### + # initial values + u0 = vec(mean(get_inputs(input_output_pairs), dims = 2)) + println("initial parameters: ", u0) + + # First let's run a short chain to determine a good step size + truth_sample = truth_sample + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, priors, emulator; init_params = u0) + new_step = optimize_stepsize(mcmc; init_stepsize = 0.1, N = 2000, discard_initial = 0) + + # Now begin the actual MCMC + println("Begin MCMC - with step size ", new_step) + chain = MarkovChainMonteCarlo.sample(mcmc, 100_000; stepsize = new_step, discard_initial = 2_000) + posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain) + + post_mean = mean(posterior) + post_cov = cov(posterior) + println("post_mean") + println(post_mean) + println("post_cov") + println(post_cov) + println("D util") + println(det(inv(post_cov))) + println(" ") + + + # Plot the posteriors together with the priors and the true parameter values (in the constrained space) + n_params = length(truth_params) + save_directory = joinpath(figure_save_directory) + + posterior_distribution_unconstrained_dict = get_distribution(posterior) #as it is a Samples, the distribution are samples + posterior_samples_unconstrained_dict = get_distribution(posterior) #as it is a Samples, the distribution are samples + param_names = get_name(posterior) + posterior_samples_unconstrained_arr = vcat([posterior_samples_unconstrained_dict[n] for n in param_names]...) + posterior_samples = transform_unconstrained_to_constrained(priors, posterior_samples_unconstrained_arr) + + for idx in 1:n_params + if idx == 1 + param = "F" + xbounds = [F_true - 1.0, F_true + 1.0] + xs = collect(xbounds[1]:((xbounds[2] - xbounds[1]) / 100):xbounds[2]) + elseif idx == 2 + param = "A" + xbounds = [A_true - 1.0, A_true + 1.0] + xs = collect(xbounds[1]:((xbounds[2] - xbounds[1]) / 100):xbounds[2]) + elseif idx == 3 + param = "ω" + xs = collect((ω_true - 0.2):0.01:(ω_true + 0.2)) + xbounds = [xs[1], xs[end]] + else + throw("not implemented") + end + + label = "true " * param + + histogram(posterior_samples[idx, :], bins = 100, normed = true, fill = :slategray, lab = "posterior") + prior_plot = get_distribution(mcmc.prior) + + # This requires StatsPlots + xs_rows = permutedims(repeat(xs, 1, size(posterior_samples, 1)), (2, 1)) # This hack will enable us to apply the transformation using the full prior + xs_rows_unconstrained = transform_constrained_to_unconstrained(priors, xs_rows) + plot!( + xs, + pdf.(prior_plot[param_names[idx]], xs_rows_unconstrained[idx, :]), + w = 2.6, + color = :blue, + lab = "prior", + ) + #plot!(xs, mcmc.prior[idx].dist, w=2.6, color=:blue, lab="prior") + plot!([truth_params[idx]], seriestype = "vline", w = 2.6, lab = label) + plot!(xlims = xbounds) + + title!(param) + + figpath = joinpath(figure_save_directory, "posterior_$(param)_" * exp_name * ".png") + savefig(figpath) + linkfig(figpath) + end + +end + +main()