From b562d4c8ca48d9dcef457bf577a0b7ef76111de4 Mon Sep 17 00:00:00 2001
From: Ray Kim <msca8h@naver.com>
Date: Mon, 6 Nov 2023 19:46:57 +0000
Subject: [PATCH] add two-sample test

---
 Project.toml       |  6 +++++
 src/MCMCTesting.jl | 13 ++++++++++-
 src/twosample.jl   | 58 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 src/twosample.jl

diff --git a/Project.toml b/Project.toml
index ee5ebb1..bbf90cb 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,7 +3,13 @@ uuid = "9963b6a1-5d46-439c-8efc-3a487843c7fa"
 authors = ["Ray Kim <msca8h@naver.com> and contributors"]
 version = "1.0.0-DEV"
 
+[deps]
+HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
 [compat]
+HypothesisTests = "0.11"
+Random = "1.9"
 julia = "1.6"
 
 [extras]
diff --git a/src/MCMCTesting.jl b/src/MCMCTesting.jl
index 120c827..a1342ba 100644
--- a/src/MCMCTesting.jl
+++ b/src/MCMCTesting.jl
@@ -1,5 +1,16 @@
+
 module MCMCTesting
 
-# Write your package code here.
+using Random
+
+function sample_prior_predictive end
+function sample_markov_chain end
+
+struct TestSubject{M, K}
+    model ::M
+    kernel::K
+end
+
+include("twosample.jl")
 
 end
diff --git a/src/twosample.jl b/src/twosample.jl
new file mode 100644
index 0000000..5a28382
--- /dev/null
+++ b/src/twosample.jl
@@ -0,0 +1,58 @@
+
+struct TwoSampleTest
+    n_control   ::Int
+    n_treatment ::Int
+    n_mcmc_steps::Int
+end
+
+function markovchain_multiple_transition(
+    rng::Random.AbstractRNG, subject::TestSubject, n_steps::Int, θ, y
+)
+    for _ = 1:n_steps
+        θ = markovchain_transition(rng, subject, θ, y)
+    end
+    θ
+end
+
+"""
+# Two-Sample Test 
+
+Algorithm 1 in Gandy & Scott 2021.
+"""
+function simulate(rng::Random.AbtractRNG, test::TwoSampleTest, subject::TestSubject)
+    # Treatment Group
+    trtm = map(hcat, 1:test.n_treatment) do _
+        θ, y   = sample_prior_predictive(rng, subject)
+        θ_trtm = markovchain_multiple_transition(rng, subject, test.n_mcmc_steps, θ, y)
+        vcat(θ_trtm, y)
+    end
+
+    # Control Group
+    ctrl = map(hcat, 1:test.n_control) do _
+        θ_ctrl, y = sample_prior_predictive(rng, subject)
+        vcat(θ_ctrl, y)
+    end
+    trtm, ctrl
+end
+
+function default_two_sample_test(x, y)
+    @error("The default two-sample test only supports real vectors. Please supply your own compute_pvalue function.")
+end
+
+function default_two_sample_test(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
+    HypothesisTests.pvalue(
+        HypothesisTests.ApproximateTwoSampleKSTest(x, y)
+    )
+end
+
+function test(
+    rng    ::Random.AbtractRNG,
+    test   ::TwoSampleTest,
+    subject::TestSubject;
+    compute_pvalue = default_two_sample_test,
+)
+    trtm, ctrl = simulate(rng, test, subject)
+    pvals      = map(trtm, ctrl) do param_trtm, param_ctrl  
+        compute_pvalue(param_trtm, param_ctrl)
+    end
+end