sketch_correlation.nx source
↩ module page · 191 lines · 6308 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"
40import "nx_vecmath.nx"
41
42const NX_CORR_VALUE_LIMIT: i64 = 32768 // |x|, |y| < 2^15 for safe sum_xx
43
44struct Correlation {
45 n: i64,
46 sum_x: i64,
47 sum_y: i64,
48 sum_xy: i64,
49 sum_xx: i64,
50 sum_yy: i64,
51 has_data: i64,
52}
53
54// === construction =================================================
55
56func nx_corr_alloc() -> *Correlation {
57 let raw: *u8 = sys_mmap(64)
58 let c: *Correlation = raw as *Correlation
59 c.n = 0
60 c.sum_x = 0
61 c.sum_y = 0
62 c.sum_xy = 0
63 c.sum_xx = 0
64 c.sum_yy = 0
65 c.has_data = 0
66 return c
67}
68
69// === isqrt =======================================================
70
71func nx_corr_isqrt(x: i64) -> i64 { return vm_isqrt(x) }
72
73// === overflow guard ==============================================
74
75func nx_corr_safe_p(x: i64, y: i64) -> i64 {
76 var ax: i64 = x
77 if ax < 0 { ax = -ax }
78 var ay: i64 = y
79 if ay < 0 { ay = -ay }
80 if ax >= NX_CORR_VALUE_LIMIT { return 0 }
81 if ay >= NX_CORR_VALUE_LIMIT { return 0 }
82 return 1
83}
84
85// === add ==========================================================
86
87func nx_corr_add(c: *Correlation, x: i64, y: i64) -> i64 {
88 if nx_corr_safe_p(x, y) == 0 { return -1 }
89 c.n = c.n + 1
90 c.sum_x = c.sum_x + x
91 c.sum_y = c.sum_y + y
92 c.sum_xy = c.sum_xy + x * y
93 c.sum_xx = c.sum_xx + x * x
94 c.sum_yy = c.sum_yy + y * y
95 c.has_data = 1
96 return 0
97}
98
99// === correlation coefficient (in PPM) ============================
100//
101// r = cov_xy / sqrt(var_x * var_y) where cov, var are the *centered*
102// values multiplied by n (so we avoid recomputing means).
103// Returns PPM in [-1_000_000, 1_000_000].
104
105func nx_corr_r_ppm(c: *Correlation) -> i64 {
106 if c.n < 2 { return 0 }
107 let cov_xy: i64 = c.n * c.sum_xy - c.sum_x * c.sum_y
108 let var_x: i64 = c.n * c.sum_xx - c.sum_x * c.sum_x
109 let var_y: i64 = c.n * c.sum_yy - c.sum_y * c.sum_y
110 if var_x <= 0 { return 0 }
111 if var_y <= 0 { return 0 }
112 // r = cov_xy / sqrt(var_x * var_y); compute in PPM.
113 // Avoid overflow in (var_x * var_y) by scaling cov upfront.
114 // r_ppm = (cov_xy * 1_000_000) / sqrt(var_x * var_y)
115 // Use staged isqrt: sqrt(var_x) * sqrt(var_y) (approximate but safe).
116 let sx: i64 = nx_corr_isqrt(var_x)
117 let sy: i64 = nx_corr_isqrt(var_y)
118 if sx == 0 { return 0 }
119 if sy == 0 { return 0 }
120 let denom: i64 = sx * sy
121 if denom == 0 { return 0 }
122 let r: i64 = (cov_xy * 1000000) / denom
123 // Clamp to [-1_000_000, 1_000_000].
124 if r > 1000000 { return 1000000 }
125 if r < -1000000 { return -1000000 }
126 return r
127}
128
129// === regression slope and intercept (in PPM) =====================
130//
131// slope_ppm = (n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x²) * 1_000_000
132// intercept = mean_y - slope * mean_x (slope here is fractional;
133// scaled-PPM math: intercept = (sum_y * 1_000_000 - slope_ppm * sum_x) / (n * 1_000_000))
134
135func nx_corr_slope_ppm(c: *Correlation) -> i64 {
136 if c.n < 2 { return 0 }
137 let cov_xy: i64 = c.n * c.sum_xy - c.sum_x * c.sum_y
138 let var_x: i64 = c.n * c.sum_xx - c.sum_x * c.sum_x
139 if var_x <= 0 { return 0 }
140 return (cov_xy * 1000000) / var_x
141}
142
143func nx_corr_intercept(c: *Correlation) -> i64 {
144 if c.n < 2 { return 0 }
145 let slope_ppm: i64 = nx_corr_slope_ppm(c)
146 // intercept = mean_y - slope * mean_x
147 // = (sum_y - slope_ppm * sum_x / 1_000_000) / n
148 let slope_times_sum_x: i64 = (slope_ppm * c.sum_x) / 1000000
149 return (c.sum_y - slope_times_sum_x) / c.n
150}
151
152// === typed queries ===============================================
153
154func nx_corr_query_r(c: *Correlation) -> *ApproxI64 {
155 let r: i64 = nx_corr_r_ppm(c)
156 return nx_approx_new(r, NX_ENV_ABS, 1,
157 1000000000,
158 NX_MATURITY_PRODUCTION,
159 NX_ADV_HONEST)
160}
161
162func nx_corr_query_slope(c: *Correlation) -> *ApproxI64 {
163 let s: i64 = nx_corr_slope_ppm(c)
164 return nx_approx_new(s, NX_ENV_ABS, 1,
165 1000000000,
166 NX_MATURITY_PRODUCTION,
167 NX_ADV_HONEST)
168}
169
170// === merge ========================================================
171
172func nx_corr_merge(a: *Correlation, b: *Correlation) -> *Correlation {
173 let out: *Correlation = nx_corr_alloc()
174 out.n = a.n + b.n
175 out.sum_x = a.sum_x + b.sum_x
176 out.sum_y = a.sum_y + b.sum_y
177 out.sum_xy = a.sum_xy + b.sum_xy
178 out.sum_xx = a.sum_xx + b.sum_xx
179 out.sum_yy = a.sum_yy + b.sum_yy
180 if a.has_data == 1 { out.has_data = 1 }
181 if b.has_data == 1 { out.has_data = 1 }
182 return out
183}
184
185func nx_corr_memory_bytes(c: *Correlation) -> i64 {
186 return 64
187}
188
189func nx_corr_count(c: *Correlation) -> i64 {
190 return c.n
191}