-
Notifications
You must be signed in to change notification settings - Fork 0
/
exactranktest.jl
173 lines (149 loc) · 5.47 KB
/
exactranktest.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
struct ExactRankTest <: AbstractMCMCTest
n_samples ::Int
n_mcmc_steps::Int
n_mcmc_thin ::Int
end
"""
ExactRankTest(n_samples, n_mcmc_steps; n_mcmc_thin)
Exact rank hypothesis testing strategy for reversible MCMC kernels. Algorithm 2 in Gandy & Scott (2021).
# Arguments
- `n_samples::Int`: Number of ranks to be simulated.
- `n_mcmc_steps::Int`: Number of MCMC states to be simulated for simulating a single rank.
- `n_mcmc_thin::Int`: Number of thinning applied to the MCMC chain.
# Returns
- `pvalues`: P-value computed for each dimension of the statistic returned from `statistics`.
# Requirements
This test requires the following functions for `model` and `kernel` to be implemented:
- `markovchain_transition`
- `sample_joint`
Furthermore, this test explicitly assumes the following
- `kernel` is reversible.
Applying this tests to an irreversible `kernel` will result in false negatives even if its stationary distribution is correct.
# Keyword Arguments for Tests
When calling `mcmctest` or `seqmcmctest`, this tests has an additional keyword argument:
- `uniformity_test_pvalue`: The p-value calculation strategy.
The default strategy is an \$\\chi^2\$ test. Any function returning a single p-value from a uniformity hypothesis test will work. The format is as follows:
```julia
uniformity_test_pvalue(x::AbstractVector)::Real
```
"""
function ExactRankTest(
n_samples ::Int,
n_mcmc_steps::Int;
n_mcmc_thin ::Int = 1,
)
@assert n_samples > 1
@assert n_mcmc_steps > 1
@assert n_mcmc_thin ≥ 1
ExactRankTest(n_samples, n_mcmc_steps, n_mcmc_thin)
end
function simulate_rank(
rng ::Random.AbstractRNG,
test ::ExactRankTest,
subject ::TestSubject,
statistics,
tie_epsilon
)
model = subject.model
kernel = subject.kernel
n_mcmc_steps = test.n_mcmc_steps
n_mcmc_thin = test.n_mcmc_thin
idx_mid = sample(rng, 1:n_mcmc_steps)
n_mcmc_fwd = n_mcmc_steps - idx_mid
n_mcmc_bwd = idx_mid - 1
θ_mid, y = sample_joint(rng, model)
stat_mid = statistics(θ_mid)
rank_wins = ones( Int, length(stat_mid))
rank_ties = zeros(Int, length(stat_mid))
# Forward Transitions
θ_fwd = copy(θ_mid)
for _ in 1:n_mcmc_fwd
θ_fwd = markovchain_multiple_transition(
rng, model, kernel, n_mcmc_thin, θ_fwd, y
)
rank_wins += statistics(θ_fwd) .< (stat_mid .- tie_epsilon)
rank_ties += abs.(stat_mid - statistics(θ_fwd)) .≤ tie_epsilon
end
# Backward Transition
θ_bwd = copy(θ_mid)
for _ in 1:n_mcmc_bwd
θ_bwd = markovchain_multiple_transition(
rng, model, kernel, n_mcmc_thin, θ_bwd, y
)
rank_wins += statistics(θ_bwd) .< stat_mid
rank_ties += abs.(stat_mid - statistics(θ_bwd)) .≤ tie_epsilon
end
# Tie Resolution
map(rank_wins, rank_ties) do rank_wins_param, rank_ties_param
rank_wins_param + sample(rng, 1:rank_ties_param+1) - 1
end
end
function compute_rank_count(ranks::AbstractVector{Int}, maxrank::Int)
param_freq = zeros(Int, maxrank)
for rank in ranks
param_freq[rank] += 1
end
param_freq
end
"""
simulate_ranks([rng,] test, subject; kwargs...)
Simulate ranks according to the exact rank test strategy.
# Arguments
- `rng::Random.AbstractRNG`: Random number generator. (Default: `Random.default_rng()`.)
- `test::AbstractMCMCTest`: Test strategy.
- `subject::TestSubject`: MCMC algorithm and model subject to test.
# Keyword Arguments
- `tie_epsilon::Real`: The tolerance for declaring a difference in statistic as a tie. (Default: `eps(Float32)`.)
- `statistics`: Function for computing test statistics from samples generated from the tests. (See section below for additional description.)
- `show_progress::Bool`: Whether to show progress.
# Returns
- ranks::Matrix: The simualted ranks. Each row are the rank samples of each statistic.
"""
function simulate_ranks(
rng ::Random.AbstractRNG,
test ::ExactRankTest,
subject ::TestSubject;
statistics = default_statistics,
show_progress::Bool = true,
tie_epsilon::Real = eps(Float32)
)
n_samples = test.n_samples
prog = ProgressMeter.Progress(
n_samples; barlen = 31, showspeed = true, enabled = show_progress
)
mapreduce(hcat, 1:n_samples) do n
next!(prog, showvalues=[(:simulated_ranks, "$(n)/$(n_samples)")])
simulate_rank(rng, test, subject, statistics, tie_epsilon)
end
end
function simulate_ranks(
test ::ExactRankTest,
subject ::TestSubject;
statistics = default_statistics,
show_progress::Bool = true,
)
simulate_ranks(Random.default_rng(), test,subject; statistics, show_progress)
end
function mcmctest(
rng ::Random.AbstractRNG,
test ::ExactRankTest,
subject::TestSubject;
statistics = default_statistics,
uniformity_test_pvalue = default_uniformity_test_pvalue,
show_progress::Bool = true,
)
n_mcmc_steps = test.n_mcmc_steps
ranks = simulate_ranks(rng, test, subject; statistics, show_progress)
map(eachrow(ranks)) do param_ranks
param_freqs = compute_rank_count(param_ranks, n_mcmc_steps)
uniformity_test_pvalue(param_freqs)
end
end
function mcmctest(
test ::ExactRankTest,
subject::TestSubject;
statistics = default_statistics,
show_progress::Bool = true,
)
mcmctest(Random.default_rng(), test, subject; statistics, show_progress)
end