sketch_thompson_test.nx source
↩ module page · 140 lines · 5427 B
1// sketch_thompson_test.nx -- Thompson Sampling verification.
2
3import "syscalls.nx"
4import "sketch_thompson.nx"
5import "sketch_types.nx"
6
7func iabs(x: i64) -> i64 {
8 if x < 0 { return -x }
9 return x
10}
11
12func main() -> i64 {
13 // ---- alloc rejection ----
14 let bad_small: *Thompson = nx_thompson_alloc(1, 1)
15 if bad_small != (0 as *Thompson) { return __syscall(93, 1, 0, 0, 0, 0, 0) }
16 let bad_huge: *Thompson = nx_thompson_alloc(999, 1)
17 if bad_huge != (0 as *Thompson) { return __syscall(93, 2, 0, 0, 0, 0, 0) }
18
19 let t: *Thompson = nx_thompson_alloc(3, 42)
20 if t == (0 as *Thompson) { return __syscall(93, 3, 0, 0, 0, 0, 0) }
21 if t.n_arms != 3 { return __syscall(93, 4, 0, 0, 0, 0, 0) }
22
23 // ---- initial state ----
24 if nx_thompson_successes_for(t, 0) != 0 { return __syscall(93, 5, 0, 0, 0, 0, 0) }
25 if nx_thompson_failures_for(t, 0) != 0 { return __syscall(93, 6, 0, 0, 0, 0, 0) }
26 if t.total_pulls != 0 { return __syscall(93, 7, 0, 0, 0, 0, 0) }
27
28 // ---- order-stat sampler: Beta(1, 1) is uniform [0, Q14) ----
29 // Sample 200 times from arm 0 (s=0, f=0) and check mean ~ Q14/2.
30 var sum: i64 = 0
31 var k: i64 = 0
32 while k < 200 {
33 sum = sum + nx_thompson_sample_q14(t, 0)
34 k = k + 1
35 }
36 let mean: i64 = sum / 200
37 // Q14/2 = 8192. Allow +- 1500 for sample variance.
38 if iabs(mean - 8192) > 1500 {
39 return __syscall(93, 10, 0, 0, 0, 0, 0)
40 }
41
42 // ---- update increments correctly ----
43 nx_thompson_update(t, 0, 1)
44 nx_thompson_update(t, 0, 1)
45 nx_thompson_update(t, 0, 0)
46 if nx_thompson_successes_for(t, 0) != 2 { return __syscall(93, 11, 0, 0, 0, 0, 0) }
47 if nx_thompson_failures_for(t, 0) != 1 { return __syscall(93, 12, 0, 0, 0, 0, 0) }
48 if t.total_pulls != 3 { return __syscall(93, 13, 0, 0, 0, 0, 0) }
49
50 // ---- update rejects bad input ----
51 if nx_thompson_update(t, -1, 0) != -1 { return __syscall(93, 14, 0, 0, 0, 0, 0) }
52 if nx_thompson_update(t, 99, 0) != -1 { return __syscall(93, 15, 0, 0, 0, 0, 0) }
53 if nx_thompson_update(t, 0, 2) != -1 { return __syscall(93, 16, 0, 0, 0, 0, 0) }
54 if nx_thompson_update(t, 0, -1) != -1 { return __syscall(93, 17, 0, 0, 0, 0, 0) }
55
56 // ---- 2-arm convergence: arm 0 = 80%, arm 1 = 20% ----
57 // After many pulls Thompson should overwhelmingly favor arm 0.
58 let bandit: *Thompson = nx_thompson_alloc(2, 7)
59 var step: i64 = 0
60 var rng_state: i64 = 999 // independent rng for simulated arm
61 while step < 1000 {
62 let chosen: i64 = nx_thompson_select(bandit)
63 // Step LCG for simulated reward
64 rng_state = ((rng_state * 1103515245) + 12345) & 0x7FFFFFFF
65 let u: i64 = rng_state & 1023 // 0..1023
66 var reward: i64 = 0
67 if chosen == 0 {
68 if u < 819 { reward = 1 } // 80% success
69 }
70 if chosen == 1 {
71 if u < 205 { reward = 1 } // 20% success
72 }
73 nx_thompson_update(bandit, chosen, reward)
74 step = step + 1
75 }
76 let pulls0: i64 = nx_thompson_successes_for(bandit, 0) + nx_thompson_failures_for(bandit, 0)
77 let pulls1: i64 = nx_thompson_successes_for(bandit, 1) + nx_thompson_failures_for(bandit, 1)
78 if pulls0 + pulls1 != 1000 { return __syscall(93, 20, 0, 0, 0, 0, 0) }
79 // Arm 0 should win at least 750/1000 pulls
80 if pulls0 < 750 { return __syscall(93, 21, 0, 0, 0, 0, 0) }
81 // Best-arm identification should pick arm 0
82 if nx_thompson_best_arm(bandit) != 0 { return __syscall(93, 22, 0, 0, 0, 0, 0) }
83
84 // ---- determinism ----
85 let a: *Thompson = nx_thompson_alloc(3, 42)
86 let b: *Thompson = nx_thompson_alloc(3, 42)
87 var i: i64 = 0
88 while i < 50 {
89 let sa: i64 = nx_thompson_select(a)
90 let sb: i64 = nx_thompson_select(b)
91 if sa != sb { return __syscall(93, 30, 0, 0, 0, 0, 0) }
92 nx_thompson_update(a, sa, i & 1)
93 nx_thompson_update(b, sb, i & 1)
94 i = i + 1
95 }
96
97 // ---- normal-approx path: artificially force s + f + 1 > THRESH ----
98 // Build a bandit with one arm pulled 100 times (50 succ, 50 fail).
99 let big: *Thompson = nx_thompson_alloc(2, 5)
100 var p: i64 = 0
101 while p < 50 {
102 nx_thompson_update(big, 0, 1)
103 nx_thompson_update(big, 0, 0)
104 p = p + 1
105 }
106 // s+f+1 = 101, well above THRESH=64 -> normal-approx path.
107 // Sample 100 times, average should approach 0.5 ± small.
108 var nsum: i64 = 0
109 var ni: i64 = 0
110 while ni < 100 {
111 nsum = nsum + nx_thompson_sample_q14(big, 0)
112 ni = ni + 1
113 }
114 let nmean: i64 = nsum / 100
115 // True μ ≈ 0.5 = 8192 in Q14. Allow +- 1500 (variance now ~1/√100 = 0.1).
116 if iabs(nmean - 8192) > 1500 {
117 return __syscall(93, 40, 0, 0, 0, 0, 0)
118 }
119
120 // ---- typed envelope ----
121 let q: *ApproxI64 = nx_thompson_query(big, 0)
122 if q.envelope_kind != NX_ENV_REL_STDDEV {
123 return __syscall(93, 50, 0, 0, 0, 0, 0)
124 }
125 if q.conf_ppb != 682700000 {
126 return __syscall(93, 51, 0, 0, 0, 0, 0)
127 }
128 if q.maturity != NX_MATURITY_REFERENCE_IMPL {
129 return __syscall(93, 52, 0, 0, 0, 0, 0)
130 }
131 if q.adv_safety != NX_ADV_HONEST {
132 return __syscall(93, 53, 0, 0, 0, 0, 0)
133 }
134 // Posterior mean: (50 + 1) / (50 + 50 + 2) = 51/102 = 0.5 → 500000 PPM.
135 if iabs(q.value - 500000) > 5000 {
136 return __syscall(93, 54, 0, 0, 0, 0, 0)
137 }
138
139 return 0
140}