nx_root.nx source
↩ module page · 238 lines · 8278 B
1// nx_root.nx -- Newton-Raphson nth-root for i64 + Q14 fixed point.
2//
3// Bits-up math primitive sibling of nx_isqrt.nx. Computes floor of
4// the integer nth-root and the Q14-scaled fractional nth-root via
5// the same iteration shape. Unblocks two queued upgrades:
6//
7// * nx_sampler.nx Karras schedule (Karras et al. 2022 EDM): needs
8// (sigma_min / sigma_max)^(1/rho) with rho = 7.
9// * nx_ms_ssim.nx Wang weighted geometric mean (Wang et al. 2003):
10// needs fractional alpha^a * c^b * s^g powers per scale.
11//
12// ===== Math =======================================================
13//
14// Newton-Raphson iteration for x^n = T, x > 0:
15//
16// x_{k+1} = ((n-1) * x_k + T / x_k^(n-1)) / n
17//
18// Starting from x_0 = max(T, 1) (an upper bound), each iteration
19// strictly DECREASES x until convergence at floor(T^(1/n)).
20// Quadratic convergence near the root: O(log log n) iterations
21// after the initial linear phase, O(log n) total.
22//
23// Integer convergence test: terminate when x_{k+1} >= x_k.
24//
25// ===== Q14 wrapper ===============================================
26//
27// For fractional nth-root with target represented in Q14:
28//
29// nth_root_q14(target_q14, n) ≈ target^(1/n) in Q14
30//
31// We iterate directly in Q14 using a Q14-aware power helper
32// `nx_pow_q14` that does (x_q14 * x_q14 / Q14) chain. This keeps
33// intermediate values bounded near the target magnitude rather
34// than overflowing via a target * Q14^(n-1) scale-up.
35//
36// genealogy_id: heron_alexandria_iterative_root + newton_raphson_1671 +
37// wikipedia_nth_root_algorithm
38// lineage_id: substrate_root_v1
39
40// nx_safety_envelope:
41// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
42// sil_target: SIL1
43// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
44// verdict: NOT_YET_EVALUATED
45
46import "nx_syscalls.nx"
47import "nx_tier.nx"
48import "nx_loop.nx"
49
50// ===== Constants ==================================================
51
52const NX_ROOT_Q14_ONE: nx_int = 16384 // 1 << 14
53const NX_ROOT_MAX_ITERS: nx_int = 64 // Newton converges in <30 for i64 range
54
55// ===== Sealed-enum: RootVerdict ===================================
56
57const NX_ROOT_OK: nx_int = 0
58const NX_ROOT_ERR_NEG_TARGET: nx_int = 1
59const NX_ROOT_ERR_BAD_N: nx_int = 2
60const NX_ROOT_N_VERDICTS: nx_int = 3
61
62func nx_root_verdict_is_valid(v: nx_int) -> nx_int {
63 if v < 0 { return 0 }
64 if v >= NX_ROOT_N_VERDICTS { return 0 }
65 return 1
66}
67
68// ===== Integer power helper =======================================
69//
70// Returns base^exp for non-negative exp. No saturation -- caller
71// guarantees result fits in i64.
72
73func nx_root_ipow(base: nx_int, exp: nx_int) -> nx_int {
74 var result: nx_int = 1
75 var b: nx_int = base
76 var e: nx_int = exp
77 var iter: nx_int = 0
78 var verdict: nx_int = NX_LOOP_RUNNING
79 let BUDGET: nx_int = 64
80 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
81 if e <= 0 { verdict = NX_LOOP_DONE_EXIT }
82 if verdict == NX_LOOP_RUNNING {
83 let parity: nx_int = e - (e / 2) * 2
84 if parity == 1 { result = result * b }
85 b = b * b
86 e = e / 2
87 }
88 iter = iter + 1
89 }
90 return result
91}
92
93// ===== Integer nth-root: floor(target^(1/n)) ======================
94//
95// For target >= 0 and n >= 1, returns floor(target^(1/n)). Special
96// cases: n == 1 returns target; target == 0 returns 0.
97
98func nx_int_nth_root(target: nx_int, n: nx_int) -> nx_int {
99 if target <= 0 { return 0 }
100 if n <= 0 { return 0 }
101 if n == 1 { return target }
102
103 // Initial guess: target itself is an upper bound (target^(1/n) <= target for target >= 1, n >= 1).
104 var x: nx_int = target
105 if x < 1 { x = 1 }
106
107 var iter: nx_int = 0
108 var verdict: nx_int = NX_LOOP_RUNNING
109 let BUDGET: nx_int = NX_ROOT_MAX_ITERS
110 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
111 let x_pow_nm1: nx_int = nx_root_ipow(x, n - 1)
112 if x_pow_nm1 <= 0 { verdict = NX_LOOP_DONE_EXIT }
113 if verdict == NX_LOOP_RUNNING {
114 let div_part: nx_int = target / x_pow_nm1
115 let new_x: nx_int = ((n - 1) * x + div_part) / n
116 if new_x >= x { verdict = NX_LOOP_DONE_EXIT }
117 if verdict == NX_LOOP_RUNNING { x = new_x }
118 }
119 iter = iter + 1
120 }
121 return x
122}
123
124// ===== Q14 nth-root: target_q14^(1/n) in Q14 =====================
125//
126// Returns x such that x_q14^n / Q14^(n-1) ≈ target_q14.
127// Newton-Raphson iterated in Q14:
128//
129// x_pow_nm1_q14 = x_q14^(n-1) / Q14^(n-2) (in Q14 scale)
130// div_part_q14 = target_q14 * Q14 / x_pow_nm1_q14
131// x_new_q14 = ((n-1) * x_q14 + div_part_q14) / n
132//
133// For numerical stability on small targets (target_q14 < Q14),
134// we start at x_0 = max(target_q14, Q14). Newton then descends.
135
136func nx_pow_q14(x_q14: nx_int, exp: nx_int) -> nx_int {
137 var result: nx_int = NX_ROOT_Q14_ONE
138 var iter: nx_int = 0
139 var verdict: nx_int = NX_LOOP_RUNNING
140 let BUDGET: nx_int = exp
141 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
142 result = (result * x_q14) / NX_ROOT_Q14_ONE
143 iter = iter + 1
144 }
145 return result
146}
147
148func nx_nth_root_q14(target_q14: nx_int, n: nx_int) -> nx_int {
149 if target_q14 <= 0 { return 0 }
150 if n <= 0 { return 0 }
151 if n == 1 { return target_q14 }
152
153 var x: nx_int = target_q14
154 if x < NX_ROOT_Q14_ONE { x = NX_ROOT_Q14_ONE }
155
156 var iter: nx_int = 0
157 var verdict: nx_int = NX_LOOP_RUNNING
158 let BUDGET: nx_int = NX_ROOT_MAX_ITERS
159 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
160 let x_pow_nm1_q14: nx_int = nx_pow_q14(x, n - 1)
161 if x_pow_nm1_q14 <= 0 { verdict = NX_LOOP_DONE_EXIT }
162 if verdict == NX_LOOP_RUNNING {
163 // div_part_q14 = target_q14 * Q14 / x_pow_nm1_q14
164 let div_part: nx_int = (target_q14 * NX_ROOT_Q14_ONE) / x_pow_nm1_q14
165 let new_x: nx_int = ((n - 1) * x + div_part) / n
166 // Convergence: x stopped descending.
167 if new_x >= x { verdict = NX_LOOP_DONE_EXIT }
168 if verdict == NX_LOOP_RUNNING { x = new_x }
169 }
170 iter = iter + 1
171 }
172 return x
173}
174
175// ===== Self-test ==================================================
176//
177// Closed-form invariants:
178//
179// Integer:
180// (a) cbrt(64) == 4
181// (b) cbrt(27) == 3
182// (c) cbrt(8) == 2
183// (d) 10th-root(1024) == 2
184// (e) 7th-root(128) == 2 (128 = 2^7)
185// (f) 5th-root(1000) in [3, 4] (3^5=243, 4^5=1024)
186//
187// Q14:
188// (g) nth_root_q14(8 * Q14, 3) ≈ 2 * Q14 (within ±1 Q14 unit)
189// (h) nth_root_q14(Q14, n) ≈ Q14 (root of unity)
190// (i) nth_root_q14(128 * Q14, 7) ≈ 2 * Q14
191// (j) ipow check: nx_root_ipow(2, 10) == 1024
192
193func main() -> i64 {
194 // --- Integer ---
195 if nx_int_nth_root(64, 2) != 8 { return 10 }
196 if nx_int_nth_root(64, 3) != 4 { return 11 }
197 if nx_int_nth_root(27, 3) != 3 { return 12 }
198 if nx_int_nth_root(8, 3) != 2 { return 13 }
199 if nx_int_nth_root(1024, 10) != 2 { return 14 }
200 if nx_int_nth_root(128, 7) != 2 { return 15 }
201 let r5: nx_int = nx_int_nth_root(1000, 5)
202 if r5 < 3 { return 16 }
203 if r5 > 4 { return 17 }
204
205 // --- ipow ---
206 if nx_root_ipow(2, 10) != 1024 { return 20 }
207 if nx_root_ipow(3, 5) != 243 { return 21 }
208 if nx_root_ipow(7, 0) != 1 { return 22 }
209
210 // --- Q14 ---
211 let q14: nx_int = NX_ROOT_Q14_ONE
212 // root of unity == unity
213 let r_one: nx_int = nx_nth_root_q14(q14, 5)
214 let drift_one: nx_int = r_one - q14
215 if drift_one > 2 { return 30 }
216 if drift_one < -2 { return 31 }
217
218 // cbrt(8) in Q14 ≈ 2 * Q14
219 let r_cbrt_8: nx_int = nx_nth_root_q14(8 * q14, 3)
220 let drift_cbrt: nx_int = r_cbrt_8 - 2 * q14
221 if drift_cbrt > 2 { return 40 }
222 if drift_cbrt < -2 { return 41 }
223
224 // 7th-root(128) in Q14 ≈ 2 * Q14
225 let r_7th: nx_int = nx_nth_root_q14(128 * q14, 7)
226 let drift_7th: nx_int = r_7th - 2 * q14
227 if drift_7th > 4 { return 50 }
228 if drift_7th < -4 { return 51 }
229
230 // --- Verdict gate ---
231 var v: nx_int = 0
232 while v < NX_ROOT_N_VERDICTS {
233 if nx_root_verdict_is_valid(v) != 1 { return 60 + v }
234 v = v + 1
235 }
236
237 return 0
238}