nx_word_sketch_test.nx source

1// nx_word_sketch_test.nx -- exercise positional word sketch primitive. 2 3import "nx_syscalls.nx" 4import "nx_tier.nx" 5import "nx_ngram.nx" 6import "nx_collocation.nx" 7import "nx_word_sketch.nx" 8 9func main() -> nx_int { 10 // ===== Tiny corpus: "the cat sat on the mat the cat ran fast" ===== 11 // 12 // Interned IDs: the=1, cat=2, sat=3, on=4, mat=5, ran=6, fast=7. 13 // Length = 10 tokens. target=cat (id=2) appears 2 times. 14 15 let stream: *i64 = (sys_mmap(80)) as *i64 16 stream[0] = 1 // the 17 stream[1] = 2 // cat 18 stream[2] = 3 // sat 19 stream[3] = 4 // on 20 stream[4] = 1 // the 21 stream[5] = 5 // mat 22 stream[6] = 1 // the 23 stream[7] = 2 // cat 24 stream[8] = 6 // ran 25 stream[9] = 7 // fast 26 27 let cap: nx_int = 32 28 let keys: *i64 = (sys_mmap(cap * 8)) as *i64 29 let counts: *i64 = (sys_mmap(cap * 8)) as *i64 30 31 // ===== Sketch "cat" with window=1 ================================= 32 // 33 // Occurrences of cat at i=1 and i=7. 34 // window=1 neighbors: 35 // i=1 -> stream[0]=the(1), stream[2]=sat(3) 36 // i=7 -> stream[6]=the(1), stream[8]=ran(6) 37 // Expected counts: the=2, sat=1, ran=1. 38 39 let n_target_1: nx_int = nx_word_sketch_count(stream, 10, 2, 1, 40 keys, counts, cap) 41 if n_target_1 != 2 { return 1 } 42 43 // Walk the table; verify the expected collocates + counts 44 var n_the: nx_int = 0 45 var n_sat: nx_int = 0 46 var n_ran: nx_int = 0 47 var n_other: nx_int = 0 48 var i: nx_int = 0 49 while i < cap { 50 if keys[i] != NX_WS_EMPTY_KEY { 51 if keys[i] == 1 { n_the = counts[i] } 52 if keys[i] == 3 { n_sat = counts[i] } 53 if keys[i] == 6 { n_ran = counts[i] } 54 if keys[i] != 1 { 55 if keys[i] != 3 { 56 if keys[i] != 6 { 57 n_other = n_other + 1 58 } 59 } 60 } 61 } 62 i = i + 1 63 } 64 if n_the != 2 { return 2 } 65 if n_sat != 1 { return 3 } 66 if n_ran != 1 { return 4 } 67 if n_other != 0 { return 5 } 68 69 // ===== Sketch "cat" with window=2 ================================= 70 // 71 // window=2 neighbors: 72 // i=1 -> stream[0]=the(1) [skip -2 since i-2=-1 OOB], stream[2]=sat(3), stream[3]=on(4) 73 // (offset -2 is OOB; offsets -1, +1, +2 valid) 74 // i=7 -> stream[5]=mat(5), stream[6]=the(1), stream[8]=ran(6), stream[9]=fast(7) 75 // Expected: the=2, sat=1, on=1, mat=1, ran=1, fast=1 76 77 let n_target_2: nx_int = nx_word_sketch_count(stream, 10, 2, 2, 78 keys, counts, cap) 79 if n_target_2 != 2 { return 10 } 80 81 var t_the: nx_int = 0 82 var t_sat: nx_int = 0 83 var t_on: nx_int = 0 84 var t_mat: nx_int = 0 85 var t_ran: nx_int = 0 86 var t_fast: nx_int = 0 87 var distinct: nx_int = 0 88 var j: nx_int = 0 89 while j < cap { 90 if keys[j] != NX_WS_EMPTY_KEY { 91 distinct = distinct + 1 92 if keys[j] == 1 { t_the = counts[j] } 93 if keys[j] == 3 { t_sat = counts[j] } 94 if keys[j] == 4 { t_on = counts[j] } 95 if keys[j] == 5 { t_mat = counts[j] } 96 if keys[j] == 6 { t_ran = counts[j] } 97 if keys[j] == 7 { t_fast = counts[j] } 98 } 99 j = j + 1 100 } 101 if t_the != 2 { return 11 } 102 if t_sat != 1 { return 12 } 103 if t_on != 1 { return 13 } 104 if t_mat != 1 { return 14 } 105 if t_ran != 1 { return 15 } 106 if t_fast != 1 { return 16 } 107 if distinct != 6 { return 17 } 108 109 // ===== Sketch verdict bands ======================================= 110 if nx_ws_classify(0) != NX_WS_KIND_EMPTY { return 20 } 111 if nx_ws_classify(2) != NX_WS_KIND_SPARSE { return 21 } 112 if nx_ws_classify(10) != NX_WS_KIND_HEALTHY { return 22 } 113 if nx_ws_classify(500) != NX_WS_KIND_DENSE { return 23 } 114 115 // ===== Score a single collocate =================================== 116 // 117 // Assume corpus stats: n_target=2 (cat occurrences), n_the=3 (the 118 // total in corpus), n_cooccur=2, n_total=10. Score with log-Dice. 119 let dice: nx_int = nx_word_sketch_metric_q10(2, 3, 2, 10, 120 NX_WS_METRIC_LOG_DICE) 121 // log-Dice = 14 + log2(2*2 / (2+3)) = 14 + log2(4/5) ~ 14 - 0.32 ~ 13.68 122 // Q10 -> ~ 14008 123 // We sanity-check a wide band (>0 and < 20*1024 = 20480). 124 if dice <= 0 { return 30 } 125 if dice >= 20480 { return 31 } 126 127 // Score "weak" collocate: high target freq + low coccur -> small score 128 let dice_weak: nx_int = nx_word_sketch_metric_q10(100, 1000, 2, 100000, 129 NX_WS_METRIC_LOG_DICE) 130 // weak should be < strong 131 if dice_weak >= dice { return 32 } 132 133 // ===== Full-metric emit returns all 4 metrics + 4 bands ========== 134 let out: *i64 = (sys_mmap(NX_COL_OUT_FIELDS * 8)) as *i64 135 nx_word_sketch_score_one(2, 3, 2, 10, out) 136 // pmi_band must be valid 137 if nx_col_band_is_valid(out[NX_COL_OUT_PMI_BAND]) == 0 { return 40 } 138 if nx_col_band_is_valid(out[NX_COL_OUT_LL_BAND]) == 0 { return 41 } 139 if nx_col_band_is_valid(out[NX_COL_OUT_T_BAND]) == 0 { return 42 } 140 if nx_col_band_is_valid(out[NX_COL_OUT_DICE_BAND]) == 0 { return 43 } 141 142 // ===== Target absent from stream ================================== 143 let n_target_3: nx_int = nx_word_sketch_count(stream, 10, 99, 2, 144 keys, counts, cap) 145 if n_target_3 != 0 { return 50 } 146 if nx_ws_classify(n_target_3) != NX_WS_KIND_EMPTY { return 51 } 147 148 return 0 149}