From a87aef956cc06c431c99a5cbc3018c89c261e666 Mon Sep 17 00:00:00 2001
From: Nat Efrat-Henrici <60049837+nefrathenrici@users.noreply.github.com>
Date: Fri, 23 Feb 2024 13:55:18 -0800
Subject: [PATCH] Add Emulate Sample functions (#51)

Add emulate + sample functionality, docs, and tests
---
 Project.toml                              |   1 +
 docs/Project.toml                         |   2 +
 docs/make.jl                              |   1 +
 docs/src/emulate_sample.md                |  64 ++++++++++++++++
 src/CalibrateAtmos.jl                     |   1 +
 src/ekp_interface.jl                      |  14 ++++
 src/emulate_sample.jl                     |  87 ++++++++++++++++++++++
 test/Project.toml                         |   4 +
 test/runtests.jl                          |   4 +
 test/test_case_inputs/eki_test.jld2       | Bin 0 -> 29854 bytes
 test/test_case_inputs/sphere_hs_rhoe.toml |   4 +
 test/test_emulate_sample.jl               |  36 +++++++++
 12 files changed, 218 insertions(+)
 create mode 100644 docs/src/emulate_sample.md
 create mode 100644 src/emulate_sample.jl
 create mode 100644 test/test_case_inputs/eki_test.jld2
 create mode 100644 test/test_case_inputs/sphere_hs_rhoe.toml
 create mode 100644 test/test_emulate_sample.jl

diff --git a/Project.toml b/Project.toml
index a45f8cbf..20f460ff 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,7 @@ authors = ["Climate Modeling Alliance"]
 version = "0.1.0"
 
 [deps]
+CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
 ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
diff --git a/docs/Project.toml b/docs/Project.toml
index 92da694a..3ca47d29 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,3 +1,5 @@
 [deps]
+CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
+EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
diff --git a/docs/make.jl b/docs/make.jl
index df070e61..b554e982 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -25,6 +25,7 @@ makedocs(
         "Home" => "index.md",
         "Getting Started" => "quickstart.md",
         "Experiment Setup Guide" => "experiment_setup_guide.md",
+        "Emulate and Sample" => "emulate_sample.md",
         "API" => "api.md",
     ],
 )
diff --git a/docs/src/emulate_sample.md b/docs/src/emulate_sample.md
new file mode 100644
index 00000000..8d02714c
--- /dev/null
+++ b/docs/src/emulate_sample.md
@@ -0,0 +1,64 @@
+# Emulate and Sample
+Once you have run a successful calibration, we can fit an emulator to the resulting input/output pairs.
+
+First, import the necessary packages:
+```julia
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+
+import EnsembleKalmanProcesses as EKP
+using EnsembleKalmanProcesses.ParameterDistributions
+using EnsembleKalmanProcesses.TOMLInterface
+
+import JLD2
+import CalibrateAtmos as CAL
+```
+
+Next, load in the data, EKP object, and prior distribution. These values are taken
+from the perfect model experiment with experiment ID `sphere_held_suarez_rhoe_equilmoist`.
+```julia
+y_obs = [261.5493]
+y_noise_cov = [0.02619;;]
+ekp = JLD2.load_object(
+    joinpath(
+        pkgdir(CAL),
+        "docs",
+        "src",
+        "assets",
+        "eki_file_for_emulate_example.jld2",
+    ),
+)
+init_params = [EKP.get_u_final(ekp)[1]]
+
+prior_path = joinpath(
+    pkgdir(CAL),
+    "experiments",
+    "sphere_held_suarez_rhoe_equilmoist",
+    "prior.toml",
+)
+
+prior = CAL.get_prior(prior_path)
+```
+Get the input-output pairs which will be used to train the emulator. 
+The inputs are the parameter values, and the outputs are the result of the observation map. 
+In thise case, the outputs are the average air temperature at roughly 500 meters.
+```julia
+input_output_pairs = CAL.get_input_output_pairs(ekp)
+```
+Next, create the Gaussian Process-based emulator and Markov chain. 
+The samples from the chain can be used in future predictive model runs with the same configuration.
+The posterior distribution can be saved to a JLD2 file using `save_posterior`. Samples can be extracted from the posterior using ClimaParams.
+```julia
+emulator = CAL.gp_emulator(input_output_pairs, y_noise_cov)
+(; mcmc, chain) = CAL.sample(emulator, y_obs, prior, init_params)
+constrained_posterior = CAL.save_posterior(mcmc, chain; filename = "samples.jld2")
+```
+
+Finally, you can plot the prior and posterior distributions to see results:
+```julia
+using Plots
+plot(prior)
+posterior = get_posterior(mcmc, chain)
+plot!(posterior)
+vline!([65.0])
+```
diff --git a/src/CalibrateAtmos.jl b/src/CalibrateAtmos.jl
index 3773bbd7..1ef3a074 100644
--- a/src/CalibrateAtmos.jl
+++ b/src/CalibrateAtmos.jl
@@ -2,5 +2,6 @@ module CalibrateAtmos
 
 include("ekp_interface.jl")
 include("atmos_interface.jl")
+include("emulate_sample.jl")
 
 end # module CalibrateAtmos
diff --git a/src/ekp_interface.jl b/src/ekp_interface.jl
index dd0215a6..c7d65284 100644
--- a/src/ekp_interface.jl
+++ b/src/ekp_interface.jl
@@ -13,6 +13,20 @@ Returns the path to the iteration folder within `output_dir` for the given itera
 path_to_iteration(output_dir, iteration) =
     joinpath(output_dir, join(["iteration", lpad(iteration, 3, "0")], "_"))
 
+"""
+    get_prior(prior_path; names = nothing)
+
+Constructs the combined prior distribution from the TOML file at the `prior_path`.
+If no parameter names are passed in, all parameters in the TOML are used in the distribution.
+"""
+function get_prior(prior_path; names = nothing)
+    param_dict = TOML.parsefile(prior_path)
+    names = isnothing(names) ? keys(param_dict) : names
+    prior_vec = [get_parameter_distribution(param_dict, n) for n in names]
+    prior = combine_distributions(prior_vec)
+    return prior
+end
+
 """
     initialize(
         experiment_id;
diff --git a/src/emulate_sample.jl b/src/emulate_sample.jl
new file mode 100644
index 00000000..28d2a452
--- /dev/null
+++ b/src/emulate_sample.jl
@@ -0,0 +1,87 @@
+import CalibrateEmulateSample as CES
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+
+import EnsembleKalmanProcesses as EKP
+using EnsembleKalmanProcesses.ParameterDistributions
+using EnsembleKalmanProcesses.TOMLInterface
+
+import JLD2
+
+"""
+    get_input_output_pairs(ekp)
+
+Helper function for getting the input/output pairs from an EKP object.
+"""
+function get_input_output_pairs(ekp; N_iterations = nothing)
+    N_iterations = isnothing(N_iterations) ? length(ekp.g) : N_iterations
+    input_output_pairs = CES.Utilities.get_training_points(ekp, N_iterations)
+    return input_output_pairs
+end
+
+"""
+    gp_emulator(input_output_pairs, obs_noise_cov)
+
+Constructs a gaussian process emulator from the given input/output pairs and noise.
+"""
+function gp_emulator(input_output_pairs, obs_noise_cov)
+    gppackage = GPJL()
+    gauss_proc = GaussianProcess(gppackage, noise_learn = false)
+    emulator = Emulator(gauss_proc, input_output_pairs; obs_noise_cov)
+    optimize_hyperparameters!(emulator)
+    return emulator
+end
+
+"""
+    sample(
+        emulator, 
+        y_obs,
+        prior, 
+        init_params; 
+        n_samples = 100_000, 
+        init_stepsize = 0.1, 
+        discard_initial = 0
+    )
+
+Constructs a MarkovChainMonteCarlo object, optimizes its stepsize, and takes
+`n_samples` number of samples. 
+The initial stepsize can be specified by `init_stepsize`, 
+and the number of initial samples to discard can be set by `discard_initial`.
+Returns both the MCMC object and the samples in a NamedTuple.
+"""
+function sample(
+    emulator,
+    y_obs,
+    prior,
+    init_params;
+    n_samples = 100_000,
+    init_stepsize = 0.1,
+    discard_initial = 0,
+)
+    mcmc = MCMCWrapper(RWMHSampling(), y_obs, prior, emulator; init_params)
+    new_step = optimize_stepsize(mcmc; init_stepsize, N = 2000, discard_initial)
+    chain = MarkovChainMonteCarlo.sample(
+        mcmc,
+        n_samples;
+        stepsize = new_step,
+        discard_initial = 0,
+    )
+    return (; mcmc, chain)
+end
+
+"""
+    save_posterior(mcmc, chain; filename = "samples.jld2")
+
+Given an MCMC object, a Markov chain of samples, and a prior distribution, 
+constructs the posterior distribution and saves it to `filename`. 
+Returns the samples in constrained (physical) parameter space.
+"""
+function save_posterior(mcmc, chain; filename = "samples.jld2")
+    posterior = MarkovChainMonteCarlo.get_posterior(mcmc, chain)
+    constrained_posterior = transform_unconstrained_to_constrained(
+        posterior,
+        MarkovChainMonteCarlo.get_distribution(posterior),
+    )
+    JLD2.save_object(filename, posterior)
+    return constrained_posterior
+end
diff --git a/test/Project.toml b/test/Project.toml
index 84a24c9e..2af571dc 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,9 +1,13 @@
 [deps]
 CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 CalibrateAtmos = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
+CalibrateEmulateSample = "95e48a1f-0bec-4818-9538-3db4340308e3"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/runtests.jl b/test/runtests.jl
index ece451f8..eaa321b5 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,4 +1,8 @@
 using Test
 
+import Random
+Random.seed!(1234)
+
 include("test_init.jl")
 include("test_atmos_config.jl")
+include("test_emulate_sample.jl")
diff --git a/test/test_case_inputs/eki_test.jld2 b/test/test_case_inputs/eki_test.jld2
new file mode 100644
index 0000000000000000000000000000000000000000..924eb89e311a344a11e8405a273fc380dc7d9f29
GIT binary patch
literal 29854
zcmeHwd0frU_y28CNw_Uqm0OX}t|Z&JQK^JfvP;WNrA>>J5ZX|(WQ&qiN+F`MlorWa
zcnjGgh01dGefiDp-ur&z>-~MdeLvrSe&0tg56yke%sDgjoS8G{oVky4ogJKPO#OVL
z2?03Q*wA2KoP%$SFU~0<D$F;=1h;_@6&)NAj<YbcG_%Cuj73$@XA3h6+!Sk5zu*|0
zo8tt`;Pd}00y~`@oQ5lED@rRUU@+36Q3%9OUX19SwO*7qysEyx(AmL58j~h73xl76
zs^d^exNjH%gRxVV!eG&x1^O?A{$nr-m`L9!R2G945e93Aa)jl!10@)Yl09lqY8XZl
zBZ~<Ut~{TZ_(%dqMq0=~_y{{Q9R0s7d_j+hspzxh@06UOZlcd~=PYy^s3B|vJyPQO
zLJq=bPDB*JOt`kagslj<i})E2lwvTB;n9RJzfi(F-_S7MaQCPPe?oLLA==D=5a=5l
z8sq67L<ootB}B=H>6FBD1ALuABYb0~SPQdc#pMXifoHL@C><isHx=jqU3<UGM>KL`
zdPOn)K)WvC;#M&jd2zY^Kk#=67w>6tJq0n{#DAa@<6!UaPe3U6#zaJk2!@L3B{)jP
zyH!<G-YQO?i|0jGO_Z-FrW*haW&Lv8Lj2K`&4~z)@eK|~%jm~fA*W$tK0o8F^uNS^
z_}}m!A<own4F}`l8y*l5X6E4OG~Y|~*C#s0H-<1$TyLbf-fSdl&F0Q`H=7;27EK@8
z5~6iDOk6%Nubwec!QpEMa+Jm8>OZzKAAvy#4<~rV1xKU(g~6zZ>xsCDwxN(jRh(}m
zrvE#foqU5sW1|SpXc%aIF&H&5-S~f?^Yjgi3?+ER`uorIjg5}Rj1tp{ryC=ge^#i7
zk-E5i(m(Ju3`8A@_6;QX#6|`n91X360tAE65YwCd1HD^>zi()8f^g0J`B6xsDW<pl
z_w)-Q#2kEM&{-&^)B0(g$HWDa^{*{1|0n$Z>|C^&p{1b}Mn_Ed!#=l<it>#YH9Y#K
zb|vk-iV&6SipwSYAyb?$?l$JFgs~)JFdal-5P?Ai1`!xUU=V>p1O^cpMBx7b0wO;}
zl*|O=r$9E5P#tc~j`0Z)ne!NoB@PBMF&I3?dKku9AoDL>5<)}wwiqCxi$Z8{N#14A
z=Mck2b5Y|CLr}TQP$`rvtSF<1k@?HJp)gkHE5Kodegp|Q>cW3{?5AJ09$+WcCM82h
zgFyt#P_5X%JSh==(LP}W-*AkA#J3{M36BVlCiwVAY`|b8*$aFK;n5gP3k4A#!lAFl
z;3TDls3?rF<Rx0-@4#66Yxoepn5hy5QOJJ9OJ0Pw95OC_g2dLPv!ozYVx?i+C0WQl
z;TI82z$_Jd21M4Z!oPDC30*4^D2JfW!vB)@q3E+xhHxUVXxf!9fx(2(03T#n`U<^H
z(SLfSq><l9I3@p+N8;l6x3r&iloWnrr3ei-$#CR;(MC{FwCr!N)Bld0{4dz8c+uMV
z8|;MsBT4TRe!=c-*SpaocAI~-EegeMGKE#bZLth9XOxIXjFu>^6{iyGb%eGkpFaBg
zwurD$SM=_O6Q#~!qO?z9x1WEn@H;387olWW!nVKP?LwE$PWR1SqG82~hlN&;a1jYb
zf$&j64ghK7rBe9O%_&*!Lcu^4yToLWFHVv#5f9d4zmKfO05{ANaat($J56&lwHLJ?
zD5LwAsEbDaIf7_?|0O5>mB;n_0Sh<6_X8ICQsWxFQA7=z{AwPX#Pg<rXl2o$*NOW-
zNsTLpeUF$74O$vGVI>nP+UZD2BFEFOWV3>okfA#!3#@h=o~(vfsOC5@c*=O*oC30G
zk_vvI;q<t}=>10Mgn(izHT<~C+F|kUM&XS^Ji1QaR>i9fi`!0UQNiQ(t+jk#r;49@
z*6cH9qY7R(H#_+^#xq@ve@>pdc<=s)<0=36NQpHOiyyv7Q}bye7C-W+%_ZN3Sp53c
z>GYrkEWRb)Sbwq!7O&Fcy?<3F7GJKOTCqb5i@z{yblnGCEPm*PO!if6EPmXJy(VLq
zWATcXHCKOQJj2B6$i?3Hx`^yQ9M7~B6&ZzfSbR{({UL4-u=uRBE5o<dWAR^0hP!XP
zhsCq(j_cV!#p2h<nN~X8#NzjreQLgT1B=&m7-zQ<rP?l!8k-xicv9-OXFE~3&Z*H8
ze`7o^#p7Yuy|4UV84uTX^OyK$EPk>aQTA067C&dohbm{3%9@rAy?hIcAI9C-n%#oM
zFY<UZZ6@mXa0Bx-A8N69_tYbVv#4IC`-zq+ls<bQd*?D5kHO1+Z&Al){U|X0aRZ2V
zlu!=bT^fle9^alNV&wGWj*>8vq<_1k{$hw*^G6P(Xbx!?F%0{$XA-sKe=5q5h?Yfu
z*Hp2{_SO*&0fR9Ti*6m`UE4+N{ws-6D`sP-ts%aQ`VUtu>I!Z~X(%40_34%o!Zc~p
zj&Na$H@6#j$IF}<C{N;6t`p{Cu(DyoG)YrwpdRK7cHlifFC!dpUZM8oH^!4Gh9h^a
zN4|*cKOfKA)di8l{^gI@G>|6s?Hx#Q8i}HGcfP3n*}Q@D?fTOLDW<h}AYF4xoO<qz
z#G6DDbADqytHjf4B`p>Bzcrrx=_>~Mm9)-tAjS2G-_Lo8-tW1I%CnaYq?mon2GV*>
zQN8;0Xu}{pa3oY;*;QqpsACI$z0Uu3!-yAh^u!I!PaoUyojNMDey*GF9*sLNobpbH
zYj3Lkj~JsKqK61Sr?0y%d=xQWEi>w`Lw<vA=x-kGLg6UUW=4ow`Qh(Ti@D-CD?}jR
ze|4q%uat<9-?bzB{PUwUN2nwXAGrICiwcg3A%y?x_E)H2NxCb!Px_vf-Z=fY1Gq>X
zl9G9VfI-T{QdA+DKnVB?nSpCKjDp0B`2jWbg07||N>^}r3^I@Yd<)p<gu#7pECmN_
zbio*ZpW_>_0b}t~hEGUrSR`iZufRjqeFmCz{+VW=Y4`7Q{}s?e?Mi}>ScCo^AdFI?
zQ2YF6lSiV3Nr+3tYP972pHgMLU)M<9W4J)HaudaS?FVDHUQDL2Qg}ox#q{DtY>f0p
zX^^QXWxI(|V^Q+w-#+otBs}UQFsS@lior;}$UwVW<|XQe)sI{L&odHT7h51gk=BmG
z2)_&=i~c=E85vSV8KPA-Md?RsX>s#1Lh1nuVcUpWNeR&^DB(JG24iEnu<-=uYl0RR
zhH99c*S*UFeN|h}ArUlSZ5-dPxsVUcv2pFKbRKwS7Ou0fVZg4Y<MBOfIbcdG=$n|v
zg_Z0#$L|j%0r&db>z0-*@LbY#U)`DyE16q>(<K0v$uD||D+I7^S@FRL#7D}jC;N~c
z4IX^-;WW4t;r<0zj~#2+@J?p<E@}q@*6xpAEc1j3rrOHgwJ98E>E637eJLM!W~s%h
z_gSzuBtZF*JQqC17}ckIV#B$2hZ5%dP~h@8#qnidSzsTQ<XHNi2!S#et`*MbK~{22
zbNzZMESj|ZbD${^%mT~yHs5E1_59_LtL9PR)oJpx_$n3@5w+%@TTF!mGnz@W&^Q%>
zlp^*v^1*;sI7QZg4Gi_T3G(mxK=)=pu{YsBdRf<E%y}lzs`myv_Oqb(_1=vA=`=9$
zkRtl*A%RN4quDpKXpnNZ;Equ~8!l?>7|k6mfQ`E~jvcjT!v>;F%9e+8DA$y#e}<vM
zy+;MB49C$SFQ8yIBog8Fuy$%u03FiGilSszQ^DbRM9B~%7T`aprIW|<;C@HUx+Tg)
zD92nWu5O}0=*A7>XjfTa9Qy9g!#*OcpEf(Z-j@sJS<x4*W9V>gKgpt`ga@)x$8Kn0
z*icn<BeVPz6ZCof&HQ8{BpI3<^Cht$Cwek<4Uq@WYB&em3|Nrg(J-_$i3eI`{4X!Z
zQep4)D({KoC@?2uZqQ{FHW+c&Fh1NS!RR{$-{wrG1N)FcHY=S35gU*66f5yzZN|$b
zi$77IAl)@@+Cm;kPcB({VJsbn-K0<A$Z{a$YQcvxSw6&gl*?KnysB~n$9x>ghR%EX
zD$94$A)3~<bb2cT9`%kq(1!ZW-I^8Zeuf3xwKf)Lb~6B$&koO8OaVvo>euQDba?Kv
z%4+f$7JM1`bX>a!4;J?944Ju&4x?BOu`Z*?a1?v+*s(eyco<LKEme%l7n+|T2eM(?
zw<j(}`8?SBWLw77TO7!Kba(CW*;EMJG{uM+DuBw(ou`jFvEVCdpXJMmY*?~;UY3Cz
z8Qd?eH0b`ogVX*tm!HIuz>(!P_U&>3r0TsQS!gprd$iFhb4Nbhd+nWw&E`O}_QLgA
zY!ZC@<4SR18w*lB7!~GvWU$@3b8DSB9UjVHCy;|UpjUIb+<pra<ZXt{HYL$PYH{qv
zFS0y%bE@jDH=0-O+MAW1i`h_*JL3J|4INx2-s<0UlMHiHa$Aa$sSx&U^B?QtnJ~-b
z==Ha%d?=d0w{N-5hXeK--VfQuf{sg8UHhIBp{J;Km_U{ZHTn*bn2QuJRWY&A)Mr3j
zVxG*&Iy&SJA@1sW&4upJsKY5_I;?1(;n;)Yz=I2_d{-71nv06(tte$ctzc&W+mHll
z?I(wH&gVnBw}VHuJs0A}D)yv}LU@@L*v?MnLF$(Uk4IglgXg0)(!m#aFic}pzo9JZ
z?@l#^JH8}1mlk?VXDS6K*WC}g^blcd!t(el@<iC5|CPFkM+Skz0k7OB8VqqK^vu=f
zfV7dxU6V7Yy^lAqdtG6J;bZBA8*U3=TTnn=Rx}&3OC569Gs*BJc68d*TU1!*zi{4I
z3JYFeBDBYaa$(K1zM8^2e0Y=Ad%pz3gvF3EMwiEfwcS3{sSF|{gxu;beMkcRAvKHC
zklct^@wWJ@2?gY+j0cTYEZDoRIAqIJ9;kdu3_V@Jh7Kbm>gJOyh-?}ezT2MxTNh_M
zojH~denk$sIrReQ4W2&hVG|j??kTo6yh{NV=j-c$`y|lTd91q~&w<3(`;yO+n6Orx
zYLLE|3<aN)O3zbya5>pLbZaOb>{@QP?o~%}%aWO5eU6XDTN`CIlMA$@;qN~aP=9yw
zMn|qBLTS?+nbct<cvj(Ndixd`Zs`j=G+uE*A#!J<+5?39;;mM?Z@6&fb@}q=$EmO^
zem3_jhX=jkN%RZ8OnAQF!9+I>8}fFlywdwb17h@)lw~(Ku(2<6O13Kn-Y2G;g*9=Y
ze4CZQ(6wBc`hJy1OC29p$^Ee*D4GJj2VyqsX>(z;(vz~NRthL5`Uh4c`SC}p*7Wxr
z8oWrD!1hOY?Z$iMZmQ>jwue%klnx0lU>_9Ux8}i~zA%e3whSozAhlF}2?MzNij<gz
zY-mkQl{bFDgUuA1m9Nipphz04el!uCGl%!hs(HkNBYDS3-KSa5w_uvvB{K%>?@%Z(
zzRiI2eH7Z1BqH=XroQtxWkP#pTv?p~1zv<yO<fa8f&01IX&ej(mf!7&Cu@=+)3KwH
z8qWcz9ZA(ASagtiecRE>kqh^hm#8ss6G6w$?@k$x0xfMLhh?#NFm>gPX;&~bDAUyF
zVv&jPN9VXXpDc-B+~-t(`4JgBG7V1MEaJdf*-X~Z$s}l&A8ruoD1b{%?TiK4Y`8DG
z>nzum1O6Y|ywnPr@M=k-XNEBwcG%ASz;j_il*t~l?+!Y=Q7G_Q)XD+Bj_rZEBbmV6
zqCAv=<`qM=%D<>ig2mX>lS3D9!M@;O=?MlCrd_=AFjH0lFQ3N%%%%e<&531==EKFV
zIPHx{j@mg}rO~RH5P!CP+oyYMSnCsWVVp4$RxkOewVuQQyKjdYci&)wW17O`Z>yM~
zsg$N1GeZCzv#RK>E-KW_lA4lwk_57jW)oe_h%mmoiIDP`2w}K+_Vb?cp;lwHD@m6G
zZOpyioi`b9kgaT!caaEs%cyU!&tL(gRVP^SssJ?3S*C9<MRHwfRmZ8bbZEMJG28qI
z7x0ZM7SyB*U~T=&Yj2`CATze_)9OtWm?%H%IBq<`+3%a5b~PHW${m%yN(R(#t_F%U
z5uWtNSDN>6;D~PaXQLHV;2658`P6Zsb{BQ(g5xZpjV0P9Od)}^t`i*br9jc*b*VOc
zcn~V<Zudl)391G4O%5N~pr2?mauzxd+v56MabLLbV2kc!FLdrdNWIhF8_fdlL2}ws
zCKHZUkUdQ<@}T3+(i^5<SkT^Zd1d)xK3IKNu(%V=_iZO^^zl*($R=DIYIGjWyX;f!
zBn=L9(3wfDAK2iQGwb9BYow=C#;siRn1=Rg(5Po~Xi(qnuT*)D1LpagT1e>J)TYf;
z%X~uuOZgi^{G7<(ePYALH9BNS(5hOZp~{0g*Ytq8oos08h{$%?L<RYg*g$Et4_Rb~
zv5Z+97&)&d?Ffkt&gLFUmbN_LEvg*Tt1E!8wD@+i3k5t{+)79d40!c9?||${CTzH%
zn9i*efOLl6`WbOl2;a8v*&;O(TpgBnE!dp{d%lEd&g$gAy}*#GEhQ9Kz{_cU`<w=c
zA2-xw&!9r*3WFtl0U0u)aRrPI3~*u|w5s6<;FR*bunIJP<~FB^_j(0zHn%x^3yTJZ
zL+#`||DeJIA~|69a}vB7t=+wp#sRZY8>1IrWkasq)w5aGsNne3eI~h*3k!puj1EHU
zy{Dk`THsw4X!YOY(6Z6_nydf7D~<-i4ef9A6<KgBxut&QLk1|`x`m4fpu(MqUh63W
zDi{ck*Kca0!iyQL3rjc9!A%c?-RRDR;>+bb6>sz5gTiO61r#QHoj6Q!x)TZ1U|!Z?
zPXY9;_1a|B!H1JKgLgjK&w;+Ny)utFkp7%J=4ujw33k?RQmTxpU=(P}TeY7L3zDpJ
z-|s+rYHvo(#6lKqX-!q@)gvLj(7bHM7XcVYt%@tnXMu?hPG#p-2JGIp=Y6a_1;R{n
z-t3slf)B2W9an-$kaCIFZ6Cmg&LQhQPc-3yor|iM?npWmXqL#5x1xT(EM^a1&wvv1
zNuC9B1aR%N{jwT$CN!MYj6a_CyzI!~T7uh5#qup;^I}SRRfoueNFDQ|1tBqD{?_H-
zH3nG{?qfc$k=m{m!-B&`<C*szxG=9+o3UUd3!=YNMr~iofGuD1=X+h^!qf8!xd+ei
z;ML@Y%A+<E@GN09jcK96_CM+}hEXW6aQnOfZ=_+lJhoJP?8}BZTCu0gOo$M2XG)GT
zmIR@bBXeIZ5WwZS<ih7m+2Ch_<JO{+(^IaIlpR6^O)eqIHVkRnQSZM>eImkKeX?W2
zCLZ8)mp__mBY=<6?gv~V1dv@}7SN+FfDWgb!=g@cL8`b*C*F_;xu@JV_U|Tw%FMB&
z?01mhmR+hEaVQO%Pj3#`K8y?+4WTw=TpBz|oPz1-;lnw{OE0^Jpc9I_L&XU7pVmM3
zTCfEXsKHum4__xjjB{Vs0wOA(GFic?iUT&3jNQcBRIqHZ&F(N~fPLvZVgQB@cIy|+
zIa*7Eq35d~Z$OrX@`>F3#Ah^!RtGBiDHR04S!a~#G_c{i_9O(6Aw6LDhz%>q(1T~S
zwWU+xflk)r^9gj=syB+F%4fg|o$l4!cCjHgz|T_}&xDu0XRf}VM1|ap3yJtN7Nq6m
z8GJrV1t*)kABW^o0pqHX#(B?zk>N5QTal$=kZGB=dIUP@N<81F9wx$=qwwmu8yoHn
z-Nv*<8ptPen}mWuCd@rjfA22ZaRPb8)$RRsQ2RK0QH?tZ&bau5FEOXUgme1WPTI0y
z;yBA0r#Dl;`t07Op1U+y@V3f&pA!WNr80DrzVKo4$9uhTGk9><+O+!WArc%d=BH~I
zAq!>feZ!nFWJp;E9808OoS!FqE#MOmLKuAN(gL(D*wS&yMO?6a#;X|jfeve$G^{%&
zBTd?WNa#gm2~Dm0xVKV`3R;=_wmoDsU~kcdM<&}yaH`{6`AyV+s*kx}I#A(ZXb?~v
zNkC0rn74Zw4~DF>txZN2-?A6#Q+w9Z;HrLD-@R56++Vq(1BWbLgEpJ5c|knznBjcB
zb1eyu9lI9qi!}EZmCKz<t~5A3pB%zsZ~%Sb`14!DA9p<Zux$wyGOvuR^mu?YlPVLl
zN9d$?F^s)vZ@>hV_}Wi<vlwt7*L%tkR6gpD^@;MMD1d3-6vD~n0aqyoQ$Ae)erny5
zHrJqiRFL9_-^_rd+oMLjPC@&^>|17EA<~5F^Ld^%0!X+kMbK}jga3rp5oZ$_kXSio
ztMx%1*w=SP+7g*y)pA&BMn4yv^EF$Vrx4*=JwNCE8D#Ot?VYm%M~09lcb85i5ep0D
zuU3wxz?-mZ@}~{?K=Rlsb$JvCvQFm3j@dv2Ylkg@lL{<Q_Aks~WV7JS))%(7UQ%Ip
z8Owqx%>*U(m*i<#bck{BjPPAThBVSXlknwikct_)YZzMqpK-B648AZRCKc<+jblTQ
ztCQTF6-c6`+9zig(_vwersn?7WO#aT;{9oUJkU97|G-a%0Wa%5e;cca_LIMX_bn<7
z+KmEQ8mp;5>A#h)K9&PlHAi%NE~3Gh=bvAj7t*2L<NopWB`l~4yJ&IqG7(bd_Vwi#
z^Fb!G(Q_z`1{Z0VEo?NuDY9`J6p#h+z#zqI4Ul0(-3}Gq`6P($d$2DrnE?fthD8N`
z6hJjswl?Yl4c0rTFntTjaE`ch;`4sAUWRL8i66NjI5jR+#)|{are~gy0VdqKApiEo
zbrwvB(Ahi7h7aSvE!%Lvj{)lJXrJ<M5>$1x_f^NzVBE<27ZxGS+iS<st#1#pp|r3N
zWW9KBO?8NgJBI--*7?hlqy*qTtf4G;p#ah}1Z_o4EI1hVAuueS2rqJ6gH|)xV7%<b
zXZLnCC`Fzg+p5oj>Q6paHX0<zlRB*3u!aIode&`8ZDa^P)p~N}buw6*UoSg+nE`oU
zT;tQ9qV-Sk?|C?l3LCyfY~;4^;8Sd%+UBi1c<Q*6%10;oP)GZPmKHQ{nz`iUVhRUd
zYEF?Z9M6PR?K)j1m-&!(!vEUF3NpypD9wLf!GaC9;wti^xp4GJzTs>nzg7*WI$YdN
z1HpLHEgFh+_+yXj5>G!eoNrT~9-fOdU-Oog^N_{ign3x)E5(F&%$njc$uv-@mD~3A
zEg$NtO&;I!V!~9)3g?bi78DQ@1dH9tfEhDAED@c2560;oRgWaX()-rYX1zoxoV%s%
zKnxGuHm=$=GFbp8Ei8|ZM;gGJafefbmmzsjGr5094jE=aYJyjp0EiK~t6R`I;jc6)
zRjsDOe9z=9Y7hAE^;xH>^gSMoI$1T{cbNc;k}F_5I%g)Y(9xTY<k`97*RoCp6JdAo
z&0_p2Hu$QZ@Y;ne;Y%wCU6CtjaO#*^-hqB9c$j~9%vUDE#@zavAfZLz>ZK7okp@;T
z=O3PH!GepL9w!}iIWUo@m!*l;!yQAv*y4Q@P&Mlx=fy-Gv@CjZy8;^?zrai~BO^Jp
zw>F_GjStrfZ6a02u^{PEck1g49Oz&3b#1T<((I&8M$wOQ;FDa9{TgHu`dm`1JlV$r
zB5QZH%Sm)@s`l7!m_veawtiXBG%~DAj$>h|RM>XTAo0~VDp<|%(kc%ZfTc@ATq1)D
z7Y@{RuIr}5#C^}M*1qC_oLTB-_An}xu2KKE>M0W}Hq`Q5kz9}Odg6RBhy;3g`(mAq
zbnw{tn%9KdH%vK~m=K0MdKw29&+?hDTz=Z?PA?iP`PMw8>l_CXaFwwZNFK@D$ev|l
zhdg~IPxy^|E_C^{&170q-~_X7&P5k0n4QZFVI5*Z)z*Wg^GLp}W9VlV++u;()skbR
zy>yUUZf1KGS-M+&x@8+Wx!||Fp`i^~gsU^69CeUIJmfO=&1h#j6uoY-+)v<u;GL&K
z@=Y4V-<50ItSW$%+KR|AL?RS7>xG{|dWOg6>PKgg-k^HwiPeGe44^JcXkF4tf}&G*
z3w4kjcWkX%tiTn3e5Xnx;Sln49nWZ-ua5Sc>4+_LP8^tNtFSshg9-)jT<;vQ<G`66
zMH&95sW5j$b^Pjk3|N1tU26l-;f&)Z&nRakkNut&JikPQlH9ffl?zEwvEvL*=?NX?
z@n_Un`BPv|xO4g0MI7joiSv&$r$NYs90S=a6bN^&v%QQoc9;9#)>DzC6BfwWGUo~h
zp5AcvnSeYZGs5RzckyF@+#gM6EKU<)(`thkBabpcIlv*N4CzrbLu*4crXmld+NBj+
z8`)4zJ5^lM!i0ngHL^)|WN42y82@-U!hslirPhuLA%`Z+yNgF2z;iv%X6>hd9rn0=
zA^~~4%qnUslaa=2o2pp-iUm6sy-1D%bPfl+7{YU<L&n1ld`kimj>M9tMLyz!F0Fiy
zN+}(J+Wn60N9XDAt!sN#GMPa4z)z|`I23k-kF`-CLXA?X$u%Y90a|$==MFlj?@iU=
z438ng2<r0k(HD?jV0X$Y-<Ss#-$t}nB9F^B->L5wkD<f1<Q!^_Ed_W_?R2cpvY|YS
zcKT2Y(lcqY3j(~zaJ_i<*M`XgU^y2)C2G@Qv+VI{w~)SBG5x$(x(~A0Jr4z!BTs|+
z+L~`8_fp`J(~a<G9XfR76jF7YNgzK-ZMX6WCUg?h?xk%a!lTGIyO)_<u(rEkIDP>c
z^r960Q0YaU%^@nD*#QCwcTZd0s4W2R97@kcq<`I-J#>`Zej?IavMy~m7QnK;-ueJ!
zF)rC3Q0F{^0wr1Yq<1Am7@PTCb;4mXJeE-&qv^(nLk%|TR916=kd-il-%W;T@=Y1C
zXBkj+w(Vdb!liA^_8gxoB!AiMw>&a=a5i~e*6sN;m@wR_L>hSr80yuRmnE>l?$b5e
z_-HEh&y2sC;zk6{+=6}2mT<r%38(cqnhqNlHFxFo(P2c|yqD29K8$7LDKyGa!0h6@
z39aX;Fm1iG^zaZqXpX}s&pX0{-FwJy64a6Bc(wGS(Le%oMT6O14ispd&s}6FK%UYH
z)3~?$>CiYbSV=E~1UpJi7U0Y{u$Z}V+QDENe0koiu)>W3oUAKN3ucqx@PyX9YiVeo
zL?6E8UBrZOJLsD?&*p(k*;L#vB=>QfGnwxn5aG;{UYE6jT)6SX?DLBABrv#bS6=^!
z1g$~ER`O$M;F;d#=$^=e;gmwNw`hE-&Iv8DNL~lw%8L|tQ-GS$5fBzg0X=rF)G#D3
z?qLbq!Du|L8?VGPdJDic>0Qe8qb$gpw)k4z1{PGRn7_@@L;cJP8O7IQLAT4RlKe|F
zaC_j^MqsnRao90Tfh`9Xne|-wM&ZI)y-d}HR5qw;$Gx-I!U9u|^7$kr2j47_nxj=A
zfUZe<nB$+(!MfSL_N5FBNa-tEXCe>6{)_jT$D^BuisA=p&BZjBDwvu#X%7ubyZcjk
z=v=uF(=cBbS;o1X2}4{PIbflMH*(cN`dGQq=9zMQxcRMq{UIcG;%Ussb$hu;Z@9nR
z2jQ_~k#e(C4H+sz`}{48ctCu(m%aKO9lBN2DV>%~h*M6a+1*B-VcmBeWwc+X&8RvO
z*g^r7%=?YINyxL2SG(nkF$*v{hL($s>5%iJjCI1F4v+H=J?C^%LHccA@K$FMc;EKz
zI2lET=a0Ab`Ze=lU#7Q4@F-+YbvH~)LGt_6otMfBV+C+^{<=C}Z{!(w$HoMYCPMsz
zwB&?Ed<c6p#+2=h)}3$5Zp=O++?zOV<!0nrI9AqUYvE0Xk5#eSM5Jd0mH2&mPoO~S
z`<U0qvXMv1PF*+PC>gf(U$CbYqxH3S=;gV{qkp2vQz;JVf9GFNl?`??z|GO;$h-(-
zSKce>efbC4AM-kExd?C4*Jf#5l)v}Z`D3N70tnLNxJ*t&o|Y3*U!=Q{=klR?*yub4
z^hqnwR9^~U`k7U;vAekN;-mVbHCNHPoIbNS5qa+BXJ6RC`-VJT#+WUe{F%^ocvyw`
zJPsT$&HfyDj1P^mZw|gHr@?fb^f7OA{=f4{D&(A?!DvoyNuDtUw!BMzyNySJOb6wC
zHBmh9o!UoNMdx4qnJXdwNUu~)&hT7>>_~Sneh3@gz%`t8eY>WY4toyhXGl#VL&J@J
zj@}kJcvWwzO-w=UmrhDQWJ!V~fBn=RS;RLlEJcmVMdxNw)FpIodn{FIR7ZNn4%@=X
zuUohfD^sM-;R(R`6SNtZ(cxu7LC$DA8RR}*=DT0vga5Kstv$Uwh%?d~Vz@*A5nIUz
z+vp4kIpcV(W*Y;1$833$qt64UE8Dl^BYW^LO}?M}f&uH=%Y2QI$9GBI;*tKwOt^S^
z8D|fY3$fFcte+zLrd03I=R)LPNWSs@Qv!_zj~+c5eNmbVCYI~dPa=J1f`_v2wOz=A
zJEmo}Mh)^1>$LfuSjU0G8qWo0Ymr=vyN|`8@+kJFcTdqiWb9d4YFI@BBgMzIot``}
z>dtO%dq;!kUy6sD9H78WpO;0R=*Hv#FXBrC@(iEqsZO1NJkD`N6@?S6NFWGeF&)tj
z*!aA6;X7=(@NJq)$DMgxDA?h9jku2o5nakrr;umZ<jSN|*-B*C^UdOIBeGwQ<i;J>
zE+)YZ`Rt%3IX1B5=XBU>@ZqYLDM<x+$`}nLd%H?V5ZKjPKtu8>{?te<@Ab$d{xx*>
zIHb=gCWU$*{KA8sUSFl`C{!4Aznp7?`~`Q5WghM{qe9x<?)geculhJ{Otzqf14rv7
zZ--)ZZhu=O(>nplHHCw2a+fG@Z>w5DeIp&>xWoxXXOKtj#NvbY`6R%p_%0_Duwi)P
z_$Jl|4%#;u&J_ow=l1KYR9jDmVQaF^PcI`w?4G{mSakk9oV0AocwYg~dJewMSwn`=
zpZ6Eto<f6+qRq!{A`iCPqBG~SS21Aa#OtZv!DOgCwPn{Aw6FRPEIgy>%mMbU75lSn
zkSBiBwijOe(M=M&_}SWrOxTwe{D?P+13uX0vleY)LnbYpb9O!*TKCUsx0ua=DzB{9
zSJ)i*TF8Gk(}V`o-W=rNXHj9-mj(RFn<RKURA!vgawarZ_@8aWF@blXe056(6QVwi
z+Wr>#m*&^4>$v`m420!6mKS3vaCMgV?#fs;?6fw_7|vwDQ>x8`Ms&02ueI;v(p)Me
zO(wiMl1+yF5n7#{=mxp@+}C20UJ5)_HH>S%$OK$;=e72JCc1Gxt6>pJhF4tUQ{#3a
z&*rG1#`brSpC`s)6|82%zMg20N(&BLqBO5rGasE}cGH{uPjEp+cC>T4B@go1`&#(X
z$OEcD^{hhjg}qLrcM0M<cvC-kQ$Ki9KX_9=cvC-kQ$Ki9KX_9=cvC-kQ$Ki9|9|(U
z-q<DAM|2~<<)<5Yn)uFs;5)tO+q|OhohM-w#5eK6Z}<*Ag9r>FFo?h)0)q$)A~1-+
zAOeF33?eXyz#syH2n-@Hh`=BMg9r>FFo?h)0)q$)A~1-+AOeF33?eXyz#syH2n-_d
zBLeyr-mU{T@t6=9{l8+{{q*ZUe?@cqS7FhFKXxRE2PXXeDPfdd$-6K>?VW<wQ=&%y
zRg7-oZ}>>a#j$!(&R@^Fto5ZLcE5~ZnT*2p2_H#JTz;77y}Ngkyof9Yh5p6cpj_dX
z5B}8%g+HGp>5@3au867luhJOdsO3<kYgr5-EHWnkhu>ur;vsd$PrFLg{Ma92vq_5r
zm2VaIQUQfP6NW__-t;wFRNY@j`>%<TMDYI288Sf}CH)tM)ct;jgrT0>BMt6~TKiWs
zB>Xie2|IBFSd{bE)4F9_k%-+dXUJ3>9a=I&qasD57KGC{5t*-y=3ir8M~Dh1{BG>U
z?-!9UgsE_<C0!Oq+hhA=Mf}Eyg%SHC<iZG|-vcO&z{%+JQWdc?`(bu5l3;ZTs1gQ)
z|8rFe83>;=h5A?j4OIFOqG}ST+{92(eDK;Lven}D!~vcqd?o4MhPoIlTxU2D_g@6T
z6cwPLIsH89KtobfrKAUf{Yj$yNm9jursPojQdlVk3(;dK4DTm=e#*C!L~Z{mf3hSZ
z`A_-Qk_gv7<=aR?CjOK^MG_kEr+iyU?8Be(r%I&7Px*F|{rF?PrG;ee{*-SinJDo$
zm_&mY|Mi9aO`^%l`EMb_4M3<Cjed+Slz@KEZXl?5gkK23KPFTx2<6ZeAeiL_;&(=0
V7Ow`3aDV^KgZ;0SedH}!{|_1+BVqsm

literal 0
HcmV?d00001

diff --git a/test/test_case_inputs/sphere_hs_rhoe.toml b/test/test_case_inputs/sphere_hs_rhoe.toml
new file mode 100644
index 00000000..a2288540
--- /dev/null
+++ b/test/test_case_inputs/sphere_hs_rhoe.toml
@@ -0,0 +1,4 @@
+["equator_pole_temperature_gradient_wet"]
+prior = "Parameterized(Normal(4.779568,0.31223328))"
+constraint = "[bounded_below(0)]"
+type = "float"
diff --git a/test/test_emulate_sample.jl b/test/test_emulate_sample.jl
new file mode 100644
index 00000000..c392212f
--- /dev/null
+++ b/test/test_emulate_sample.jl
@@ -0,0 +1,36 @@
+import JLD2
+import Statistics: mean
+
+using CalibrateEmulateSample.Emulators
+using CalibrateEmulateSample.MarkovChainMonteCarlo
+import EnsembleKalmanProcesses as EKP
+using EnsembleKalmanProcesses.ParameterDistributions
+using EnsembleKalmanProcesses.TOMLInterface
+
+import CalibrateAtmos as CAL
+
+y_obs = [261.5493]
+y_noise_cov = [0.02619;;]
+ekp = JLD2.load_object(joinpath("test_case_inputs", "eki_test.jld2"))
+init_params = [EKP.get_u_final(ekp)[1]]
+
+prior_path = joinpath("test_case_inputs", "sphere_hs_rhoe.toml")
+
+prior = CAL.get_prior(prior_path)
+
+input_output_pairs = CAL.get_input_output_pairs(ekp)
+
+@test input_output_pairs.inputs.stored_data ==
+      hcat([ekp.u[i].stored_data for i in 1:(length(ekp.u) - 1)]...)
+@test input_output_pairs.outputs.stored_data ==
+      hcat([ekp.g[i].stored_data for i in 1:length(ekp.g)]...)
+
+emulator = CAL.gp_emulator(input_output_pairs, y_noise_cov)
+
+
+(; mcmc, chain) = CAL.sample(emulator, y_obs, prior, init_params)
+@test mean(chain.value[1:100000]) ≈ 4.19035299 rtol = 0.0001
+
+constrained_posterior = CAL.save_posterior(mcmc, chain)
+@test mean(constrained_posterior["equator_pole_temperature_gradient_wet"]) ≈
+      66.046965013381 rtol = 0.0001