diff --git a/examples/Darcy/GModel.jl b/examples/Darcy/GModel.jl new file mode 100644 index 00000000..e18ecccf --- /dev/null +++ b/examples/Darcy/GModel.jl @@ -0,0 +1,219 @@ +################## +# Copied on 3/16/23 and modified from +# https://github.com/Zhengyu-Huang/InverseProblems.jl/blob/master/Fluid/Darcy-2D.jl +################## + +using JLD2 +using Statistics +using LinearAlgebra +using Distributions +using Random +using SparseArrays + + + + +mutable struct Setup_Param{FT <: AbstractFloat, IT <: Int} + # physics + N::IT # number of grid points for both x and y directions (including both ends) + Δx::FT + xx::AbstractVector{FT} # uniform grid [a, a+Δx, a+2Δx ... b] (in each dimension) + + #for source term + f_2d::AbstractMatrix{FT} + + κ::AbstractMatrix{FT} + + # observation locations is tensor product x_locs × y_locs + x_locs::AbstractVector{IT} + y_locs::AbstractVector{IT} + + N_y::IT +end + + +function Setup_Param( + xx::AbstractVector{FT}, + obs_ΔN::IT, + κ::AbstractMatrix; + seed::IT = 123, +) where {FT <: AbstractFloat, IT <: Int} + + N = length(xx) + Δx = xx[2] - xx[1] + + # logκ_2d, φ, λ, θ_ref = generate_θ_KL(xx, N_KL, d, τ, seed=seed) + f_2d = compute_f_2d(xx) + + x_locs = Array(obs_ΔN:obs_ΔN:(N - obs_ΔN)) + y_locs = Array(obs_ΔN:obs_ΔN:(N - obs_ΔN)) + N_y = length(x_locs) * length(y_locs) + + Setup_Param(N, Δx, xx, f_2d, κ, x_locs, y_locs, N_y) +end + + + +#= +A hardcoding source function, +which assumes the computational domain is +[0 1]×[0 1] +f(x,y) = f(y), +which dependes only on y +=# +function compute_f_2d(yy::AbstractVector{FT}) where {FT <: AbstractFloat} + N = length(yy) + f_2d = zeros(FT, N, N) + for i in 1:N + if (yy[i] <= 4 / 6) + f_2d[:, i] .= 1000.0 + elseif (yy[i] >= 4 / 6 && yy[i] <= 5 / 6) + f_2d[:, i] .= 2000.0 + elseif (yy[i] >= 5 / 6) + f_2d[:, i] .= 3000.0 + end + end + return f_2d +end + +""" + run_G_ensemble(darcy,κs::AbstractMatrix) + +Computes the forward map `G` (`solve_Darcy_2D` followed by `compute_obs`) over an ensemble of `κ`'s, stored flat as columns of `κs` +""" +function run_G_ensemble(darcy, κs::AbstractMatrix) + N_ens = size(κs, 2) # ens size + nd = darcy.N_y #num obs + g_ens = zeros(nd, N_ens) + for i in 1:N_ens + # run the model with the current parameters, i.e., map θ to G(θ) + κ_i = reshape(κs[:, i], darcy.N, darcy.N) # unflatten + h_i = solve_Darcy_2D(darcy, κ_i) # run model + g_ens[:, i] = compute_obs(darcy, h_i) # observe solution + end + return g_ens +end + + + +#= + return the unknow index for the grid point + Since zero-Dirichlet boundary conditions are imposed on + all four edges, the freedoms are only on interior points +=# +function ind(darcy::Setup_Param{FT, IT}, ix::IT, iy::IT) where {FT <: AbstractFloat, IT <: Int} + return (ix - 1) + (iy - 2) * (darcy.N - 2) +end + +function ind_all(darcy::Setup_Param{FT, IT}, ix::IT, iy::IT) where {FT <: AbstractFloat, IT <: Int} + return ix + (iy - 1) * darcy.N +end + +#= + solve Darcy equation with finite difference method: + -∇(κ∇h) = f + with Dirichlet boundary condition, h=0 on ∂Ω +=# +function solve_Darcy_2D(darcy::Setup_Param{FT, IT}, κ_2d::AbstractMatrix{FT}) where {FT <: AbstractFloat, IT <: Int} + Δx, N = darcy.Δx, darcy.N + + indx = IT[] + indy = IT[] + vals = FT[] + + f_2d = darcy.f_2d + + 𝓒 = Δx^2 + for iy in 2:(N - 1) + for ix in 2:(N - 1) + + ixy = ind(darcy, ix, iy) + + #top + if iy == N - 1 + #ft = -(κ_2d[ix, iy] + κ_2d[ix, iy+1])/2.0 * (0 - h_2d[ix,iy]) + push!(indx, ixy) + push!(indy, ixy) + push!(vals, (κ_2d[ix, iy] + κ_2d[ix, iy + 1]) / 2.0 / 𝓒) + else + #ft = -(κ_2d[ix, iy] + κ_2d[ix, iy+1])/2.0 * (h_2d[ix,iy+1] - h_2d[ix,iy]) + append!(indx, [ixy, ixy]) + append!(indy, [ixy, ind(darcy, ix, iy + 1)]) + append!( + vals, + [(κ_2d[ix, iy] + κ_2d[ix, iy + 1]) / 2.0 / 𝓒, -(κ_2d[ix, iy] + κ_2d[ix, iy + 1]) / 2.0 / 𝓒], + ) + end + + #bottom + if iy == 2 + #fb = (κ_2d[ix, iy] + κ_2d[ix, iy-1])/2.0 * (h_2d[ix,iy] - 0) + push!(indx, ixy) + push!(indy, ixy) + push!(vals, (κ_2d[ix, iy] + κ_2d[ix, iy - 1]) / 2.0 / 𝓒) + else + #fb = (κ_2d[ix, iy] + κ_2d[ix, iy-1])/2.0 * (h_2d[ix,iy] - h_2d[ix,iy-1]) + append!(indx, [ixy, ixy]) + append!(indy, [ixy, ind(darcy, ix, iy - 1)]) + append!( + vals, + [(κ_2d[ix, iy] + κ_2d[ix, iy - 1]) / 2.0 / 𝓒, -(κ_2d[ix, iy] + κ_2d[ix, iy - 1]) / 2.0 / 𝓒], + ) + end + + #right + if ix == N - 1 + #fr = -(κ_2d[ix, iy] + κ_2d[ix+1, iy])/2.0 * (0 - h_2d[ix,iy]) + push!(indx, ixy) + push!(indy, ixy) + push!(vals, (κ_2d[ix, iy] + κ_2d[ix + 1, iy]) / 2.0 / 𝓒) + else + #fr = -(κ_2d[ix, iy] + κ_2d[ix+1, iy])/2.0 * (h_2d[ix+1,iy] - h_2d[ix,iy]) + append!(indx, [ixy, ixy]) + append!(indy, [ixy, ind(darcy, ix + 1, iy)]) + append!( + vals, + [(κ_2d[ix, iy] + κ_2d[ix + 1, iy]) / 2.0 / 𝓒, -(κ_2d[ix, iy] + κ_2d[ix + 1, iy]) / 2.0 / 𝓒], + ) + end + + #left + if ix == 2 + #fl = (κ_2d[ix, iy] + κ_2d[ix-1, iy])/2.0 * (h_2d[ix,iy] - 0) + push!(indx, ixy) + push!(indy, ixy) + push!(vals, (κ_2d[ix, iy] + κ_2d[ix - 1, iy]) / 2.0 / 𝓒) + else + #fl = (κ_2d[ix, iy] + κ_2d[ix-1, iy])/2.0 * (h_2d[ix,iy] - h_2d[ix-1,iy]) + append!(indx, [ixy, ixy]) + append!(indy, [ixy, ind(darcy, ix - 1, iy)]) + append!( + vals, + [(κ_2d[ix, iy] + κ_2d[ix - 1, iy]) / 2.0 / 𝓒, -(κ_2d[ix, iy] + κ_2d[ix - 1, iy]) / 2.0 / 𝓒], + ) + end + + end + end + + + + df = sparse(indx, indy, vals, (N - 2)^2, (N - 2)^2) + # Multithread does not support sparse matrix solver + h = df \ (f_2d[2:(N - 1), 2:(N - 1)])[:] + + h_2d = zeros(FT, N, N) + h_2d[2:(N - 1), 2:(N - 1)] .= reshape(h, N - 2, N - 2) + + return h_2d +end + +#= +Compute observation values +=# +function compute_obs(darcy::Setup_Param{FT, IT}, h_2d::AbstractMatrix{FT}) where {FT <: AbstractFloat, IT <: Int} + # X---X(1)---X(2) ... X(obs_N)---X + obs_2d = h_2d[darcy.x_locs, darcy.y_locs] + + return obs_2d[:] +end diff --git a/examples/Darcy/Project.toml b/examples/Darcy/Project.toml new file mode 100644 index 00000000..ae3882ed --- /dev/null +++ b/examples/Darcy/Project.toml @@ -0,0 +1,10 @@ +[deps] +CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +GaussianRandomFields = "e4b2fa32-6e09-5554-b718-106ed5adafe9" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/examples/Darcy/calibrate.jl b/examples/Darcy/calibrate.jl new file mode 100644 index 00000000..3043239d --- /dev/null +++ b/examples/Darcy/calibrate.jl @@ -0,0 +1,208 @@ +# # [Learning the Pearmibility field in a Darcy flow from noisy sparse observations] + +# In this example we hope to illustrate function learning. One may wish to use function learning in cases where the underlying parameter of interest is actual a finite-dimensional approximation (e.g. spatial discretization) of some "true" function. Treating such an object directly will lead to increasingly high-dimensional learning problems as the spatial resolution is increased, resulting in poor computational scaling and increasingly ill-posed inverse problems. Treating the object as a discretized function from a function space, one can learn coefficients not in the standard basis, but instead in a basis of this function space, it is commonly the case that functions will have relatively low effective dimension, and will be depend only on the spatial discretization due to discretization error, that should vanish as resolution is increased. + +# We will solve for an unknown permeability field ``\kappa`` governing the pressure field of a Darcy flow on a square 2D domain. To learn about the permeability we shall take few pointwise measurements of the solved pressure field within the domain. The forward solver is a simple finite difference scheme taken and modified from code [here](https://github.com/Zhengyu-Huang/InverseProblems.jl/blob/master/Fluid/Darcy-2D.jl). + +# First we load standard packages +using LinearAlgebra +using Distributions +using Random +using JLD2 + +# the package to define the function distributions +import GaussianRandomFields # we wrap this so we don't want to use "using" +const GRF = GaussianRandomFields + +# and finally the EKP packages +using CalibrateEmulateSample +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.ParameterDistributions +const EKP = CalibrateEmulateSample.EnsembleKalmanProcesses + +# We include the forward solver here +include("GModel.jl") + +# Then link some outputs for figures and plotting +fig_save_directory = joinpath(@__DIR__, "output") +data_save_directory = joinpath(@__DIR__, "output") +if !isdir(fig_save_directory) + mkdir(fig_save_directory) +end +if !isdir(data_save_directory) + mkdir(data_save_directory) +end# TOML interface for fitting parameters of a sinusoid + +PLOT_FLAG = true +if PLOT_FLAG + using Plots + @info "Plotting enabled, this will reduce code performance. Figures stored in $fig_save_directory" +end + +# Set a random seed. +seed = 100234 +rng = Random.MersenneTwister(seed) + + +function main() +# Define the spatial domain and discretization +dim = 2 +N, L = 80, 1.0 +pts_per_dim = LinRange(0, L, N) +obs_ΔN = 10 + +# To provide a simple test case, we assume that the true function parameter is a particular sample from the function space we set up to define our prior. More precisely we choose a value of the truth that doesnt have a vanishingly small probability under the prior defined by a probability distribution over functions; here taken as a family of Gaussian Random Fields (GRF). The function distribution is characterized by a covariance function - here a Matern kernel which assumes a level of smoothness over the samples from the distribution. We define an appropriate expansion of this distribution, here based on the Karhunen-Loeve expansion (similar to an eigenvalue-eigenfunction expansion) that is truncated to a finite number of terms, known as the degrees of freedom (`dofs`). The `dofs` define the effective dimension of the learning problem, decoupled from the spatial discretization. Explicitly, larger `dofs` may be required to represent multiscale functions, but come at an increased dimension of the parameter space and therefore a typical increase in cost and difficulty of the learning problem. + +smoothness = 1.0 +corr_length = 0.25 +dofs = 5 + +grf = GRF.GaussianRandomField( + GRF.CovarianceFunction(dim, GRF.Matern(smoothness, corr_length)), + GRF.KarhunenLoeve(dofs), + pts_per_dim, + pts_per_dim, +) + +# We define a wrapper around the GRF, and as the permeability field must be positive we introduce a domain constraint into the function distribution. Henceforth, the GRF is interfaced in the same manner as any other parameter distribution with regards to interface. +pkg = GRFJL() +distribution = GaussianRandomFieldInterface(grf, pkg) # our wrapper from EKP +domain_constraint = bounded_below(0) # make κ positive +pd = ParameterDistribution(Dict("distribution" => distribution, "name" => "kappa", "constraint" => domain_constraint)) # the fully constrained parameter distribution + +# Now we have a function distribution, we sample a reasonably high-probability value from this distribution as a true value (here all degrees of freedom set with `u_{\mathrm{true}} = -0.5`). We use the EKP transform function to build the corresponding instance of the ``\kappa_{\mathrm{true}}``. +u_true = -1.5 * ones(dofs, 1) # the truth parameter +println("True coefficients: ") +println(u_true) +κ_true = transform_unconstrained_to_constrained(pd, u_true) # builds and constrains the function. +κ_true = reshape(κ_true, N, N) + +# Now we generate the data sample for the truth in a perfect model setting by evaluating the the model here, and observing it by subsampling in each dimension every `obs_ΔN` points, and add some observational noise +darcy = Setup_Param(pts_per_dim, obs_ΔN, κ_true) +println(" Number of observation points: $(darcy.N_y)") +h_2d_true = solve_Darcy_2D(darcy, κ_true) +y_noiseless = compute_obs(darcy, h_2d_true) +obs_noise_cov = 0.05^2 * I(length(y_noiseless)) * (maximum(y_noiseless) - minimum(y_noiseless)) +truth_sample = vec(y_noiseless + rand(rng, MvNormal(zeros(length(y_noiseless)), obs_noise_cov))) + + +# Now we set up the Bayesian inversion algorithm. The prior we have already defined to construct our truth +prior = pd + + +# We define some algorithm parameters, here we take ensemble members larger than the dimension of the parameter space +N_ens = 50 # number of ensemble members +N_iter = 5 # number of EKI iterations + +# We sample the initial ensemble from the prior, and create the EKP object as an EKI algorithm using the `Inversion()` keyword +initial_params = construct_initial_ensemble(rng, prior, N_ens) +ekiobj = EKP.EnsembleKalmanProcess(initial_params, truth_sample, obs_noise_cov, Inversion()) + +# We perform the inversion loop. Remember that within calls to `get_ϕ_final` the EKP transformations are applied, thus the ensemble that is returned will be the positively-bounded permeability field evaluated at all the discretization points. +println("Begin inversion") +err = [] +final_it = [N_iter] +for i in 1:N_iter + params_i = get_ϕ_final(prior, ekiobj) + g_ens = run_G_ensemble(darcy, params_i) + terminate = EKP.update_ensemble!(ekiobj, g_ens) + push!(err, get_error(ekiobj)[end]) #mean((params_true - mean(params_i,dims=2)).^2) + println("Iteration: " * string(i) * ", Error: " * string(err[i])) + if !isnothing(terminate) + final_it[1] = i - 1 + break + end +end +n_iter = final_it[1] +# We plot first the prior ensemble mean and pointwise variance of the permeability field, and also the pressure field solved with the ensemble mean. Each ensemble member is stored as a column and therefore for uses such as plotting one needs to reshape to the desired dimension. +if PLOT_FLAG + gr(size = (1500, 400), legend = false) + prior_κ_ens = get_ϕ(prior, ekiobj, 1) + κ_ens_mean = reshape(mean(prior_κ_ens, dims = 2), N, N) + p1 = contour(pts_per_dim, pts_per_dim, κ_ens_mean', fill = true, levels = 15, title = "kappa mean", colorbar = true) + κ_ens_ptw_var = reshape(var(prior_κ_ens, dims = 2), N, N) + p2 = contour( + pts_per_dim, + pts_per_dim, + κ_ens_ptw_var', + fill = true, + levels = 15, + title = "kappa var", + colorbar = true, + ) + h_2d = solve_Darcy_2D(darcy, κ_ens_mean) + p3 = contour(pts_per_dim, pts_per_dim, h_2d', fill = true, levels = 15, title = "pressure", colorbar = true) + l = @layout [a b c] + plt = plot(p1, p2, p3, layout = l) + savefig(plt, joinpath(fig_save_directory, "output_prior.png")) # pre update + +end + +# Now we plot the final ensemble mean and pointwise variance of the permeability field, and also the pressure field solved with the ensemble mean. +if PLOT_FLAG + gr(size = (1500, 400), legend = false) + final_κ_ens = get_ϕ_final(prior, ekiobj) # the `ϕ` indicates that the `params_i` are in the constrained space + κ_ens_mean = reshape(mean(final_κ_ens, dims = 2), N, N) + p1 = contour(pts_per_dim, pts_per_dim, κ_ens_mean', fill = true, levels = 15, title = "kappa mean", colorbar = true) + κ_ens_ptw_var = reshape(var(final_κ_ens, dims = 2), N, N) + p2 = contour( + pts_per_dim, + pts_per_dim, + κ_ens_ptw_var', + fill = true, + levels = 15, + title = "kappa var", + colorbar = true, + ) + h_2d = solve_Darcy_2D(darcy, κ_ens_mean) + p3 = contour(pts_per_dim, pts_per_dim, h_2d', fill = true, levels = 15, title = "pressure", colorbar = true) + l = @layout [a b c] + plt = plot(p1, p2, p3; layout = l) + savefig(plt, joinpath(fig_save_directory, "output_it_" * string(n_iter) * ".png")) # pre update + +end +println("Final coefficients (ensemble mean):") +println(get_u_mean_final(ekiobj)) + +# We can compare this with the true permeability and pressure field: +if PLOT_FLAG + gr(size = (1000, 400), legend = false) + p1 = contour(pts_per_dim, pts_per_dim, κ_true', fill = true, levels = 15, title = "kappa true", colorbar = true) + p2 = contour( + pts_per_dim, + pts_per_dim, + h_2d_true', + fill = true, + levels = 15, + title = "pressure true", + colorbar = true, + ) + l = @layout [a b] + plt = plot(p1, p2, layout = l) + savefig(plt, joinpath(fig_save_directory, "output_true.png")) +end + +# Finally the data is saved +u_stored = get_u(ekiobj, return_array = false) +g_stored = get_g(ekiobj, return_array = false) + + save( + joinpath(data_save_directory, "calibrate_results.jld2"), + "inputs", + u_stored, + "outputs", + g_stored, + "prior", + prior, + "eki", + ekiobj, + "truth_sample", + truth_sample, #data + "truth_input_constrained", # the discrete true parameter field + κ_true, + "truth_input_unconstrained", # the discrete true KL coefficients + u_true, + ) +end + +main() diff --git a/examples/Darcy/emulate_sample.jl b/examples/Darcy/emulate_sample.jl new file mode 100644 index 00000000..655b46bd --- /dev/null +++ b/examples/Darcy/emulate_sample.jl @@ -0,0 +1,203 @@ +# Import modules +include(joinpath(@__DIR__, "..", "ci", "linkfig.jl")) + +# Import modules +using Distributions # probability distributions and associated functions +using LinearAlgebra +ENV["GKSwstype"] = "100" +using Plots +using Random +using JLD2 + +# CES +using CalibrateEmulateSample.Emulators +using CalibrateEmulateSample.MarkovChainMonteCarlo +using CalibrateEmulateSample.Utilities +using CalibrateEmulateSample.EnsembleKalmanProcesses +using CalibrateEmulateSample.EnsembleKalmanProcesses.Localizers +using CalibrateEmulateSample.ParameterDistributions +using CalibrateEmulateSample.DataContainers + +function main() + + cases = [ + "GP", # diagonalize, train scalar GP, assume diag inputs + "RF-vector-svd-nonsep", + ] + + #### CHOOSE YOUR CASE: + mask = [1] # 1:8 # e.g. 1:8 or [7] + for (case) in cases[mask] + + + println("case: ", case) + min_iter = 1 + max_iter = 5 # number of EKP iterations to use data from is at most this + overrides = Dict( + "verbose" => true, + "scheduler" => DataMisfitController(terminate_at = 100.0), + "cov_sample_multiplier" => 0.5, + "n_iteration" => 20, + "n_ensemble" => 120, + "localization" => Localizers.SECNice(1000,0.1,1.0), + ) + # we do not want termination, as our priors have relatively little interpretation + + # Should be loaded: + dim = 2 + N, L = 80, 1.0 + pts_per_dim = LinRange(0, L, N) + obs_ΔN = 10 + smoothness = 1.0 + corr_length = 0.25 + dofs = 20 # i.e. the input space dimension + #### + + exp_name = "darcy" + rng_seed = 940284 + rng = Random.MersenneTwister(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"] + prior = load(data_save_file)["prior"] + truth_sample = load(data_save_file)["truth_sample"] + truth_params_constrained = load(data_save_file)["truth_input_constrained"] #true parameters in constrained space + truth_params = load(data_save_file)["truth_input_unconstrained"] #true parameters in unconstrained space + Γy = get_obs_noise_cov(ekiobj) + + + n_params = length(truth_params) # "input dim" + output_dim = size(Γy, 1) + ### + ### Emulate: Gaussian Process Regression + ### + + # Emulate-sample settings + # choice of machine-learning tool in the emulation stage + nugget = 1e-12 + if case == "GP" +# gppackage = Emulators.SKLJL() + gppackage = Emulators.GPJL() + mlt = GaussianProcess( + gppackage; + noise_learn = false, + ) + elseif case ∈ ["RF-vector-svd-nonsep"] + kernel_structure = NonseparableKernel(LowRankFactor(1, nugget)) + n_features = 100 + + mlt = VectorRandomFeatureInterface( + n_features, + n_params, + output_dim, + rng = rng, + kernel_structure = kernel_structure, + optimizer_options = overrides, + ) + end + + + # Get training points from the EKP iteration number in the second input term + N_iter = min(max_iter, length(get_u(ekiobj)) - 1) # number of paired iterations taken from EKP + min_iter = min(max_iter, max(1, min_iter)) + input_output_pairs = Utilities.get_training_points(ekiobj, min_iter:(N_iter - 1)) + input_output_pairs_test = Utilities.get_training_points(ekiobj, N_iter:(length(get_u(ekiobj)) - 1)) # "next" iterations + # Save data + @save joinpath(data_save_directory, "input_output_pairs.jld2") input_output_pairs + + retained_svd_frac = 1.0 + normalized = true + # do we want to use SVD to decorrelate outputs + decorrelate = case ∈ ["RF-vector-nosvd-diag", "RF-vector-nosvd-nondiag"] ? false : true + + emulator = Emulator( + mlt, + input_output_pairs; + obs_noise_cov = Γy, + normalize_inputs = normalized, + retained_svd_frac = retained_svd_frac, + decorrelate = decorrelate, + ) + 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) + y_mean_test, y_var_test = + Emulators.predict(emulator, get_inputs(input_output_pairs_test), transform_to_real = true) + + println("ML prediction on true parameters: ") + println(vec(y_mean)) + println("true data: ") + println(truth_sample) # what was used as truth + println(" ML predicted standard deviation") + println(sqrt.(diag(y_var[1], 0))) + println("ML MSE (truth): ") + println(mean((truth_sample - vec(y_mean)) .^ 2)) + println("ML MSE (next ensemble): ") + println(mean((get_outputs(input_output_pairs_test) - y_mean_test) .^ 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 + mcmc = MCMCWrapper(RWMHSampling(), truth_sample, prior, 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(" ") + + param_names = get_name(posterior) + + posterior_samples = vcat([get_distribution(posterior)[name] for name in get_name(posterior)]...) #samples are columns + constrained_posterior_samples = + mapslices(x -> transform_unconstrained_to_constrained(posterior, x), posterior_samples, dims = 1) + + + # Save data + save( + joinpath(data_save_directory, "posterior.jld2"), + "posterior", + posterior, + "input_output_pairs", + input_output_pairs, + "truth_params", + truth_params, + ) + end +end + +main() diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index b3d358c3..8fe9aeaf 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -601,14 +601,21 @@ function get_posterior(mcmc::MCMCWrapper, chain::MCMCChains.Chains) p_names = get_name(mcmc.prior) p_slices = batch(mcmc.prior) flat_constraints = get_all_constraints(mcmc.prior) - # live in same space as prior - p_constraints = [flat_constraints[slice] for slice in p_slices] # Cast data in chain to a ParameterDistribution object. Data layout in Chain is an # (N_samples x n_params x n_chains) AxisArray, so samples are in rows. p_chain = Array(Chains(chain, :parameters)) # discard internal/diagnostic data p_samples = [Samples(p_chain[:, slice, 1], params_are_columns = false) for slice in p_slices] + # live in same space as prior + # checks if a function distribution, by looking at if the distribution is nested + p_constraints = [ + !isa(get_distribution(mcmc.prior)[pn],ParameterDistribution) ? # if not func-dist + flat_constraints[slice] : # constraints are slice + get_all_constraints(get_distribution(mcmc.prior)[ps]) # get constraints of nested dist + for (pn,slice) in zip(p_names, p_slices) + ] + # distributions created as atoms and pieced together posterior_distribution = combine_distributions([ ParameterDistribution(ps, pc, pn) for (ps, pc, pn) in zip(p_samples, p_constraints, p_names)