sketch_correlation.nx source
↩ module page · 205 lines · 6617 B
1// sketch_correlation.nx -- streaming Pearson correlation + linear regression.
2//
3// Bivariate streaming primitive. For paired observations (x, y) consumed
4// one at a time, computes:
5// r -- Pearson correlation coefficient, in PPM in [-1_000_000, 1_000_000]
6// slope -- least-squares regression slope m: y = m*x + b
7// intercept -- least-squares regression intercept b
8//
9// SIX ACCUMULATORS (integer-exact under overflow budget):
10// n = count
11// sum_x = Σ x_i
12// sum_y = Σ y_i
13// sum_xy = Σ x_i * y_i
14// sum_xx = Σ x_i²
15// sum_yy = Σ y_i²
16//
17// FORMULAS:
18// cov_xy = n * sum_xy - sum_x * sum_y
19// var_x = n * sum_xx - sum_x²
20// var_y = n * sum_yy - sum_y²
21// r = cov_xy / sqrt(var_x * var_y)
22// slope = cov_xy / var_x
23// intercept = (sum_y - slope * sum_x) / n
24//
25// OVERFLOW BUDGET:
26// sum_xx grows as n * x_max². For x_max=2^15, n_max=2^32 -> sum_xx
27// max ~ 2^62. And the cross-term computations need extra margin.
28// nx_corr_safe_p flags unsafe inputs.
29//
30// USE CASES:
31// - SRE: correlate metric series (CPU vs latency)
32// - finance: returns correlation
33// - science: regression on streamed measurements
34//
35// LOSSLESS-LANGUAGE DISCIPLINE: r returned in PPM with NX_ENV_ABS,
36// param_a = 0 (exact under budget). Production tier.
37
38import "syscalls.nx"
39import "sketch_types.nx"
40
41const NX_CORR_VALUE_LIMIT: i64 = 32768 // |x|, |y| < 2^15 for safe sum_xx
42
43struct Correlation {
44 n: i64,
45 sum_x: i64,
46 sum_y: i64,
47 sum_xy: i64,
48 sum_xx: i64,
49 sum_yy: i64,
50 has_data: i64,
51}
52
53// === construction =================================================
54
55func nx_corr_alloc() -> *Correlation {
56 let raw: *u8 = sys_mmap(64)
57 let c: *Correlation = raw as *Correlation
58 c.n = 0
59 c.sum_x = 0
60 c.sum_y = 0
61 c.sum_xy = 0
62 c.sum_xx = 0
63 c.sum_yy = 0
64 c.has_data = 0
65 return c
66}
67
68// === isqrt =======================================================
69
70func nx_corr_isqrt(x: i64) -> i64 {
71 if x < 0 { return 0 }
72 if x == 0 { return 0 }
73 if x < 4 { return 1 }
74 var g: i64 = (x >> 1) + 1
75 var iter: i64 = 0
76 while iter < 64 {
77 let next_g: i64 = (g + x / g) / 2
78 if next_g >= g { iter = 64 }
79 if next_g < g {
80 g = next_g
81 iter = iter + 1
82 }
83 }
84 return g
85}
86
87// === overflow guard ==============================================
88
89func nx_corr_safe_p(x: i64, y: i64) -> i64 {
90 var ax: i64 = x
91 if ax < 0 { ax = -ax }
92 var ay: i64 = y
93 if ay < 0 { ay = -ay }
94 if ax >= NX_CORR_VALUE_LIMIT { return 0 }
95 if ay >= NX_CORR_VALUE_LIMIT { return 0 }
96 return 1
97}
98
99// === add ==========================================================
100
101func nx_corr_add(c: *Correlation, x: i64, y: i64) -> i64 {
102 if nx_corr_safe_p(x, y) == 0 { return -1 }
103 c.n = c.n + 1
104 c.sum_x = c.sum_x + x
105 c.sum_y = c.sum_y + y
106 c.sum_xy = c.sum_xy + x * y
107 c.sum_xx = c.sum_xx + x * x
108 c.sum_yy = c.sum_yy + y * y
109 c.has_data = 1
110 return 0
111}
112
113// === correlation coefficient (in PPM) ============================
114//
115// r = cov_xy / sqrt(var_x * var_y) where cov, var are the *centered*
116// values multiplied by n (so we avoid recomputing means).
117// Returns PPM in [-1_000_000, 1_000_000].
118
119func nx_corr_r_ppm(c: *Correlation) -> i64 {
120 if c.n < 2 { return 0 }
121 let cov_xy: i64 = c.n * c.sum_xy - c.sum_x * c.sum_y
122 let var_x: i64 = c.n * c.sum_xx - c.sum_x * c.sum_x
123 let var_y: i64 = c.n * c.sum_yy - c.sum_y * c.sum_y
124 if var_x <= 0 { return 0 }
125 if var_y <= 0 { return 0 }
126 // r = cov_xy / sqrt(var_x * var_y); compute in PPM.
127 // Avoid overflow in (var_x * var_y) by scaling cov upfront.
128 // r_ppm = (cov_xy * 1_000_000) / sqrt(var_x * var_y)
129 // Use staged isqrt: sqrt(var_x) * sqrt(var_y) (approximate but safe).
130 let sx: i64 = nx_corr_isqrt(var_x)
131 let sy: i64 = nx_corr_isqrt(var_y)
132 if sx == 0 { return 0 }
133 if sy == 0 { return 0 }
134 let denom: i64 = sx * sy
135 if denom == 0 { return 0 }
136 let r: i64 = (cov_xy * 1000000) / denom
137 // Clamp to [-1_000_000, 1_000_000].
138 if r > 1000000 { return 1000000 }
139 if r < -1000000 { return -1000000 }
140 return r
141}
142
143// === regression slope and intercept (in PPM) =====================
144//
145// slope_ppm = (n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x²) * 1_000_000
146// intercept = mean_y - slope * mean_x (slope here is fractional;
147// scaled-PPM math: intercept = (sum_y * 1_000_000 - slope_ppm * sum_x) / (n * 1_000_000))
148
149func nx_corr_slope_ppm(c: *Correlation) -> i64 {
150 if c.n < 2 { return 0 }
151 let cov_xy: i64 = c.n * c.sum_xy - c.sum_x * c.sum_y
152 let var_x: i64 = c.n * c.sum_xx - c.sum_x * c.sum_x
153 if var_x <= 0 { return 0 }
154 return (cov_xy * 1000000) / var_x
155}
156
157func nx_corr_intercept(c: *Correlation) -> i64 {
158 if c.n < 2 { return 0 }
159 let slope_ppm: i64 = nx_corr_slope_ppm(c)
160 // intercept = mean_y - slope * mean_x
161 // = (sum_y - slope_ppm * sum_x / 1_000_000) / n
162 let slope_times_sum_x: i64 = (slope_ppm * c.sum_x) / 1000000
163 return (c.sum_y - slope_times_sum_x) / c.n
164}
165
166// === typed queries ===============================================
167
168func nx_corr_query_r(c: *Correlation) -> *ApproxI64 {
169 let r: i64 = nx_corr_r_ppm(c)
170 return nx_approx_new(r, NX_ENV_ABS, 1,
171 1000000000,
172 NX_MATURITY_PRODUCTION,
173 NX_ADV_HONEST)
174}
175
176func nx_corr_query_slope(c: *Correlation) -> *ApproxI64 {
177 let s: i64 = nx_corr_slope_ppm(c)
178 return nx_approx_new(s, NX_ENV_ABS, 1,
179 1000000000,
180 NX_MATURITY_PRODUCTION,
181 NX_ADV_HONEST)
182}
183
184// === merge ========================================================
185
186func nx_corr_merge(a: *Correlation, b: *Correlation) -> *Correlation {
187 let out: *Correlation = nx_corr_alloc()
188 out.n = a.n + b.n
189 out.sum_x = a.sum_x + b.sum_x
190 out.sum_y = a.sum_y + b.sum_y
191 out.sum_xy = a.sum_xy + b.sum_xy
192 out.sum_xx = a.sum_xx + b.sum_xx
193 out.sum_yy = a.sum_yy + b.sum_yy
194 if a.has_data == 1 { out.has_data = 1 }
195 if b.has_data == 1 { out.has_data = 1 }
196 return out
197}
198
199func nx_corr_memory_bytes(c: *Correlation) -> i64 {
200 return 64
201}
202
203func nx_corr_count(c: *Correlation) -> i64 {
204 return c.n
205}