gam_math/jet_tower.rs
1//! Taylor-jet tower algebra: write each family's row log-likelihood ONCE,
2//! derive the entire `RowKernel<K>` derivative tower mechanically (#932).
3//!
4//! # The object
5//!
6//! [`Tower4<K>`] is a truncated multivariate Taylor scalar in `K` primary
7//! variables, carrying the value and ALL partial derivatives through fourth
8//! order as full (unsymmetrized) tensors:
9//!
10//! ```text
11//! v ℓ
12//! g[a] ∂ℓ/∂p_a
13//! h[a][b] ∂²ℓ/∂p_a∂p_b
14//! t3[abc] ∂³ℓ/∂p_a∂p_b∂p_c
15//! t4[abcd] ∂⁴ℓ/∂p_a∂p_b∂p_c∂p_d
16//! ```
17//!
18//! Arithmetic (`+ − × ÷`, scalar mixes) propagates the tower by the exact
19//! Leibniz rule; unary transcendentals propagate by the exact multivariate
20//! Faà di Bruno formula given a `[f, f′, f″, f‴, f⁗]` stack evaluated at the
21//! inner value. This is truncated Taylor ALGEBRA — exact derivatives of the
22//! evaluated expression, not finite differences, not an approximation —
23//! fully compatible with the exact-REML-only policy.
24//!
25//! One evaluation of a row NLL program at seeded variables yields, in a
26//! single pass, every channel the `gam_models::row_kernel::RowKernel` trait
27//! demands: `row_kernel` (value/∇/H), `row_third_contracted(dir)` (contract
28//! `t3` with `dir`), and `row_fourth_contracted(u, v)` (contract `t4` with
29//! `u` and `v`). The directional cross-channels that hand-written towers
30//! drop (#736's residual gap) cannot be dropped here: there is no separate
31//! "channel" to forget — every derivative of the one expression is carried.
32//!
33//! # Why this exists (the bug genus)
34//!
35//! Every family today hand-writes its tower: value in one function,
36//! gradient in another, `pdfthird_derivative`/`pdffourth_derivative`,
37//! entry/exit-specific cross blocks — thousands of lines of calculus that
38//! drift. #736 was a sign flip in a hand-written cross-Hessian block,
39//! invisible until a new consumer touched it; #948 is a derivative path
40//! that is not the derivative of the evaluated row loss (clamped-μ
41//! surrogate); the objective↔gradient desync class is the same disease at
42//! the criterion level. A tower-derived kernel is exact-by-construction:
43//! the value channel IS the production loss expression, so its derivative
44//! channels cannot desync from it.
45//!
46//! # Relation to `jet_partitions::MultiDirJet`
47//!
48//! The tree already carries a *directional* jet (bitmask coefficients over
49//! distinct seeded directions, heap-allocated, Bell-partition compose) used
50//! inside the marginal-slope and latent-survival families. It answers "the
51//! derivative along THESE specific directions" and must be re-seeded and
52//! re-evaluated per direction tuple (e.g. 10 symmetric `(a,b)` pairs for a
53//! K=4 fourth contraction). `Tower4` answers ALL of them from one
54//! evaluation: contraction happens AFTER differentiation, as plain linear
55//! algebra on the stored tensors. Use `MultiDirJet` when you need a handful
56//! of directions of a huge-K expression; use `Tower4` when you need the
57//! complete small-K tower — which is exactly the `RowKernel<K≤4>` shape.
58//! The `[f64; 5]` unary-derivative stacks
59//! (`unary_derivatives_neglog_phi`, …) are signature-compatible with
60//! [`Tower4::compose_unary`], so the families' existing special-function
61//! stacks are directly reusable.
62//!
63//! # Stability discipline (why this is NOT autodiff)
64//!
65//! Differentiating the primal code path inherits its instabilities: a jet
66//! pushed through a naive `ln(1 + e^η)` is garbage in the saturated tail
67//! even though the true derivative σ(η) is benign there. This module
68//! therefore splits responsibility: **humans own primitive stability,
69//! the algebra owns combinatorics**. Tail-critical special functions enter
70//! a program ONLY as hand-certified `[f64; 5]` derivative stacks through
71//! [`Tower4::compose_unary`] — the same stacks the families already write
72//! (`unary_derivatives_neglog_phi` and friends, built on erfcx/log_ndtr) —
73//! and the tower mechanizes only the Leibniz/Faà di Bruno composition,
74//! which is where hand-written towers actually fail (#736 was a
75//! composition sign flip, not a primitive error). Program authors must use
76//! a stable primitive stack wherever the f64 production loss does; the
77//! convenience methods (`exp`, `ln`, `sqrt`, …) are for expressions whose
78//! arguments are tame by construction.
79//!
80//! # Storage convention
81//!
82//! Tensors are stored FULL, not symmetric-packed: `t4` for K=4 is 256
83//! doubles where 35 would do. This is deliberate clarity-over-speed for the
84//! oracle role — indexing is trivially auditable, contraction loops are
85//! obvious, and the redundancy is itself a checked invariant (the algebra
86//! only ever writes symmetric values). Symmetric packing is a later,
87//! profile-justified optimization behind the same API.
88//!
89//! # Deployment ladder (#932)
90//!
91//! 1. This module: the algebra + the program seam + the oracle.
92//! 2. Universal oracle: every hand-written `RowKernel` gains a CI test
93//! asserting channel-by-channel agreement with a [`RowProgram`] written
94//! once — see `verify_kernel_channels`. This alone would have caught
95//! #736 at introduction.
96//! 3. Derive every channel through [`program_row_kernel`],
97//! [`program_third_contracted`], [`program_fourth_contracted`], or
98//! [`program_full_tower`], selecting only the representation its consumer
99//! needs while retaining one expression.
100//! 4. New families (#914/#916/#917 ZI/ordinal/expectile, #921's location-
101//! scale port) implement ONLY [`RowProgram`] and get an exact fourth-order
102//! tower for the price of writing the likelihood.
103
104use crate::jet_algebra;
105
106/// Truncated fourth-order multivariate Taylor scalar in `K` variables.
107///
108/// See the module documentation for semantics and conventions. `Copy` is
109/// intentional despite the size (2 KiB at K=4): towers are per-row
110/// temporaries that live entirely in registers/stack during a row program,
111/// and value semantics keep program code readable (`a * b + c`).
112#[derive(Clone, Copy, Debug)]
113pub struct Tower4<const K: usize> {
114 /// Value ℓ.
115 pub v: f64,
116 /// Gradient ∂ℓ/∂p_a.
117 pub g: [f64; K],
118 /// Hessian ∂²ℓ/∂p_a∂p_b (symmetric).
119 pub h: [[f64; K]; K],
120 /// Third derivatives ∂³ℓ/∂p_a∂p_b∂p_c (fully symmetric).
121 pub t3: [[[f64; K]; K]; K],
122 /// Fourth derivatives ∂⁴ℓ/∂p_a∂p_b∂p_c∂p_d (fully symmetric).
123 pub t4: [[[[f64; K]; K]; K]; K],
124}
125
126impl<const K: usize> Tower4<K> {
127 /// The additive identity.
128 pub fn zero() -> Self {
129 Self {
130 v: 0.0,
131 g: [0.0; K],
132 h: [[0.0; K]; K],
133 t3: [[[0.0; K]; K]; K],
134 t4: [[[[0.0; K]; K]; K]; K],
135 }
136 }
137
138 /// A constant: value `c`, all derivatives zero.
139 pub fn constant(c: f64) -> Self {
140 let mut out = Self::zero();
141 out.v = c;
142 out
143 }
144
145 /// The seeded variable `p_idx` with current value `value`:
146 /// unit first derivative in slot `idx`, zero elsewhere and above.
147 pub fn variable(value: f64, idx: usize) -> Self {
148 let mut out = Self::constant(value);
149 out.g[idx] = 1.0;
150 out
151 }
152
153 /// Read the (fully symmetric) derivative tensor entry whose differentiation
154 /// axes are `labels` (length 0..=4): value, `g`, `h`, `t3`, `t4`.
155 #[inline]
156 fn deriv(&self, labels: &[usize]) -> f64 {
157 assert!(
158 labels.len() <= 4,
159 "Tower4 carries at most fourth-order derivatives"
160 );
161 match labels.len() {
162 0 => self.v,
163 1 => self.g[labels[0]],
164 2 => self.h[labels[0]][labels[1]],
165 3 => self.t3[labels[0]][labels[1]][labels[2]],
166 _ => self.t4[labels[0]][labels[1]][labels[2]][labels[3]],
167 }
168 }
169
170 /// Exact truncated Leibniz product `D_S(ab) = Σ_{T ⊆ S} D_T(a) · D_{S∖T}(b)`.
171 ///
172 /// # Codegen
173 ///
174 /// Each output entry's `2^m` subset sum is written as a compact straight-line
175 /// expression instead of the shared `jet_algebra::leibniz_product` subset
176 /// walker (which, per entry, builds `SlotBuf`s and `match`-dispatches the
177 /// `deriv` closure across all `2^m` subsets). The loop nest over `(i,j,k,l)`
178 /// is unchanged — only the inner per-entry sum is unrolled — so this does NOT
179 /// unroll over `K` and does NOT bloat code: on a `Tower4<9>` mul-and-read
180 /// consumer the new form is faster AND smaller (asm: 34 outlined walker `bl`
181 /// calls → 0, 21.1 KiB → 14.3 KiB, +100 NEON `.2d` ops).
182 ///
183 /// BIT-IDENTICAL to the walker: each entry's terms are in the walker's exact
184 /// subset-enumeration order (subset bit `b` ↔ position `b`, `sub = 0..2^m`),
185 /// and the per-entry `acc` accumulator mirrors the walker's `total = 0.0`
186 /// start so a signed-zero leading product collapses to `+0.0` identically —
187 /// which matters because real jets carry exact-`0.0` channels
188 /// (`constant`/`variable` towers). Proven `to_bits`-identical on
189 /// `v`/`g`/`h`/`t3`/`t4` across `K ∈ {2,3,4,9}`, 5000 inputs each with ~30 %
190 /// exact-`0.0` channels and signed values (a no-leading-`0.0` form fails this
191 /// stress — the accumulator start is load-bearing).
192 pub fn mul(&self, o: &Self) -> Self {
193 let a = self;
194 let b = o;
195 let mut out = Self::zero();
196 out.v = a.v * b.v;
197 for i in 0..K {
198 // subsets of {i}: {} {i}
199 let mut acc = 0.0;
200 acc += a.v * b.g[i];
201 acc += a.g[i] * b.v;
202 out.g[i] = acc;
203 }
204 // Hessian is symmetric under i↔j; compute the upper triangle and mirror
205 // (see [`Tower2::mul`] — same term order, enforces exact symmetry).
206 for i in 0..K {
207 for j in i..K {
208 // subsets of {i,j}: {} {i} {j} {ij}
209 let mut acc = 0.0;
210 acc += a.v * b.h[i][j];
211 acc += a.g[i] * b.g[j];
212 acc += a.g[j] * b.g[i];
213 acc += a.h[i][j] * b.v;
214 out.h[i][j] = acc;
215 out.h[j][i] = acc;
216 }
217 }
218 for i in 0..K {
219 for j in 0..K {
220 for k in 0..K {
221 // subsets of {i,j,k}: {} {i} {j} {ij} {k} {ik} {jk} {ijk}
222 let mut acc = 0.0;
223 acc += a.v * b.t3[i][j][k];
224 acc += a.g[i] * b.h[j][k];
225 acc += a.g[j] * b.h[i][k];
226 acc += a.h[i][j] * b.g[k];
227 acc += a.g[k] * b.h[i][j];
228 acc += a.h[i][k] * b.g[j];
229 acc += a.h[j][k] * b.g[i];
230 acc += a.t3[i][j][k] * b.v;
231 out.t3[i][j][k] = acc;
232 }
233 }
234 }
235 for i in 0..K {
236 for j in 0..K {
237 for k in 0..K {
238 for l in 0..K {
239 // subsets of {i,j,k,l} in bit order sub = 0..16
240 let mut acc = 0.0;
241 acc += a.v * b.t4[i][j][k][l];
242 acc += a.g[i] * b.t3[j][k][l];
243 acc += a.g[j] * b.t3[i][k][l];
244 acc += a.h[i][j] * b.h[k][l];
245 acc += a.g[k] * b.t3[i][j][l];
246 acc += a.h[i][k] * b.h[j][l];
247 acc += a.h[j][k] * b.h[i][l];
248 acc += a.t3[i][j][k] * b.g[l];
249 acc += a.g[l] * b.t3[i][j][k];
250 acc += a.h[i][l] * b.h[j][k];
251 acc += a.h[j][l] * b.h[i][k];
252 acc += a.t3[i][j][l] * b.g[k];
253 acc += a.h[k][l] * b.h[i][j];
254 acc += a.t3[i][k][l] * b.g[j];
255 acc += a.t3[j][k][l] * b.g[i];
256 acc += a.t4[i][j][k][l] * b.v;
257 out.t4[i][j][k][l] = acc;
258 }
259 }
260 }
261 }
262 out
263 }
264
265 /// Ref-taking elementwise sum, the by-ref twin of the `std::ops::Add`
266 /// operator (which consumes by value). Mirrors the inherent `mul`/`scale`
267 /// API so a chain like `a.mul(&b).add(&c)` reads uniformly without moving
268 /// out of the borrowed operands.
269 pub fn add(&self, o: &Self) -> Self {
270 *self + *o
271 }
272
273 /// Ref-taking elementwise difference, the by-ref twin of `std::ops::Sub`.
274 pub fn sub(&self, o: &Self) -> Self {
275 *self + o.scale(-1.0)
276 }
277
278 /// Exact multivariate Faà di Bruno composition `f ∘ self`.
279 ///
280 /// `d = [f(u), f′(u), f″(u), f‴(u), f⁗(u)]` evaluated at `u = self.v` —
281 /// the SAME `[f64; 5]` stack shape the families' existing
282 /// `unary_derivatives_*` helpers produce, so those special-function
283 /// stacks (Φ, log-Φ, normal pdf, …) plug in directly.
284 ///
285 /// The order-m output sums over the set partitions of the m indices
286 /// (Bell(3) = 5 terms at order 3, Bell(4) = 15 at order 4), grouped by
287 /// block count: each partition into r blocks contributes
288 /// `f⁽ʳ⁾ · Π_blocks D_block(u)`.
289 ///
290 /// # Codegen
291 ///
292 /// Evaluated as a compact closed form (the Bell(4)=15 set-partitions of
293 /// `t4`, Bell(3)=5 of `t3`, …) instead of routing through the recursive
294 /// [`jet_algebra::faa_di_bruno`] walker (per-output `for_each_partition`
295 /// recursion + per-block `SlotBuf` + closure dispatch). The loop nest is
296 /// identical to the walker's (`for i,j,k,l`); only the per-entry partition
297 /// sum is straight-line, so this does NOT unroll over `K` and does NOT
298 /// bloat code — measured on a `Tower4<9>` compose-and-read consumer the new
299 /// form is both faster and SMALLER (asm: 94 outlined walker `bl` calls → 0,
300 /// 47.5 KiB → 16.7 KiB, +197 NEON `.2d` ops).
301 ///
302 /// BIT-IDENTICAL to the walker: each channel's terms are emitted in the
303 /// walker's exact partition-enumeration order, each term's block products
304 /// are left-associated exactly as the walker's `prod *= block`, and the
305 /// per-channel `acc` accumulator mirrors the walker's `total = 0.0` start
306 /// (so signed-zero products collapse to `+0.0` identically). The order-4
307 /// term sequence was generated from the walker's own enumeration. Proven
308 /// `to_bits`-identical on `v`/`g`/`h`/`t3`/`t4` across `K ∈ {2,3,4,9}`,
309 /// 5000 random inputs each (zeroed / sign-varied stacks included).
310 pub fn compose_unary(&self, d: [f64; 5]) -> Self {
311 let mut out = Self::zero();
312 out.v = d[0];
313 for i in 0..K {
314 let mut acc = 0.0;
315 acc += d[1] * self.g[i];
316 out.g[i] = acc;
317 }
318 for i in 0..K {
319 for j in 0..K {
320 let mut acc = 0.0;
321 acc += d[1] * self.h[i][j];
322 acc += d[2] * self.g[i] * self.g[j];
323 out.h[i][j] = acc;
324 }
325 }
326 for i in 0..K {
327 for j in 0..K {
328 for k in 0..K {
329 // walker partitions: {ijk} {ij}{k} {ik}{j} {i}{jk} {i}{j}{k}
330 let mut acc = 0.0;
331 acc += d[1] * self.t3[i][j][k];
332 acc += d[2] * self.h[i][j] * self.g[k];
333 acc += d[2] * self.h[i][k] * self.g[j];
334 acc += d[2] * self.g[i] * self.h[j][k];
335 acc += d[3] * self.g[i] * self.g[j] * self.g[k];
336 out.t3[i][j][k] = acc;
337 }
338 }
339 }
340 for i in 0..K {
341 for j in 0..K {
342 for k in 0..K {
343 for l in 0..K {
344 // Bell(4)=15 partitions, walker enumeration order.
345 let mut acc = 0.0;
346 acc += d[1] * self.t4[i][j][k][l];
347 acc += d[2] * self.t3[i][j][k] * self.g[l];
348 acc += d[2] * self.t3[i][j][l] * self.g[k];
349 acc += d[2] * self.h[i][j] * self.h[k][l];
350 acc += d[3] * self.h[i][j] * self.g[k] * self.g[l];
351 acc += d[2] * self.t3[i][k][l] * self.g[j];
352 acc += d[2] * self.h[i][k] * self.h[j][l];
353 acc += d[3] * self.h[i][k] * self.g[j] * self.g[l];
354 acc += d[2] * self.h[i][l] * self.h[j][k];
355 acc += d[2] * self.g[i] * self.t3[j][k][l];
356 acc += d[3] * self.g[i] * self.h[j][k] * self.g[l];
357 acc += d[3] * self.h[i][l] * self.g[j] * self.g[k];
358 acc += d[3] * self.g[i] * self.h[j][l] * self.g[k];
359 acc += d[3] * self.g[i] * self.g[j] * self.h[k][l];
360 acc += d[4] * self.g[i] * self.g[j] * self.g[k] * self.g[l];
361 out.t4[i][j][k][l] = acc;
362 }
363 }
364 }
365 }
366 out
367 }
368
369 /// Multiply every channel by a plain scalar.
370 pub fn scale(&self, s: f64) -> Self {
371 let mut out = *self;
372 out.v *= s;
373 for i in 0..K {
374 out.g[i] *= s;
375 for j in 0..K {
376 out.h[i][j] *= s;
377 for k in 0..K {
378 out.t3[i][j][k] *= s;
379 for l in 0..K {
380 out.t4[i][j][k][l] *= s;
381 }
382 }
383 }
384 }
385 out
386 }
387
388 /// e^self.
389 pub fn exp(&self) -> Self {
390 let e = self.v.exp();
391 self.compose_unary([e, e, e, e, e])
392 }
393
394 /// ln(self). Caller guarantees positivity (likelihood programs do).
395 pub fn ln(&self) -> Self {
396 let u = self.v;
397 let r = 1.0 / u;
398 self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
399 }
400
401 /// 1/self.
402 pub fn recip(&self) -> Self {
403 let r = 1.0 / self.v;
404 let r2 = r * r;
405 self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
406 }
407
408 /// √self. Caller guarantees positivity.
409 pub fn sqrt(&self) -> Self {
410 let u = self.v;
411 let s = u.sqrt();
412 self.compose_unary([
413 s,
414 0.5 / s,
415 -0.25 / (u * s),
416 0.375 / (u * u * s),
417 -0.9375 / (u * u * u * s),
418 ])
419 }
420
421 /// self^a for real exponent `a`. Caller guarantees a positive base.
422 pub fn powf(&self, a: f64) -> Self {
423 let u = self.v;
424 let f0 = u.powf(a);
425 let f1 = a * u.powf(a - 1.0);
426 let f2 = a * (a - 1.0) * u.powf(a - 2.0);
427 let f3 = a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0);
428 let f4 = a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0);
429 self.compose_unary([f0, f1, f2, f3, f4])
430 }
431
432 /// ln Γ(self). Caller guarantees positivity.
433 pub fn ln_gamma(&self) -> Self {
434 self.compose_unary(ln_gamma_derivative_stack(self.v))
435 }
436
437 /// Contract `t3` with one primary-space direction:
438 /// `out[a][b] = Σ_c t3[a][b][c] · dir[c]` — exactly the
439 /// `row_third_contracted` shape.
440 ///
441 /// The output is symmetric in `(a, b)`: `t3` is fully index-symmetric, so
442 /// `t3[a][b][c] == t3[b][a][c]` and the `Σ_c` contraction gives
443 /// `out[a][b] == out[b][a]` term-for-term, in the same `c` order. We compute
444 /// only the upper triangle `a ≤ b` (the inner contraction is unchanged and
445 /// stays contiguous/vectorisable) and mirror into the lower triangle — this
446 /// is BIT-IDENTICAL to the full `a, b ∈ 0..K` nest while doing ~2× fewer
447 /// inner contractions, with no dense scatter (the mirror is a `K × K` copy).
448 pub fn third_contracted(&self, dir: &[f64; K]) -> [[f64; K]; K] {
449 let mut out = [[0.0; K]; K];
450 for a in 0..K {
451 for b in a..K {
452 let mut acc = 0.0;
453 for c in 0..K {
454 acc += self.t3[a][b][c] * dir[c];
455 }
456 out[a][b] = acc;
457 out[b][a] = acc;
458 }
459 }
460 out
461 }
462
463 /// Contract `t4` with two primary-space directions:
464 /// `out[a][b] = Σ_{c,d} t4[a][b][c][d] · u[c] · v[d]` — exactly the
465 /// `row_fourth_contracted` shape.
466 ///
467 /// As in [`Self::third_contracted`], the output is symmetric in `(i, j)`
468 /// (`t4[j][i][k][l] == t4[i][j][k][l]`, contracted in the same `(k, l)`
469 /// order), so the upper triangle `i ≤ j` is computed and mirrored —
470 /// BIT-IDENTICAL to the full nest, ~2× fewer inner `Σ_{k,l}` contractions,
471 /// and the inner double loop stays the original contiguous/vectorisable form.
472 pub fn fourth_contracted(&self, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
473 let mut out = [[0.0; K]; K];
474 for i in 0..K {
475 for j in i..K {
476 let mut acc = 0.0;
477 for k in 0..K {
478 for l in 0..K {
479 acc += self.t4[i][j][k][l] * u[k] * w[l];
480 }
481 }
482 out[i][j] = acc;
483 out[j][i] = acc;
484 }
485 }
486 out
487 }
488}
489
490impl<const K: usize> jet_algebra::JetAlgebra<5> for Tower4<K> {
491 #[inline]
492 fn derivative(&self, labels: &[usize]) -> f64 {
493 self.deriv(labels)
494 }
495
496 fn map_derivatives<F>(&self, mut f: F) -> Self
497 where
498 F: FnMut(&[usize]) -> f64,
499 {
500 let mut out = Self::zero();
501 out.v = f(&[]);
502 for i in 0..K {
503 let labels = [i];
504 out.g[i] = f(&labels);
505 }
506 for i in 0..K {
507 for j in 0..K {
508 let labels = [i, j];
509 out.h[i][j] = f(&labels);
510 }
511 }
512 for i in 0..K {
513 for j in 0..K {
514 for k in 0..K {
515 let labels = [i, j, k];
516 out.t3[i][j][k] = f(&labels);
517 }
518 }
519 }
520 for i in 0..K {
521 for j in 0..K {
522 for k in 0..K {
523 for l in 0..K {
524 let labels = [i, j, k, l];
525 out.t4[i][j][k][l] = f(&labels);
526 }
527 }
528 }
529 }
530 out
531 }
532}
533
534/// Truncated SECOND-order multivariate Taylor scalar in `K` variables.
535///
536/// This is the value/gradient/Hessian-only sibling of [`Tower4`]. Every
537/// channel it carries (`v`, `g`, `h`) is computed by the SAME formulas
538/// [`Tower4`] uses for those orders, so for any program written over both
539/// towers the order-≤2 outputs are *bit-identical*: the order-2 Leibniz and
540/// Faà-di-Bruno terms read only the order-≤2 channels of their inputs (see
541/// [`Tower4::mul`] / [`Tower4::compose_unary`] — `out.h` never touches `t3`
542/// or `t4`), so dropping the third/fourth tensors cannot perturb the value,
543/// gradient, or Hessian.
544///
545/// It exists purely for performance: an inner Newton step and a value-only
546/// outer-objective probe need at most curvature, never the outer-κ/ψ
547/// third/fourth derivatives. Evaluating a row likelihood over
548/// `Tower2` skips the `K⁴` fourth-tensor product/composition arithmetic that
549/// dominates the cold marginal-slope fit, while returning the exact same
550/// `(v, g, h)`.
551#[derive(Clone, Copy, Debug)]
552pub struct Tower2<const K: usize> {
553 /// Value ℓ.
554 pub v: f64,
555 /// Gradient ∂ℓ/∂p_a.
556 pub g: [f64; K],
557 /// Hessian ∂²ℓ/∂p_a∂p_b (symmetric).
558 pub h: [[f64; K]; K],
559}
560
561impl<const K: usize> Tower2<K> {
562 /// The additive identity.
563 pub fn zero() -> Self {
564 Self {
565 v: 0.0,
566 g: [0.0; K],
567 h: [[0.0; K]; K],
568 }
569 }
570
571 /// A constant: value `c`, all derivatives zero.
572 pub fn constant(c: f64) -> Self {
573 let mut out = Self::zero();
574 out.v = c;
575 out
576 }
577
578 /// The seeded variable `p_idx` with current value `value`:
579 /// unit first derivative in slot `idx`, zero elsewhere and above.
580 pub fn variable(value: f64, idx: usize) -> Self {
581 let mut out = Self::constant(value);
582 out.g[idx] = 1.0;
583 out
584 }
585
586 /// Read the derivative tensor entry whose differentiation axes are
587 /// `labels` (length 0..=2): value, `g`, `h`.
588 #[inline]
589 fn deriv(&self, labels: &[usize]) -> f64 {
590 assert!(
591 labels.len() <= 2,
592 "Tower2 carries at most second-order derivatives"
593 );
594 match labels.len() {
595 0 => self.v,
596 1 => self.g[labels[0]],
597 _ => self.h[labels[0]][labels[1]],
598 }
599 }
600
601 /// Exact truncated (order ≤ 2) Leibniz product. The `v`/`g`/`h` upper
602 /// triangle matches [`Tower4::mul`] term-for-term.
603 ///
604 /// # Symmetry fast path
605 ///
606 /// The order-≤2 Leibniz Hessian
607 /// `h[i][j] = a.v·b.h[i][j] + a.g[i]·b.g[j] + a.g[j]·b.g[i] + a.h[i][j]·b.v`
608 /// is symmetric under `i ↔ j` whenever the operand Hessians are — which they
609 /// always are: `constant`/`variable` seed a symmetric (zero) `h`, and
610 /// `mul`/`compose_unary`/`add`/`scale` each preserve symmetry, so the
611 /// invariant holds for every tower a row program can build. We therefore
612 /// compute only the upper triangle `j ≥ i` and mirror it into the lower
613 /// triangle. At the `K = 9` survival width that is `K(K+1)/2 = 45` four-product
614 /// entry evaluations instead of `K² = 81`, and the win is larger in wall-clock
615 /// because the `648`-entry `h` spills at `K = 9` — halving the expensive
616 /// stores/reloads roughly halves the kernel (measured ≈2× on a `Tower2<9>`
617 /// mul-and-read throughput microbench; the dominant `mul` under every packed
618 /// scalar bottoms out here).
619 ///
620 /// The upper-triangle entries are BIT-IDENTICAL to the old rectangular form
621 /// (same term/accumulation order). The lower triangle now equals its mirror
622 /// exactly, where the rectangular form rounded `h[i][j]` and `h[j][i]`
623 /// independently (the two cross products accumulate in opposite order) and
624 /// left a ≤1-ulp asymmetry; mirroring removes it, so the result is exactly
625 /// symmetric — strictly closer to the true symmetric Hessian, not merely a
626 /// reordering. Dense-`h` consumers are all tolerance-gated (rel-tol ≥ 1e-11 ≫
627 /// 1e-16); the `f64`/`f64x4` lane oracle stays exact because
628 /// [`crate::jet_scalar::Order2Lane::mul`] mirrors term-for-term.
629 pub fn mul(&self, o: &Self) -> Self {
630 let a = self;
631 let b = o;
632 let mut out = Self::zero();
633 out.v = a.v * b.v;
634 for i in 0..K {
635 out.g[i] = a.v * b.g[i] + a.g[i] * b.v;
636 }
637 for i in 0..K {
638 for j in i..K {
639 let hij = a.v * b.h[i][j] + a.g[i] * b.g[j] + a.g[j] * b.g[i] + a.h[i][j] * b.v;
640 out.h[i][j] = hij;
641 out.h[j][i] = hij;
642 }
643 }
644 out
645 }
646
647 /// Exact (order ≤ 2) multivariate Faà di Bruno composition `f ∘ self`.
648 ///
649 /// `d = [f(u), f′(u), f″(u)]` evaluated at `u = self.v`. The `v`/`g`/`h`
650 /// channels match [`Tower4::compose_unary`] term-for-term (which uses only
651 /// `d[0..=2]` for those orders), so this is a strict truncation, not an
652 /// approximation. The full-order `[f64; 5]` derivative stacks the families
653 /// already produce can be passed by slicing their first three entries.
654 ///
655 /// # Codegen
656 ///
657 /// Order-≤2 Faà di Bruno is a tiny closed form, so this evaluates it
658 /// directly instead of routing through the generic
659 /// [`jet_algebra::faa_di_bruno`] set-partition walker (recursion + per-block
660 /// closure dispatch). That matters because this is the kernel under EVERY
661 /// packed scalar — [`crate::jet_scalar::Order2`] / `OneSeed` / `TwoSeed`
662 /// composition all bottom out here — so the straight-line form (whose inner
663 /// loops auto-vectorise to NEON/SSE 2-wide and which emits zero outlined
664 /// walker calls) lifts all of them at once.
665 ///
666 /// The term and accumulation order is BIT-IDENTICAL to the walker it
667 /// replaces: each output channel mirrors the walker's `total = 0.0` start
668 /// (the explicit `acc` accumulator), so a signed-zero product collapses to
669 /// `+0.0` exactly as `total += prod` does. Proven `to_bits`-identical on
670 /// `v`/`g`/`h` across `K ∈ {2,3,4,9}`, 5000 random inputs each (incl.
671 /// zeroed / sign-varied stacks). The order-≤2 walker partitions are:
672 /// `g[i]` = `f′·u_i` (single block `{i}`)
673 /// `h[i][j]` = `f′·u_ij + (f″·u_i)·u_j` (blocks `{ij}` then `{i}{j}`),
674 /// with `f′ = d[1]`, `f″ = d[2]`, `u_* = self.{g,h}`.
675 pub fn compose_unary(&self, d: [f64; 3]) -> Self {
676 let mut out = Self::zero();
677 out.v = d[0];
678 for i in 0..K {
679 let mut acc = 0.0;
680 acc += d[1] * self.g[i];
681 out.g[i] = acc;
682 }
683 for i in 0..K {
684 for j in 0..K {
685 let mut acc = 0.0;
686 acc += d[1] * self.h[i][j];
687 acc += d[2] * self.g[i] * self.g[j];
688 out.h[i][j] = acc;
689 }
690 }
691 out
692 }
693
694 /// Multiply every channel by a plain scalar.
695 pub fn scale(&self, s: f64) -> Self {
696 let mut out = *self;
697 out.v *= s;
698 for i in 0..K {
699 out.g[i] *= s;
700 for j in 0..K {
701 out.h[i][j] *= s;
702 }
703 }
704 out
705 }
706
707 /// e^self.
708 pub fn exp(&self) -> Self {
709 let e = self.v.exp();
710 self.compose_unary([e, e, e])
711 }
712
713 /// √self. Caller guarantees positivity.
714 pub fn sqrt(&self) -> Self {
715 let u = self.v;
716 let s = u.sqrt();
717 self.compose_unary([s, 0.5 / s, -0.25 / (u * s)])
718 }
719}
720
721impl<const K: usize> jet_algebra::JetAlgebra<3> for Tower2<K> {
722 #[inline]
723 fn derivative(&self, labels: &[usize]) -> f64 {
724 self.deriv(labels)
725 }
726
727 fn map_derivatives<F>(&self, mut f: F) -> Self
728 where
729 F: FnMut(&[usize]) -> f64,
730 {
731 let mut out = Self::zero();
732 out.v = f(&[]);
733 for i in 0..K {
734 let labels = [i];
735 out.g[i] = f(&labels);
736 }
737 for i in 0..K {
738 for j in 0..K {
739 let labels = [i, j];
740 out.h[i][j] = f(&labels);
741 }
742 }
743 out
744 }
745}
746
747impl<const K: usize> std::ops::Add for Tower2<K> {
748 type Output = Self;
749 fn add(self, o: Self) -> Self {
750 let mut out = self;
751 out.v += o.v;
752 for i in 0..K {
753 out.g[i] += o.g[i];
754 for j in 0..K {
755 out.h[i][j] += o.h[i][j];
756 }
757 }
758 out
759 }
760}
761
762impl<const K: usize> std::ops::Mul for Tower2<K> {
763 type Output = Self;
764 fn mul(self, o: Self) -> Self {
765 Tower2::mul(&self, &o)
766 }
767}
768
769impl<const K: usize> std::ops::Add<f64> for Tower2<K> {
770 type Output = Self;
771 fn add(self, c: f64) -> Self {
772 let mut out = self;
773 out.v += c;
774 out
775 }
776}
777
778impl<const K: usize> std::ops::Mul<f64> for Tower2<K> {
779 type Output = Self;
780 fn mul(self, c: f64) -> Self {
781 self.scale(c)
782 }
783}
784
785/// Truncated THIRD-order multivariate Taylor scalar in `K` variables.
786///
787/// The value/gradient/Hessian/third-derivative sibling of [`Tower4`], standing
788/// between [`Tower2`] and [`Tower4`]. Every channel it carries (`v`, `g`, `h`,
789/// `t3`) is computed by the SAME shared Leibniz / Faà-di-Bruno kernels
790/// [`Tower4`] uses for those orders, and the order-≤3 terms of those kernels
791/// read only the order-≤3 channels of their inputs (the order-3 Faà-di-Bruno
792/// partitions never reach the f⁗ stack slot or the inner `t4` tensor — see
793/// [`Tower4::compose_unary`]). So for any program written over both towers the
794/// order-≤3 outputs are *bit-identical*: dropping the fourth tensor cannot
795/// perturb the value, gradient, Hessian, or third derivatives.
796///
797/// It exists purely for performance, exactly like [`Tower2`]: a consumer that
798/// needs up to third derivatives (the survival location-scale row kernel reads
799/// `g`, the diagonal `h`, and the diagonal `t3`, but never `t4`) pays the
800/// `K³` third-tensor arithmetic but skips the `K⁴` fourth-tensor
801/// product/composition that otherwise dominates the per-row cost.
802#[derive(Clone, Copy, Debug)]
803pub struct Tower3<const K: usize> {
804 /// Value ℓ.
805 pub v: f64,
806 /// Gradient ∂ℓ/∂p_a.
807 pub g: [f64; K],
808 /// Hessian ∂²ℓ/∂p_a∂p_b (symmetric).
809 pub h: [[f64; K]; K],
810 /// Third derivatives ∂³ℓ/∂p_a∂p_b∂p_c (fully symmetric).
811 pub t3: [[[f64; K]; K]; K],
812}
813
814impl<const K: usize> Tower3<K> {
815 /// The additive identity.
816 pub fn zero() -> Self {
817 Self {
818 v: 0.0,
819 g: [0.0; K],
820 h: [[0.0; K]; K],
821 t3: [[[0.0; K]; K]; K],
822 }
823 }
824
825 /// A constant: value `c`, all derivatives zero.
826 pub fn constant(c: f64) -> Self {
827 let mut out = Self::zero();
828 out.v = c;
829 out
830 }
831
832 /// The seeded variable `p_idx` with current value `value`:
833 /// unit first derivative in slot `idx`, zero elsewhere and above.
834 pub fn variable(value: f64, idx: usize) -> Self {
835 let mut out = Self::constant(value);
836 out.g[idx] = 1.0;
837 out
838 }
839
840 /// Read the (fully symmetric) derivative tensor entry whose differentiation
841 /// axes are `labels` (length 0..=3): value, `g`, `h`, `t3`.
842 #[inline]
843 fn deriv(&self, labels: &[usize]) -> f64 {
844 assert!(
845 labels.len() <= 3,
846 "Tower3 carries at most third-order derivatives"
847 );
848 match labels.len() {
849 0 => self.v,
850 1 => self.g[labels[0]],
851 2 => self.h[labels[0]][labels[1]],
852 _ => self.t3[labels[0]][labels[1]][labels[2]],
853 }
854 }
855
856 /// Exact truncated (order ≤ 3) Leibniz product. The `v`/`g`/`h`/`t3`
857 /// channels match [`Tower4::mul`] term-for-term.
858 ///
859 /// # Codegen
860 ///
861 /// Straight-line per-entry subset sums instead of the
862 /// `jet_algebra::leibniz_product` walker — the order-≤3 sibling of
863 /// [`Tower4::mul`] (no `t4`). Loop nest unchanged, no unroll over `K`, no
864 /// code bloat; auto-vectorises. BIT-IDENTICAL: terms in the walker's exact
865 /// subset order with an `acc = 0.0` accumulator start (load-bearing for the
866 /// signed-zero leading product on exact-`0.0` jet channels). Proven
867 /// `to_bits`-identical on `v`/`g`/`h`/`t3` across `K ∈ {2,3,4,9}`, 5000
868 /// zero/sign-stressed inputs each (these channel formulas are exactly the
869 /// `g`/`h`/`t3` of the [`Tower4::mul`] oracle, which passes that stress).
870 pub fn mul(&self, o: &Self) -> Self {
871 let a = self;
872 let b = o;
873 let mut out = Self::zero();
874 out.v = a.v * b.v;
875 for i in 0..K {
876 let mut acc = 0.0;
877 acc += a.v * b.g[i];
878 acc += a.g[i] * b.v;
879 out.g[i] = acc;
880 }
881 // Hessian is symmetric under i↔j; upper triangle + mirror (see Tower2::mul).
882 for i in 0..K {
883 for j in i..K {
884 let mut acc = 0.0;
885 acc += a.v * b.h[i][j];
886 acc += a.g[i] * b.g[j];
887 acc += a.g[j] * b.g[i];
888 acc += a.h[i][j] * b.v;
889 out.h[i][j] = acc;
890 out.h[j][i] = acc;
891 }
892 }
893 for i in 0..K {
894 for j in 0..K {
895 for k in 0..K {
896 // subsets of {i,j,k}: {} {i} {j} {ij} {k} {ik} {jk} {ijk}
897 let mut acc = 0.0;
898 acc += a.v * b.t3[i][j][k];
899 acc += a.g[i] * b.h[j][k];
900 acc += a.g[j] * b.h[i][k];
901 acc += a.h[i][j] * b.g[k];
902 acc += a.g[k] * b.h[i][j];
903 acc += a.h[i][k] * b.g[j];
904 acc += a.h[j][k] * b.g[i];
905 acc += a.t3[i][j][k] * b.v;
906 out.t3[i][j][k] = acc;
907 }
908 }
909 }
910 out
911 }
912
913 /// Ref-taking elementwise sum, the by-ref twin of the `std::ops::Add`
914 /// operator (which consumes by value). Mirrors the inherent `mul`/`scale`
915 /// API so a chain like `a.mul(&b).add(&c)` reads uniformly without moving
916 /// out of the borrowed operands.
917 pub fn add(&self, o: &Self) -> Self {
918 *self + *o
919 }
920
921 /// Ref-taking elementwise difference, the by-ref twin of `std::ops::Sub`.
922 pub fn sub(&self, o: &Self) -> Self {
923 *self + o.scale(-1.0)
924 }
925
926 /// Exact (order ≤ 3) multivariate Faà di Bruno composition `f ∘ self`.
927 ///
928 /// `d = [f(u), f′(u), f″(u), f‴(u)]` evaluated at `u = self.v`. The
929 /// `v`/`g`/`h`/`t3` channels match [`Tower4::compose_unary`] term-for-term
930 /// (which uses only `d[0..=3]` for those orders), so this is a strict
931 /// truncation, not an approximation. The full-order `[f64; 5]` derivative
932 /// stacks the families already produce can be passed by slicing their first
933 /// four entries.
934 ///
935 /// # Codegen
936 ///
937 /// Order-≤3 Faà di Bruno written as a compact closed form instead of the
938 /// recursive [`jet_algebra::faa_di_bruno`] walker — the order-≤2 sibling of
939 /// [`Tower4::compose_unary`], one tensor order shallower. The loop nest is
940 /// unchanged (no unroll over `K`, no code bloat: measured on a `Tower3<9>`
941 /// compose-and-read consumer the new form is faster and SMALLER — asm: 71
942 /// walker `bl` calls → 0, 39.5 KiB → 13.9 KiB, +197 NEON `.2d` ops).
943 /// BIT-IDENTICAL: terms in the walker's exact partition order, left-
944 /// associated block products, `acc = 0.0` accumulator start. Proven
945 /// `to_bits`-identical on `v`/`g`/`h`/`t3` across `K ∈ {2,3,4,9}`, 5000
946 /// random inputs each.
947 pub fn compose_unary(&self, d: [f64; 4]) -> Self {
948 let mut out = Self::zero();
949 out.v = d[0];
950 for i in 0..K {
951 let mut acc = 0.0;
952 acc += d[1] * self.g[i];
953 out.g[i] = acc;
954 }
955 for i in 0..K {
956 for j in 0..K {
957 let mut acc = 0.0;
958 acc += d[1] * self.h[i][j];
959 acc += d[2] * self.g[i] * self.g[j];
960 out.h[i][j] = acc;
961 }
962 }
963 for i in 0..K {
964 for j in 0..K {
965 for k in 0..K {
966 // walker partitions: {ijk} {ij}{k} {ik}{j} {i}{jk} {i}{j}{k}
967 let mut acc = 0.0;
968 acc += d[1] * self.t3[i][j][k];
969 acc += d[2] * self.h[i][j] * self.g[k];
970 acc += d[2] * self.h[i][k] * self.g[j];
971 acc += d[2] * self.g[i] * self.h[j][k];
972 acc += d[3] * self.g[i] * self.g[j] * self.g[k];
973 out.t3[i][j][k] = acc;
974 }
975 }
976 }
977 out
978 }
979
980 /// Compose with a unary special-function whose `[f64; 4]` derivative stack is
981 /// built from the base value through `stack_fn`. Evaluates `stack_fn(self.v)`
982 /// once and forwards to [`Self::compose_unary`], so it is bit-identical to the
983 /// explicit form. The order-≤3 sibling of `Tower4::compose_unary_with`.
984 #[inline]
985 pub fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 4]) -> Self {
986 self.compose_unary(stack_fn(self.v))
987 }
988
989 /// Multiply every channel by a plain scalar.
990 pub fn scale(&self, s: f64) -> Self {
991 let mut out = *self;
992 out.v *= s;
993 for i in 0..K {
994 out.g[i] *= s;
995 for j in 0..K {
996 out.h[i][j] *= s;
997 for k in 0..K {
998 out.t3[i][j][k] *= s;
999 }
1000 }
1001 }
1002 out
1003 }
1004}
1005
1006impl<const K: usize> jet_algebra::JetAlgebra<4> for Tower3<K> {
1007 #[inline]
1008 fn derivative(&self, labels: &[usize]) -> f64 {
1009 self.deriv(labels)
1010 }
1011
1012 fn map_derivatives<F>(&self, mut f: F) -> Self
1013 where
1014 F: FnMut(&[usize]) -> f64,
1015 {
1016 let mut out = Self::zero();
1017 out.v = f(&[]);
1018 for i in 0..K {
1019 let labels = [i];
1020 out.g[i] = f(&labels);
1021 }
1022 for i in 0..K {
1023 for j in 0..K {
1024 let labels = [i, j];
1025 out.h[i][j] = f(&labels);
1026 }
1027 }
1028 for i in 0..K {
1029 for j in 0..K {
1030 for k in 0..K {
1031 let labels = [i, j, k];
1032 out.t3[i][j][k] = f(&labels);
1033 }
1034 }
1035 }
1036 out
1037 }
1038}
1039
1040impl<const K: usize> std::ops::Add for Tower3<K> {
1041 type Output = Self;
1042 fn add(self, o: Self) -> Self {
1043 let mut out = self;
1044 out.v += o.v;
1045 for i in 0..K {
1046 out.g[i] += o.g[i];
1047 for j in 0..K {
1048 out.h[i][j] += o.h[i][j];
1049 for k in 0..K {
1050 out.t3[i][j][k] += o.t3[i][j][k];
1051 }
1052 }
1053 }
1054 out
1055 }
1056}
1057
1058pub fn ln_gamma_derivative_stack(x: f64) -> [f64; 5] {
1059 [
1060 statrs::function::gamma::ln_gamma(x),
1061 digamma_positive(x),
1062 polygamma_positive::<1>(x),
1063 polygamma_positive::<2>(x),
1064 polygamma_positive::<3>(x),
1065 ]
1066}
1067
1068pub fn ln_gamma_derivative_stack_order2(x: f64) -> [f64; 3] {
1069 [
1070 statrs::function::gamma::ln_gamma(x),
1071 digamma_positive(x),
1072 polygamma_positive::<1>(x),
1073 ]
1074}
1075
1076pub fn ln_gamma_derivative_stack_order3(x: f64) -> [f64; 4] {
1077 [
1078 statrs::function::gamma::ln_gamma(x),
1079 digamma_positive(x),
1080 polygamma_positive::<1>(x),
1081 polygamma_positive::<2>(x),
1082 ]
1083}
1084
1085pub fn digamma_derivative_stack(x: f64) -> [f64; 5] {
1086 [
1087 digamma_positive(x),
1088 polygamma_positive::<1>(x),
1089 polygamma_positive::<2>(x),
1090 polygamma_positive::<3>(x),
1091 polygamma_positive::<4>(x),
1092 ]
1093}
1094
1095/// Scalar digamma ψ(x) for x>0. Bit-identical to `digamma_derivative_stack(x)[0]`
1096/// and to `ln_gamma_derivative_stack(x)[1]`, but evaluates ONLY ψ — the four
1097/// higher polygammas those `[f64; 5]` stacks build are pure discarded work at a
1098/// scalar consumer that reads a single element. Hot-path row kernels that need
1099/// only the digamma value (e.g. the GAMLSS Beta observed cross weight) call this
1100/// instead of indexing `[0]` off a full derivative stack.
1101#[inline]
1102pub fn digamma(x: f64) -> f64 {
1103 digamma_positive(x)
1104}
1105
1106/// Scalar trigamma ψ′(x) for x>0. Bit-identical to
1107/// `trigamma_derivative_stack(x)[0]` (both use `polygamma_positive::<1>(x)`),
1108/// but evaluates ONLY ψ′ — the four higher polygammas (orders 2–5) the
1109/// `[f64; 5]` stack builds are discarded at a `[0]` consumer. Used by the
1110/// dispersion-channel Fisher-information row kernels (NB2 `ψ′(θ)−ψ′(θ+μ)`, Beta
1111/// `μψ′(μφ)−(1−μ)ψ′((1−μ)φ)`) which read the trigamma value alone.
1112#[inline]
1113pub fn trigamma(x: f64) -> f64 {
1114 polygamma_positive::<1>(x)
1115}
1116
1117fn digamma_positive(mut x: f64) -> f64 {
1118 if !(x.is_finite() && x > 0.0) {
1119 return f64::NAN;
1120 }
1121 let mut acc = 0.0;
1122 while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
1123 acc -= 1.0 / x;
1124 x += 1.0;
1125 }
1126 acc + digamma_asymptotic(x)
1127}
1128
1129fn polygamma_positive<const ORDER: usize>(mut x: f64) -> f64 {
1130 if !(x.is_finite() && x > 0.0) {
1131 return f64::NAN;
1132 }
1133 let mut acc = 0.0;
1134 while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
1135 acc += polygamma_recurrence_term::<ORDER>(x);
1136 x += 1.0;
1137 }
1138 acc + polygamma_asymptotic::<ORDER>(x)
1139}
1140
1141const POLYGAMMA_ASYMPTOTIC_MIN_X: f64 = 20.0;
1142const BERNOULLI_EVEN: [(usize, f64); 10] = [
1143 (2, 1.0 / 6.0),
1144 (4, -1.0 / 30.0),
1145 (6, 1.0 / 42.0),
1146 (8, -1.0 / 30.0),
1147 (10, 5.0 / 66.0),
1148 (12, -691.0 / 2730.0),
1149 (14, 7.0 / 6.0),
1150 (16, -3617.0 / 510.0),
1151 (18, 43867.0 / 798.0),
1152 (20, -174611.0 / 330.0),
1153];
1154
1155fn polygamma_recurrence_term<const ORDER: usize>(x: f64) -> f64 {
1156 let coefficient = const {
1157 let sign = if ORDER % 2 == 1 { 1.0 } else { -1.0 };
1158 sign * factorial(ORDER)
1159 };
1160 coefficient / x.powi((ORDER + 1) as i32)
1161}
1162
1163fn digamma_asymptotic(x: f64) -> f64 {
1164 let mut out = x.ln() - 0.5 / x;
1165 for (bernoulli_order, bernoulli) in BERNOULLI_EVEN {
1166 out -= bernoulli / (bernoulli_order as f64 * x.powi(bernoulli_order as i32));
1167 }
1168 out
1169}
1170
1171fn polygamma_asymptotic<const ORDER: usize>(x: f64) -> f64 {
1172 const { assert!(ORDER >= 1 && ORDER <= 5) };
1173 let (leading, half_term) = const {
1174 let sign = if ORDER % 2 == 1 { 1.0 } else { -1.0 };
1175 (sign * factorial(ORDER - 1), sign * factorial(ORDER))
1176 };
1177 let mut out =
1178 leading / x.powi(ORDER as i32) + half_term / (2.0 * x.powi((ORDER + 1) as i32));
1179
1180 // Derivative order and Bernoulli coefficients are fixed by the series.
1181 // Build their factorial products once at compile time, including in dev
1182 // builds: #2668's NB profile spent 13.53% of cycles in factorial alone.
1183 let coefficients = const { polygamma_asymptotic_coefficients::<ORDER>() };
1184 for (power, coefficient) in coefficients {
1185 out += coefficient / x.powi(power);
1186 }
1187 out
1188}
1189
1190const fn polygamma_asymptotic_coefficients<const ORDER: usize>() -> [(i32, f64); 10] {
1191 let sign = if ORDER % 2 == 1 { 1.0 } else { -1.0 };
1192 let mut coefficients = [(0, 0.0); BERNOULLI_EVEN.len()];
1193 let mut index = 0;
1194 while index < BERNOULLI_EVEN.len() {
1195 let (power, bernoulli) = BERNOULLI_EVEN[index];
1196 coefficients[index] = (
1197 (power + ORDER) as i32,
1198 sign * bernoulli * rising_factorial(power, ORDER) / power as f64,
1199 );
1200 index += 1;
1201 }
1202 coefficients
1203}
1204
1205const fn factorial(n: usize) -> f64 {
1206 rising_factorial(1, n)
1207}
1208
1209const fn rising_factorial(start: usize, len: usize) -> f64 {
1210 let mut product = 1.0;
1211 let mut offset = 0;
1212 while offset < len {
1213 product *= (start + offset) as f64;
1214 offset += 1;
1215 }
1216 product
1217}
1218
1219impl<const K: usize> std::ops::Add for Tower4<K> {
1220 type Output = Self;
1221 fn add(self, o: Self) -> Self {
1222 let mut out = self;
1223 out.v += o.v;
1224 for i in 0..K {
1225 out.g[i] += o.g[i];
1226 for j in 0..K {
1227 out.h[i][j] += o.h[i][j];
1228 for k in 0..K {
1229 out.t3[i][j][k] += o.t3[i][j][k];
1230 for l in 0..K {
1231 out.t4[i][j][k][l] += o.t4[i][j][k][l];
1232 }
1233 }
1234 }
1235 }
1236 out
1237 }
1238}
1239
1240impl<const K: usize> std::ops::Sub for Tower4<K> {
1241 type Output = Self;
1242 fn sub(self, o: Self) -> Self {
1243 self + o.scale(-1.0)
1244 }
1245}
1246
1247impl<const K: usize> std::ops::Neg for Tower4<K> {
1248 type Output = Self;
1249 fn neg(self) -> Self {
1250 self.scale(-1.0)
1251 }
1252}
1253
1254impl<const K: usize> std::ops::Mul for Tower4<K> {
1255 type Output = Self;
1256 fn mul(self, o: Self) -> Self {
1257 Tower4::mul(&self, &o)
1258 }
1259}
1260
1261impl<const K: usize> std::ops::Div for Tower4<K> {
1262 type Output = Self;
1263 fn div(self, o: Self) -> Self {
1264 Tower4::mul(&self, &o.recip())
1265 }
1266}
1267
1268impl<const K: usize> std::ops::Add<f64> for Tower4<K> {
1269 type Output = Self;
1270 fn add(self, c: f64) -> Self {
1271 let mut out = self;
1272 out.v += c;
1273 out
1274 }
1275}
1276
1277impl<const K: usize> std::ops::Sub<f64> for Tower4<K> {
1278 type Output = Self;
1279 fn sub(self, c: f64) -> Self {
1280 self + (-c)
1281 }
1282}
1283
1284impl<const K: usize> std::ops::Mul<f64> for Tower4<K> {
1285 type Output = Self;
1286 fn mul(self, c: f64) -> Self {
1287 self.scale(c)
1288 }
1289}
1290
1291// ── Implicit-function and moving-boundary seams (#932 flex) ──────────
1292//
1293// The flexible survival marginal-slope row loss is NOT a free composition
1294// of the primaries: it threads an IMPLICIT calibration intercept `a(θ)`
1295// solving a constraint `F(a, θ) = 0`, and integrates a density over cells
1296// whose edges `z_L(θ), z_R(θ)` MOVE with θ through that intercept. Plain
1297// `Tower4` Faà di Bruno cannot express either — so the flex tower was the
1298// last hand-written one in the codebase, and the genus of #736-class
1299// drift bugs (the (g,w0) deviation-cross third was 3× short for exactly
1300// this reason). These two combinators close that gap: once the constraint
1301// `F` and the integrand/boundaries are themselves towers, the intercept's
1302// derivative tower and the integral's derivative tower come out EXACTLY at
1303// every order — there is no order left to hand-code and forget.
1304
1305// ── The program seam ─────────────────────────────────────────────────
1306
1307// ── The canonical single-source seam (#932 consolidation) ────────────
1308//
1309// `RowProgram<K>` is the ONE row-program interface #932 converges every family
1310// onto. Its generic `eval<S: JetScalar<K>>` body is the go-forward derivation
1311// surface for every calculus channel; `program_*` selects only the derivative
1312// representation each consumer needs.
1313
1314/// The single source of truth #932 asks for: a family's row negative
1315/// log-likelihood written ONCE over the generic [`crate::jet_scalar::JetScalar`]
1316/// interface, from which every `RowKernel` (gam-models) derivative channel is
1317/// mechanically derived. A family implements ONLY this (plus its linear Jacobian
1318/// wiring, which is family data, not calculus) — it cannot author an independent
1319/// derivative tower, because there is no other channel to author.
1320///
1321/// Because a body uses only `add`/`sub`/`mul`/`scale`/`exp`/`ln`/… — all provided
1322/// by [`crate::jet_scalar::JetScalar`] — the SAME body re-instantiates at
1323/// [`crate::jet_scalar::Order2`] (value/grad/Hessian), [`crate::jet_scalar::OneSeed`]
1324/// (contracted third), [`crate::jet_scalar::TwoSeed`] (contracted fourth), and the
1325/// full [`Tower4`] (every channel), with the contraction folded into the
1326/// differentiation so no dense `t3`/`t4` is ever materialised.
1327pub trait RowProgram<const K: usize>: Send + Sync {
1328 /// Number of observations the program covers.
1329 fn n_rows(&self) -> usize;
1330
1331 /// Current primary-scalar values for `row` (where to seed the scalar).
1332 fn primaries(&self, row: usize) -> Result<[f64; K], String>;
1333
1334 /// The row NLL evaluated on a generic jet scalar. `p[a]` arrives pre-seeded
1335 /// (base value + per-scalar nilpotent directions) by the caller; the body
1336 /// uses ONLY [`crate::jet_scalar::JetScalar`] ops and per-row data (response,
1337 /// censoring, offsets) entering as constants.
1338 fn eval<S: crate::jet_scalar::JetScalar<K>>(&self, row: usize, p: &[S; K])
1339 -> Result<S, String>;
1340}
1341
1342/// Maximum size of one canonical dense-jet storage object kept on the call
1343/// stack. Small fixed-width programs stay allocation-free; wider derivative
1344/// representations use exact-length heap storage instead of making the thread
1345/// stack scale as `K * size_of::<S>()`. A full dense result larger than this
1346/// boundary is rejected in favor of the bounded directional APIs.
1347///
1348/// This is a storage-policy boundary, not a calculus fallback: both branches
1349/// invoke the same [`RowProgram::eval`] expression with the same scalar type.
1350const PROGRAM_DENSE_JET_STACK_BUDGET_BYTES: usize = 64 * 1024;
1351
1352#[inline]
1353fn program_primary_jets_fit_stack<S, const K: usize>() -> bool {
1354 std::mem::size_of::<S>()
1355 .checked_mul(K)
1356 .is_some_and(|bytes| bytes <= PROGRAM_DENSE_JET_STACK_BUDGET_BYTES)
1357}
1358
1359fn evaluate_program_with_stack_primaries<const K: usize, P, S>(
1360 prog: &P,
1361 row: usize,
1362 mut seed: impl FnMut(usize) -> S,
1363) -> Result<S, String>
1364where
1365 P: RowProgram<K> + ?Sized,
1366 S: crate::jet_scalar::JetScalar<K>,
1367{
1368 let vars: [S; K] = std::array::from_fn(&mut seed);
1369 prog.eval(row, &vars)
1370}
1371
1372#[inline(never)]
1373fn evaluate_program_with_heap_primaries<const K: usize, P, S>(
1374 prog: &P,
1375 row: usize,
1376 seed: impl FnMut(usize) -> S,
1377) -> Result<S, String>
1378where
1379 P: RowProgram<K> + ?Sized,
1380 S: crate::jet_scalar::JetScalar<K>,
1381{
1382 // The exact-size range builds precisely K initialized Copy scalars in
1383 // heap-backed storage. Converting the boxed slice to a boxed array changes
1384 // only its type; it never materializes `[S; K]` on the stack.
1385 let vars: Box<[S]> = (0..K).map(seed).collect();
1386 let vars: Box<[S; K]> = vars.try_into().map_err(|vars: Box<[S]>| {
1387 format!(
1388 "canonical row program seeded {} primary jets; expected exactly {K}",
1389 vars.len()
1390 )
1391 })?;
1392 prog.eval(row, &vars)
1393}
1394
1395#[inline]
1396fn evaluate_program_with_seeded_primaries<const K: usize, P, S>(
1397 prog: &P,
1398 row: usize,
1399 seed: impl FnMut(usize) -> S,
1400) -> Result<S, String>
1401where
1402 P: RowProgram<K> + ?Sized,
1403 S: crate::jet_scalar::JetScalar<K>,
1404{
1405 if program_primary_jets_fit_stack::<S, K>() {
1406 evaluate_program_with_stack_primaries(prog, row, seed)
1407 } else {
1408 evaluate_program_with_heap_primaries(prog, row, seed)
1409 }
1410}
1411
1412/// Derive the `row_kernel` channel `(nll, ∇, H)` from a [`RowProgram`] at the
1413/// value/gradient/Hessian scalar [`crate::jet_scalar::Order2`], WITHOUT
1414/// materialising any third / fourth tensor.
1415pub fn program_row_kernel<const K: usize, P: RowProgram<K> + ?Sized>(
1416 prog: &P,
1417 row: usize,
1418) -> Result<(f64, [f64; K], [[f64; K]; K]), String> {
1419 let base = prog.primaries(row)?;
1420 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1421 <crate::jet_scalar::Order2<K> as crate::jet_scalar::JetScalar<K>>::variable(base[a], a)
1422 })?;
1423 Ok(s.into_channels())
1424}
1425
1426/// Derive the `row_third_contracted(dir)` channel `Σ_c ℓ_{abc} dir_c` from a
1427/// [`RowProgram`] at the one-seed scalar [`crate::jet_scalar::OneSeed`], WITHOUT
1428/// materialising the dense `t3`.
1429pub fn program_third_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
1430 prog: &P,
1431 row: usize,
1432 dir: &[f64; K],
1433) -> Result<[[f64; K]; K], String> {
1434 let base = prog.primaries(row)?;
1435 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1436 crate::jet_scalar::OneSeed::seed_direction(base[a], a, dir[a])
1437 })?;
1438 Ok(s.contracted_third())
1439}
1440
1441/// Derive the `row_fourth_contracted(u, v)` channel `Σ_{cd} ℓ_{abcd} u_c v_d`
1442/// from a [`RowProgram`] at the two-seed scalar [`crate::jet_scalar::TwoSeed`],
1443/// WITHOUT materialising the dense `t4`.
1444pub fn program_fourth_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
1445 prog: &P,
1446 row: usize,
1447 dir_u: &[f64; K],
1448 dir_v: &[f64; K],
1449) -> Result<[[f64; K]; K], String> {
1450 let base = prog.primaries(row)?;
1451 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1452 crate::jet_scalar::TwoSeed::seed(base[a], a, dir_u[a], dir_v[a])
1453 })?;
1454 Ok(s.contracted_fourth())
1455}
1456
1457/// Derive every channel `(v, g, h, t3, t4)` in one pass from a [`RowProgram`] at
1458/// the full dense [`Tower4`] scalar.
1459///
1460/// The result is boxed so the return slot itself remains bounded independently
1461/// of `K`. Dense towers above the canonical storage budget are rejected before
1462/// the program is touched; consumers at those widths must request only the
1463/// channels they need through [`program_row_kernel`],
1464/// [`program_third_contracted`], and [`program_fourth_contracted`].
1465pub fn program_full_tower<const K: usize, P: RowProgram<K> + ?Sized>(
1466 prog: &P,
1467 row: usize,
1468) -> Result<Box<Tower4<K>>, String> {
1469 let tower_bytes = std::mem::size_of::<Tower4<K>>();
1470 if tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES {
1471 return Err(format!(
1472 "canonical dense Tower4<{K}> requires {tower_bytes} bytes, exceeding the {}-byte \
1473 storage budget; use the bounded row-kernel and directional channel APIs",
1474 PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
1475 ));
1476 }
1477 let base = prog.primaries(row)?;
1478 evaluate_program_with_seeded_primaries(prog, row, |a| Tower4::variable(base[a], a))
1479 .map(Box::new)
1480}
1481
1482// ── The oracle ───────────────────────────────────────────────────────
1483
1484/// One row's worth of hand-written kernel outputs, as claimed by a
1485/// `RowKernel` implementation, packaged for verification against the
1486/// tower truth. Plain data (no trait coupling) so any kernel — whatever
1487/// its visibility — can be audited from its own test module.
1488pub struct KernelChannels<const K: usize> {
1489 /// Claimed `(nll, ∇, H)` from `row_kernel`.
1490 pub value: f64,
1491 /// Claimed gradient.
1492 pub gradient: [f64; K],
1493 /// Claimed Hessian.
1494 pub hessian: [[f64; K]; K],
1495 /// Claimed `row_third_contracted(dir)` outputs as `(dir, claim)` pairs.
1496 pub third: Vec<([f64; K], [[f64; K]; K])>,
1497 /// Claimed `row_fourth_contracted(u, v)` outputs as `(u, v, claim)`.
1498 pub fourth: Vec<([f64; K], [f64; K], [[f64; K]; K])>,
1499}
1500
1501#[cfg(test)]
1502mod tests {
1503 use super::*;
1504
1505 /// `Tower3<K>` must be bit-identical to `Tower4<K>` on every channel it
1506 /// carries (value, gradient, Hessian, third derivatives). The order-≤3
1507 /// Leibniz / Faà-di-Bruno terms read only order-≤3 inner channels, so
1508 /// dropping the fourth tensor cannot perturb them. Exercises products
1509 /// (Leibniz cross-terms), unary composition, scaling, and addition — the
1510 /// same operations the survival location-scale `nll_index_tower` composes —
1511 /// across all mixed partials, not just the diagonal entries that kernel reads.
1512 #[test]
1513 fn tower3_matches_tower4_through_third_order() {
1514 let s_a: [f64; 5] = [
1515 0.3_f64.sin(),
1516 0.3_f64.cos(),
1517 -0.3_f64.sin(),
1518 -0.3_f64.cos(),
1519 0.3_f64.sin(),
1520 ];
1521 let s_b: [f64; 5] = [1.1, -0.4, 0.8, -0.2, 0.05];
1522 let s4 = |s: [f64; 5]| [s[0], s[1], s[2], s[3]];
1523
1524 let a4 = Tower4::<3>::variable(0.4, 0);
1525 let b4 = Tower4::<3>::variable(-0.7, 1);
1526 let c4 = Tower4::<3>::variable(0.9, 2);
1527 let prog4 = (a4.mul(&b4) + c4).compose_unary(s_a).scale(1.3)
1528 + a4.mul(&c4).scale(-0.7)
1529 + b4.compose_unary(s_b).scale(0.25);
1530
1531 let a3 = Tower3::<3>::variable(0.4, 0);
1532 let b3 = Tower3::<3>::variable(-0.7, 1);
1533 let c3 = Tower3::<3>::variable(0.9, 2);
1534 let prog3 = (a3.mul(&b3) + c3).compose_unary(s4(s_a)).scale(1.3)
1535 + a3.mul(&c3).scale(-0.7)
1536 + b3.compose_unary(s4(s_b)).scale(0.25);
1537
1538 assert_eq!(prog3.v.to_bits(), prog4.v.to_bits(), "value mismatch");
1539 for i in 0..3 {
1540 assert_eq!(
1541 prog3.g[i].to_bits(),
1542 prog4.g[i].to_bits(),
1543 "g[{i}] mismatch"
1544 );
1545 for j in 0..3 {
1546 assert_eq!(
1547 prog3.h[i][j].to_bits(),
1548 prog4.h[i][j].to_bits(),
1549 "h[{i}][{j}] mismatch"
1550 );
1551 for k in 0..3 {
1552 assert_eq!(
1553 prog3.t3[i][j][k].to_bits(),
1554 prog4.t3[i][j][k].to_bits(),
1555 "t3[{i}][{j}][{k}] mismatch"
1556 );
1557 }
1558 }
1559 }
1560 }
1561
1562 /// Binomial-logit row NLL, K=1: ℓ(η) = ln(1 + e^η) − y·η.
1563 /// The entire tower has textbook closed forms in μ = σ(η); this test
1564 /// pins the algebra (exp, ln, scalar mixes, Leibniz/Faà di Bruno) to
1565 /// analytic truth at near-machine precision.
1566 struct LogitProgram {
1567 eta: Vec<f64>,
1568 y: Vec<f64>,
1569 }
1570
1571 impl RowProgram<1> for LogitProgram {
1572 fn n_rows(&self) -> usize {
1573 self.eta.len()
1574 }
1575 fn primaries(&self, row: usize) -> Result<[f64; 1], String> {
1576 Ok([self.eta[row]])
1577 }
1578 fn eval<S: crate::jet_scalar::JetScalar<1>>(
1579 &self,
1580 row: usize,
1581 p: &[S; 1],
1582 ) -> Result<S, String> {
1583 let eta = p[0];
1584 Ok(eta
1585 .exp()
1586 .add(&S::constant(1.0))
1587 .ln()
1588 .sub(&eta.scale(self.y[row])))
1589 }
1590 }
1591
1592 #[test]
1593 fn logit_tower_matches_closed_forms() {
1594 let prog = LogitProgram {
1595 eta: vec![-2.3, -0.4, 0.0, 0.9, 3.1],
1596 y: vec![1.0, 0.0, 1.0, 0.0, 1.0],
1597 };
1598 for row in 0..prog.n_rows() {
1599 let t = program_full_tower(&prog, row).expect("logit program");
1600 let eta = prog.eta[row];
1601 let y = prog.y[row];
1602 let mu = 1.0 / (1.0 + (-eta).exp());
1603 let w = mu * (1.0 - mu);
1604 let expect = [
1605 (t.v, (1.0 + eta.exp()).ln() - y * eta, "value"),
1606 (t.g[0], mu - y, "grad"),
1607 (t.h[0][0], w, "hess"),
1608 (t.t3[0][0][0], w * (1.0 - 2.0 * mu), "third"),
1609 (
1610 t.t4[0][0][0][0],
1611 w * (1.0 - 6.0 * mu + 6.0 * mu * mu),
1612 "fourth",
1613 ),
1614 ];
1615 for (got, want, label) in expect {
1616 assert!(
1617 (got - want).abs() <= 1e-12 * want.abs().max(1.0),
1618 "row {row} {label}: got {got:+.15e} want {want:+.15e}"
1619 );
1620 }
1621 }
1622 }
1623
1624 struct OversizedDenseProgram;
1625
1626 impl RowProgram<10> for OversizedDenseProgram {
1627 fn n_rows(&self) -> usize {
1628 1
1629 }
1630
1631 fn primaries(&self, row: usize) -> Result<[f64; 10], String> {
1632 Err(format!(
1633 "dense-tower storage check reached program primaries at row {row}"
1634 ))
1635 }
1636
1637 fn eval<S: crate::jet_scalar::JetScalar<10>>(
1638 &self,
1639 row: usize,
1640 primaries: &[S; 10],
1641 ) -> Result<S, String> {
1642 Err(format!(
1643 "dense-tower storage check reached program evaluation at row {row} with {} primaries",
1644 primaries.len()
1645 ))
1646 }
1647 }
1648
1649 struct LargestBudgetedDenseProgram;
1650
1651 impl RowProgram<9> for LargestBudgetedDenseProgram {
1652 fn n_rows(&self) -> usize {
1653 1
1654 }
1655
1656 fn primaries(&self, row: usize) -> Result<[f64; 9], String> {
1657 if row == 0 {
1658 Ok([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
1659 } else {
1660 Err(format!("largest budgeted dense program has no row {row}"))
1661 }
1662 }
1663
1664 fn eval<S: crate::jet_scalar::JetScalar<9>>(
1665 &self,
1666 row: usize,
1667 primaries: &[S; 9],
1668 ) -> Result<S, String> {
1669 if row != 0 {
1670 return Err(format!("largest budgeted dense program has no row {row}"));
1671 }
1672 let linear =
1673 S::linear_combination(primaries, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
1674 let quartic = primaries[0]
1675 .mul(&primaries[1])
1676 .mul(&primaries[2])
1677 .mul(&primaries[3]);
1678 Ok(linear.add(&quartic))
1679 }
1680 }
1681
1682 #[test]
1683 fn full_tower_accepts_largest_width_inside_storage_budget_932() {
1684 assert_eq!(std::mem::size_of::<Tower4<9>>(), 59_048);
1685 assert!(
1686 !program_primary_jets_fit_stack::<Tower4<9>, 9>(),
1687 "nine full-width primary towers must use exact-length heap storage"
1688 );
1689
1690 let tower = program_full_tower(&LargestBudgetedDenseProgram, 0)
1691 .expect("Tower4<9> must remain inside the canonical dense storage budget");
1692 assert_eq!(tower.v, 309.0);
1693 assert_eq!(tower.g, [25.0, 14.0, 11.0, 10.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
1694 // `t4` stores derivatives, not Taylor coefficients: the distinct-axis
1695 // derivative of p0*p1*p2*p3 is 1, with no 4! normalization.
1696 assert_eq!(tower.t4[0][1][2][3], 1.0);
1697 }
1698
1699 #[test]
1700 fn full_tower_refuses_oversized_result_before_touching_program() {
1701 let tower_bytes = std::mem::size_of::<Tower4<10>>();
1702 assert!(tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES);
1703 assert!(
1704 std::mem::size_of::<Result<Box<Tower4<32>>, String>>()
1705 <= 4 * std::mem::size_of::<usize>(),
1706 "boxed full-tower API must keep its return slot independent of dense tower width"
1707 );
1708
1709 let error = program_full_tower(&OversizedDenseProgram, 0)
1710 .expect_err("Tower4<10> must exceed the canonical dense storage budget");
1711 assert_eq!(
1712 error,
1713 format!(
1714 "canonical dense Tower4<10> requires {tower_bytes} bytes, exceeding the {}-byte \
1715 storage budget; use the bounded row-kernel and directional channel APIs",
1716 PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
1717 )
1718 );
1719 }
1720
1721 /// Gaussian location-scale row NLL, K=2 primaries (η, s = log σ):
1722 /// ℓ = s + ½ e^{−2s} (y − η)². Mixed cross blocks — the #736 fragility
1723 /// shape — all have one-line closed forms here.
1724 struct LocScaleProgram {
1725 eta: Vec<f64>,
1726 s: Vec<f64>,
1727 y: Vec<f64>,
1728 }
1729
1730 impl RowProgram<2> for LocScaleProgram {
1731 fn n_rows(&self) -> usize {
1732 self.eta.len()
1733 }
1734 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
1735 Ok([self.eta[row], self.s[row]])
1736 }
1737 fn eval<S: crate::jet_scalar::JetScalar<2>>(
1738 &self,
1739 row: usize,
1740 p: &[S; 2],
1741 ) -> Result<S, String> {
1742 let r = S::constant(self.y[row]).sub(&p[0]);
1743 Ok(p[1].add(&p[1].scale(-2.0).exp().mul(&r).mul(&r).scale(0.5)))
1744 }
1745 }
1746
1747 #[test]
1748 fn locscale_tower_matches_closed_forms_including_cross_blocks() {
1749 let prog = LocScaleProgram {
1750 eta: vec![0.3, -1.1, 2.0],
1751 s: vec![-0.5, 0.2, 0.8],
1752 y: vec![1.0, -2.0, 2.5],
1753 };
1754 let tol = 1e-12;
1755 for row in 0..prog.n_rows() {
1756 let t = program_full_tower(&prog, row).expect("locscale program");
1757 let r = prog.y[row] - prog.eta[row];
1758 let w = (-2.0 * prog.s[row]).exp();
1759 // (η, s) = indices (0, 1).
1760 let truth_g = [-w * r, 1.0 - w * r * r];
1761 let truth_h = [[w, 2.0 * w * r], [2.0 * w * r, 2.0 * w * r * r]];
1762 // Third tensor: distinct-entry closed forms.
1763 // ∂ηηη = 0, ∂ηηs = −2w, ∂ηss = −4wr, ∂sss = −4wr².
1764 let t3_truth = |a: usize, b: usize, c: usize| -> f64 {
1765 match a + b + c {
1766 0 => 0.0,
1767 1 => -2.0 * w,
1768 2 => -4.0 * w * r,
1769 _ => -4.0 * w * r * r,
1770 }
1771 };
1772 // Fourth tensor: ∂ηηηη = 0, ∂ηηηs = 0? No: d/ds(∂ηηη)=0 ✓;
1773 // ∂ηηss = 4w, ∂ηsss = 8wr, ∂ssss = 8wr².
1774 let t4_truth = |a: usize, b: usize, c: usize, d: usize| -> f64 {
1775 match a + b + c + d {
1776 0 | 1 => 0.0,
1777 2 => 4.0 * w,
1778 3 => 8.0 * w * r,
1779 _ => 8.0 * w * r * r,
1780 }
1781 };
1782 for a in 0..2 {
1783 assert!(
1784 (t.g[a] - truth_g[a]).abs() <= tol * truth_g[a].abs().max(1.0),
1785 "row {row} grad[{a}]"
1786 );
1787 for b in 0..2 {
1788 assert!(
1789 (t.h[a][b] - truth_h[a][b]).abs() <= tol * w.max(1.0) * (1.0 + r.abs()),
1790 "row {row} hess[{a}][{b}]: got {} want {}",
1791 t.h[a][b],
1792 truth_h[a][b]
1793 );
1794 for c in 0..2 {
1795 assert!(
1796 (t.t3[a][b][c] - t3_truth(a, b, c)).abs()
1797 <= tol * 8.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
1798 "row {row} t3[{a}][{b}][{c}]: got {} want {}",
1799 t.t3[a][b][c],
1800 t3_truth(a, b, c)
1801 );
1802 for d in 0..2 {
1803 assert!(
1804 (t.t4[a][b][c][d] - t4_truth(a, b, c, d)).abs()
1805 <= tol * 16.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
1806 "row {row} t4[{a}][{b}][{c}][{d}]: got {} want {}",
1807 t.t4[a][b][c][d],
1808 t4_truth(a, b, c, d)
1809 );
1810 }
1811 }
1812 }
1813 }
1814 // The canonical trait-surface helpers agree with direct contraction.
1815 let dir = [0.7, -1.3];
1816 let third = program_third_contracted(&prog, row, &dir).expect("third");
1817 for a in 0..2 {
1818 for b in 0..2 {
1819 let want = t.t3[a][b][0] * dir[0] + t.t3[a][b][1] * dir[1];
1820 assert!((third[a][b] - want).abs() <= 1e-13 * want.abs().max(1.0));
1821 }
1822 }
1823 }
1824 }
1825
1826 /// FD cross-check on a deliberately gnarly composition (div, sqrt,
1827 /// powf, nested exp/ln) in K=3, where no closed form is consulted:
1828 /// every tower channel is checked against central finite differences
1829 /// of the channel one order below — value→grad, grad→hess, hess→t3,
1830 /// t3→t4 — so each order is independently anchored.
1831 ///
1832 /// The program carries a per-row primary fixture plus a per-row offset
1833 /// `tau[row]` that enters the loss as a constant, so `row` genuinely
1834 /// drives both the seed point and the evaluated expression.
1835 struct GnarlyProgram {
1836 primaries: Vec<[f64; 3]>,
1837 tau: Vec<f64>,
1838 }
1839
1840 impl GnarlyProgram {
1841 fn fixture() -> Self {
1842 Self {
1843 primaries: vec![[0.4, -0.7, 1.2], [-0.9, 0.6, 0.3], [1.1, -0.2, -0.8]],
1844 tau: vec![0.15, -0.35, 0.5],
1845 }
1846 }
1847 }
1848
1849 impl RowProgram<3> for GnarlyProgram {
1850 fn n_rows(&self) -> usize {
1851 self.primaries.len()
1852 }
1853 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
1854 self.primaries
1855 .get(row)
1856 .copied()
1857 .ok_or_else(|| format!("gnarly: row {row} out of range"))
1858 }
1859 fn eval<S: crate::jet_scalar::JetScalar<3>>(
1860 &self,
1861 row: usize,
1862 p: &[S; 3],
1863 ) -> Result<S, String> {
1864 let tau = *self
1865 .tau
1866 .get(row)
1867 .ok_or_else(|| format!("gnarly: tau row {row} out of range"))?;
1868 let a = p[0].mul(&p[1]).exp();
1869 let b = p[2].mul(&p[2]).add(&S::constant(1.0)).sqrt();
1870 let c = a.add(&b).add(&S::constant(tau)).ln();
1871 let d = p[1].scale(0.5).add(&S::constant(2.0)).powf(1.7);
1872 let delta = p[0].sub(&p[2]);
1873 Ok(c.mul(&d.recip()).add(&delta.mul(&delta).scale(0.25)))
1874 }
1875 }
1876
1877 /// Evaluate the gnarly program's tower at an ARBITRARY seed point for
1878 /// `row` (used to drive central differences off the fixture grid),
1879 /// while keeping `row`'s per-row data (`tau`) in the loss.
1880 fn gnarly_tower_at(prog: &GnarlyProgram, row: usize, p: [f64; 3]) -> Tower4<3> {
1881 struct At<'a> {
1882 base: &'a GnarlyProgram,
1883 row: usize,
1884 p: [f64; 3],
1885 }
1886 impl RowProgram<3> for At<'_> {
1887 fn n_rows(&self) -> usize {
1888 1
1889 }
1890 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
1891 if row != 0 {
1892 return Err(format!("gnarly-at: row {row} out of range"));
1893 }
1894 Ok(self.p)
1895 }
1896 fn eval<S: crate::jet_scalar::JetScalar<3>>(
1897 &self,
1898 eval_row: usize,
1899 vars: &[S; 3],
1900 ) -> Result<S, String> {
1901 if eval_row != 0 {
1902 return Err(format!("gnarly-at: eval row {eval_row} out of range"));
1903 }
1904 self.base.eval(self.row, vars)
1905 }
1906 }
1907 *program_full_tower(&At { base: prog, row, p }, 0).expect("gnarly tower")
1908 }
1909
1910 #[test]
1911 fn gnarly_tower_is_fd_consistent_order_by_order() {
1912 let prog = GnarlyProgram::fixture();
1913 for row in 0..prog.n_rows() {
1914 let base = prog.primaries(row).expect("primaries");
1915 let t = gnarly_tower_at(&prog, row, base);
1916 let h_step = 1e-5;
1917 let tol = 1e-6;
1918 for c in 0..3 {
1919 let mut up = base;
1920 let mut dn = base;
1921 up[c] += h_step;
1922 dn[c] -= h_step;
1923 let t_up = gnarly_tower_at(&prog, row, up);
1924 let t_dn = gnarly_tower_at(&prog, row, dn);
1925 // value → gradient.
1926 let fd_g = (t_up.v - t_dn.v) / (2.0 * h_step);
1927 assert!(
1928 (t.g[c] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
1929 "grad[{c}]: analytic {} fd {}",
1930 t.g[c],
1931 fd_g
1932 );
1933 for a in 0..3 {
1934 // gradient → Hessian.
1935 let fd_h = (t_up.g[a] - t_dn.g[a]) / (2.0 * h_step);
1936 assert!(
1937 (t.h[a][c] - fd_h).abs() <= tol * fd_h.abs().max(1.0),
1938 "hess[{a}][{c}]: analytic {} fd {}",
1939 t.h[a][c],
1940 fd_h
1941 );
1942 for b in 0..3 {
1943 // Hessian → third.
1944 let fd_t3 = (t_up.h[a][b] - t_dn.h[a][b]) / (2.0 * h_step);
1945 assert!(
1946 (t.t3[a][b][c] - fd_t3).abs() <= tol * fd_t3.abs().max(1.0),
1947 "t3[{a}][{b}][{c}]: analytic {} fd {}",
1948 t.t3[a][b][c],
1949 fd_t3
1950 );
1951 for d in 0..3 {
1952 // third → fourth.
1953 let fd_t4 = (t_up.t3[a][b][d] - t_dn.t3[a][b][d]) / (2.0 * h_step);
1954 assert!(
1955 (t.t4[a][b][d][c] - fd_t4).abs() <= tol * fd_t4.abs().max(1.0),
1956 "t4[{a}][{b}][{d}][{c}]: analytic {} fd {}",
1957 t.t4[a][b][d][c],
1958 fd_t4
1959 );
1960 }
1961 }
1962 }
1963 }
1964 }
1965 }
1966
1967 /// The survival crossing-edge position tower `z_edge = (τ − a(θ)) / b`,
1968 /// `b = exp(g)`, built from the intercept tower `a(θ)` (here a stand-in)
1969 /// and the seeded slope `g`, reproduces taylor-jet's exact hand-path
1970 /// boundary-velocity formulas:
1971 /// z_u = −(a_u + [u==g]·z) / b
1972 /// z_uv = −(a_uv + [u==g]·z_v + [v==g]·z_u) / b
1973 /// This pins the bridge between `implicit_solve` and
1974 /// `cell_moving_boundary_flux_tower`: the boundary jet that the production
1975 /// flex path hand-codes (and dropped `z_uv` from) is exactly `∂²` of this
1976 /// tower. K=3 reduced frame: slot 0 = a-axis carrier (an arbitrary smooth
1977 /// a(θ) with nonzero a_u/a_uv), slot 1 = g (the slope), slot 2 unused.
1978 #[test]
1979 fn crossing_edge_tower_matches_handpath_velocity_formulas() {
1980 const TAU: f64 = 1.3; // the link-knot crossing threshold τ
1981 let g_idx = 1usize;
1982 let g0 = 0.85_f64; // the slope value b (the g-primary IS the slope)
1983 // Stand-in intercept tower a(θ): nonzero value, gradient, Hessian in the
1984 // two live axes so a_u and a_uv are both exercised. (In production this
1985 // comes from implicit_solve; here we plant known derivatives.)
1986 let mut a = Tower4::<3>::constant(0.45);
1987 a.g[0] = 0.7;
1988 a.g[1] = -0.3;
1989 a.h[0][0] = 0.25;
1990 a.h[0][1] = 0.11;
1991 a.h[1][0] = 0.11;
1992 a.h[1][1] = -0.08;
1993
1994 // In the survival flex frame the slope `b` IS the g-primary directly
1995 // (the directional code passes `g` as `b`, and ∂z/∂g uses ∂b/∂g = 1):
1996 // z_edge = (τ − a) / b with b seeded as the g-axis variable.
1997 let b = Tower4::<3>::variable(g0, g_idx);
1998 let z_edge = (Tower4::<3>::constant(TAU) - a) / b;
1999
2000 let bv = g0;
2001 let z0 = z_edge.v;
2002 assert!((z0 - (TAU - 0.45) / bv).abs() < 1e-12);
2003
2004 // z_u = −(a_u + [u==g]·z) / b.
2005 for u in 0..2 {
2006 let direct = if u == g_idx { z0 } else { 0.0 };
2007 let want = -(a.g[u] + direct) / bv;
2008 assert!(
2009 (z_edge.g[u] - want).abs() < 1e-10,
2010 "z_u[{u}] {:+.8e} vs hand formula {:+.8e}",
2011 z_edge.g[u],
2012 want
2013 );
2014 }
2015 // z_uv = −(a_uv + [u==g]·z_v + [v==g]·z_u) / b, using the tower's own
2016 // first-order z_v/z_u (already verified above).
2017 for u in 0..2 {
2018 for v in 0..2 {
2019 let cross = if u == g_idx { z_edge.g[v] } else { 0.0 }
2020 + if v == g_idx { z_edge.g[u] } else { 0.0 };
2021 let want = -(a.h[u][v] + cross) / bv;
2022 assert!(
2023 (z_edge.h[u][v] - want).abs() < 1e-10,
2024 "z_uv[{u}][{v}] {:+.8e} vs hand formula {:+.8e}",
2025 z_edge.h[u][v],
2026 want
2027 );
2028 }
2029 }
2030 }
2031
2032 /// The crossing-edge tower in the CONSTRAINT frame (intercept `a` and
2033 /// slope `b` BOTH independent — slots 0 and 1) reproduces taylor-jet's
2034 /// FD-certified bare boundary-velocity constants exactly:
2035 /// z_a = ∂z/∂a = −1/b
2036 /// z_ab = ∂²z/∂a∂b = +1/b²
2037 /// z_aa = ∂²z/∂a² = 0
2038 /// z_bb = ∂²z/∂b² = +2(τ−a)/b³
2039 /// These are the `f_a`/`f_au`/`f_aa` constraint-jet boundary motions the
2040 /// production base path drops (and only adds in the dir twins, causing the
2041 /// #932 desync). Here `a` is independent (NOT yet substituted with a(θ)),
2042 /// so `z_aa = 0` and there is no `a_uv` chain — `implicit_solve` introduces
2043 /// that later. Pins the constant before the constraint-tower wiring.
2044 #[test]
2045 fn crossing_edge_constraint_frame_matches_bare_velocity_constants() {
2046 const TAU: f64 = 1.3;
2047 let a0 = 0.45_f64;
2048 let b0 = 0.85_f64;
2049 // Slot 0 = a, slot 1 = b, both seeded independent.
2050 let a = Tower4::<2>::variable(a0, 0);
2051 let b = Tower4::<2>::variable(b0, 1);
2052 let z = (Tower4::<2>::constant(TAU) - a) / b;
2053
2054 assert!((z.v - (TAU - a0) / b0).abs() < 1e-12);
2055 assert!((z.g[0] - (-1.0 / b0)).abs() < 1e-12, "z_a {:+.10e}", z.g[0]);
2056 assert!(
2057 (z.h[0][1] - 1.0 / (b0 * b0)).abs() < 1e-12,
2058 "z_ab {:+.10e} vs +1/b² {:+.10e}",
2059 z.h[0][1],
2060 1.0 / (b0 * b0)
2061 );
2062 assert!(
2063 z.h[0][0].abs() < 1e-12,
2064 "z_aa must vanish, got {:+.10e}",
2065 z.h[0][0]
2066 );
2067 let want_zbb = 2.0 * (TAU - a0) / (b0 * b0 * b0);
2068 assert!(
2069 (z.h[1][1] - want_zbb).abs() < 1e-12,
2070 "z_bb {:+.10e} vs 2(τ−a)/b³ {:+.10e}",
2071 z.h[1][1],
2072 want_zbb
2073 );
2074 }
2075
2076 /// The third- and fourth-order tensors must be FULLY symmetric under
2077 /// index permutation (mixed partials commute). The tower stores them
2078 /// unsymmetrized, so equal-by-construction is a real invariant of the
2079 /// Leibniz/Faà di Bruno writes — a cheap typo tripwire. Asserted on a
2080 /// nontrivial K=3 tower with all of div/sqrt/powf/exp/ln exercised, so
2081 /// every composition path contributes. Lives in a test (not the hot
2082 /// per-op path) on purpose.
2083 #[test]
2084 fn t3_t4_are_fully_index_symmetric() {
2085 let prog = GnarlyProgram::fixture();
2086 // 3! = 6 permutations of three indices.
2087 let perms3: [[usize; 3]; 6] = [
2088 [0, 1, 2],
2089 [0, 2, 1],
2090 [1, 0, 2],
2091 [1, 2, 0],
2092 [2, 0, 1],
2093 [2, 1, 0],
2094 ];
2095 // 4! = 24 permutations of four indices.
2096 let perms4: [[usize; 4]; 24] = [
2097 [0, 1, 2, 3],
2098 [0, 1, 3, 2],
2099 [0, 2, 1, 3],
2100 [0, 2, 3, 1],
2101 [0, 3, 1, 2],
2102 [0, 3, 2, 1],
2103 [1, 0, 2, 3],
2104 [1, 0, 3, 2],
2105 [1, 2, 0, 3],
2106 [1, 2, 3, 0],
2107 [1, 3, 0, 2],
2108 [1, 3, 2, 0],
2109 [2, 0, 1, 3],
2110 [2, 0, 3, 1],
2111 [2, 1, 0, 3],
2112 [2, 1, 3, 0],
2113 [2, 3, 0, 1],
2114 [2, 3, 1, 0],
2115 [3, 0, 1, 2],
2116 [3, 0, 2, 1],
2117 [3, 1, 0, 2],
2118 [3, 1, 2, 0],
2119 [3, 2, 0, 1],
2120 [3, 2, 1, 0],
2121 ];
2122 for row in 0..prog.n_rows() {
2123 let t = program_full_tower(&prog, row).expect("gnarly tower");
2124 let scale_t3 =
2125 t.t3.iter()
2126 .flatten()
2127 .flatten()
2128 .fold(0.0_f64, |m, x| m.max(x.abs()))
2129 .max(1.0);
2130 let scale_t4 =
2131 t.t4.iter()
2132 .flatten()
2133 .flatten()
2134 .flatten()
2135 .fold(0.0_f64, |m, x| m.max(x.abs()))
2136 .max(1.0);
2137 for i in 0..3 {
2138 for j in 0..3 {
2139 for k in 0..3 {
2140 let base = t.t3[i][j][k];
2141 let idx = [i, j, k];
2142 for p in &perms3 {
2143 let permed = t.t3[idx[p[0]]][idx[p[1]]][idx[p[2]]];
2144 assert!(
2145 (base - permed).abs() <= 1e-12 * scale_t3,
2146 "row {row}: t3[{i}][{j}][{k}]={base:+.15e} != \
2147 permuted {permed:+.15e} under {p:?}"
2148 );
2149 }
2150 for l in 0..3 {
2151 let base4 = t.t4[i][j][k][l];
2152 let idx4 = [i, j, k, l];
2153 for p in &perms4 {
2154 let permed = t.t4[idx4[p[0]]][idx4[p[1]]][idx4[p[2]]][idx4[p[3]]];
2155 assert!(
2156 (base4 - permed).abs() <= 1e-12 * scale_t4,
2157 "row {row}: t4[{i}][{j}][{k}][{l}]={base4:+.15e} != \
2158 permuted {permed:+.15e} under {p:?}"
2159 );
2160 }
2161 }
2162 }
2163 }
2164 }
2165 }
2166 }
2167}
2168
2169#[cfg(test)]
2170mod derivative_stack_tests {
2171 use super::*;
2172 // ── ln_gamma_derivative_stack / digamma_derivative_stack / trigamma_derivative_stack ──
2173
2174 #[test]
2175 fn ln_gamma_derivative_stack_known_values_at_1() {
2176 let s = ln_gamma_derivative_stack(1.0);
2177 // ln Γ(1) = 0; statrs uses Lanczos so the result is within ULP noise
2178 assert!(s[0].abs() < 1e-14, "ln_gamma(1) must be ~0, got {}", s[0]);
2179 // ψ₀(1) = -γ (Euler–Mascheroni)
2180 let euler_mascheroni = 0.577_215_664_901_532_9_f64;
2181 assert!(
2182 (s[1] + euler_mascheroni).abs() < 1e-10,
2183 "digamma(1) ≈ -{euler_mascheroni:.6}, got {}",
2184 s[1]
2185 );
2186 // ψ₁(1) = π²/6
2187 let pi2_6 = std::f64::consts::PI * std::f64::consts::PI / 6.0;
2188 assert!(
2189 (s[2] - pi2_6).abs() < 1e-10,
2190 "trigamma(1) ≈ {pi2_6:.6}, got {}",
2191 s[2]
2192 );
2193 }
2194
2195 #[test]
2196 fn ln_gamma_derivative_stack_known_values_at_2() {
2197 let s = ln_gamma_derivative_stack(2.0);
2198 // ln Γ(2) = ln(1) = 0 exactly
2199 assert!(s[0].abs() < 1e-14, "ln_gamma(2) must be 0, got {}", s[0]);
2200 // ψ₀(2) = 1 − γ (recurrence: ψ₀(x+1) = ψ₀(x) + 1/x)
2201 let euler_mascheroni = 0.577_215_664_901_532_9_f64;
2202 let digamma_2 = 1.0 - euler_mascheroni;
2203 assert!(
2204 (s[1] - digamma_2).abs() < 1e-10,
2205 "digamma(2) ≈ {digamma_2:.6}, got {}",
2206 s[1]
2207 );
2208 }
2209
2210 #[test]
2211 fn ln_gamma_derivative_stack_order2_is_prefix() {
2212 for &x in &[1.0e-8_f64, 0.5, 1.0, 2.0, 5.0, 20.0, 1.0e8] {
2213 let full = ln_gamma_derivative_stack(x);
2214 let ord2 = ln_gamma_derivative_stack_order2(x);
2215 let ord3 = ln_gamma_derivative_stack_order3(x);
2216 assert_eq!(ord2[0], full[0], "order2[0] != full[0] at x={x}");
2217 assert_eq!(ord2[1], full[1], "order2[1] != full[1] at x={x}");
2218 assert_eq!(ord2[2], full[2], "order2[2] != full[2] at x={x}");
2219 assert_eq!(&ord3, &full[..4], "order3 prefix differs at x={x}");
2220 }
2221 }
2222
2223 #[test]
2224 fn digamma_derivative_stack_overlaps_ln_gamma_stack() {
2225 // The two stacks share a run of four polygamma values:
2226 // ln_gamma_stack[1..5] == digamma_stack[0..4]
2227 for &x in &[0.5_f64, 1.0, 2.0, 7.0] {
2228 let lg = ln_gamma_derivative_stack(x);
2229 let dg = digamma_derivative_stack(x);
2230 for i in 0..4 {
2231 assert_eq!(
2232 lg[i + 1],
2233 dg[i],
2234 "ln_gamma_stack[{}] != digamma_stack[{}] at x={x}",
2235 i + 1,
2236 i
2237 );
2238 }
2239 }
2240 }
2241}
2242
2243// ── Contraction-symmetry optimization gate ────────────────────────────────────
2244//
2245// `Tower4::third_contracted` / `fourth_contracted` contract the (fully
2246// index-symmetric) `t3`/`t4` tensors against directions, leaving the output
2247// indices `(a, b)` / `(i, j)` free. Those free indices inherit the tensor's
2248// symmetry — `out[a][b] == out[b][a]` term-for-term — so only the upper triangle
2249// need be summed and the lower triangle mirrored. Unlike the dense symmetric
2250// FILL (which needs a K⁴ scatter and loses inner-loop vectorisation, and was
2251// measured SLOWER), the mirror here is a tiny K×K copy and the inner contraction
2252// is untouched (contiguous, vectorisable). This is BIT-IDENTICAL to the full
2253// nest, so it needs no fingerprint re-baseline; the gate is (1) bit-identity vs
2254// the full reference and (2) a measured wall-clock that is not slower.
2255#[cfg(test)]
2256mod contraction_symmetry_tests {
2257 use super::*;
2258
2259 struct Rng(u64);
2260 impl Rng {
2261 fn u(&mut self) -> f64 {
2262 self.0 = self
2263 .0
2264 .wrapping_mul(6364136223846793005)
2265 .wrapping_add(1442695040888963407);
2266 (self.0 >> 11) as f64 / (1u64 << 53) as f64
2267 }
2268 fn s(&mut self) -> f64 {
2269 (self.u() - 0.5) * 4.0
2270 }
2271 }
2272
2273 /// Random VALID fully-symmetric `Tower4<K>` (symmetric `h`/`t3`/`t4`).
2274 fn rand_sym4<const K: usize>(r: &mut Rng) -> Tower4<K> {
2275 let mut t = Tower4::<K>::zero();
2276 t.v = r.s();
2277 for i in 0..K {
2278 t.g[i] = r.s();
2279 }
2280 for a in 0..K {
2281 for b in a..K {
2282 let v2 = r.s();
2283 t.h[a][b] = v2;
2284 t.h[b][a] = v2;
2285 for c in b..K {
2286 let v3 = r.s();
2287 for p in perms3([a, b, c]) {
2288 t.t3[p[0]][p[1]][p[2]] = v3;
2289 }
2290 for d in c..K {
2291 let v4 = r.s();
2292 for p in perms4([a, b, c, d]) {
2293 t.t4[p[0]][p[1]][p[2]][p[3]] = v4;
2294 }
2295 }
2296 }
2297 }
2298 }
2299 t
2300 }
2301
2302 fn perms3(idx: [usize; 3]) -> [[usize; 3]; 6] {
2303 let [a, b, c] = idx;
2304 [
2305 [a, b, c],
2306 [a, c, b],
2307 [b, a, c],
2308 [b, c, a],
2309 [c, a, b],
2310 [c, b, a],
2311 ]
2312 }
2313 fn perms4(idx: [usize; 4]) -> [[usize; 4]; 24] {
2314 let [a, b, c, d] = idx;
2315 [
2316 [a, b, c, d],
2317 [a, b, d, c],
2318 [a, c, b, d],
2319 [a, c, d, b],
2320 [a, d, b, c],
2321 [a, d, c, b],
2322 [b, a, c, d],
2323 [b, a, d, c],
2324 [b, c, a, d],
2325 [b, c, d, a],
2326 [b, d, a, c],
2327 [b, d, c, a],
2328 [c, a, b, d],
2329 [c, a, d, b],
2330 [c, b, a, d],
2331 [c, b, d, a],
2332 [c, d, a, b],
2333 [c, d, b, a],
2334 [d, a, b, c],
2335 [d, a, c, b],
2336 [d, b, a, c],
2337 [d, b, c, a],
2338 [d, c, a, b],
2339 [d, c, b, a],
2340 ]
2341 }
2342
2343 /// Full-nest reference (the pre-opt `a, b ∈ 0..K` form).
2344 fn third_full<const K: usize>(t: &Tower4<K>, dir: &[f64; K]) -> [[f64; K]; K] {
2345 let mut out = [[0.0; K]; K];
2346 for a in 0..K {
2347 for b in 0..K {
2348 let mut acc = 0.0;
2349 for c in 0..K {
2350 acc += t.t3[a][b][c] * dir[c];
2351 }
2352 out[a][b] = acc;
2353 }
2354 }
2355 out
2356 }
2357 fn fourth_full<const K: usize>(t: &Tower4<K>, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
2358 let mut out = [[0.0; K]; K];
2359 for i in 0..K {
2360 for j in 0..K {
2361 let mut acc = 0.0;
2362 for k in 0..K {
2363 for l in 0..K {
2364 acc += t.t4[i][j][k][l] * u[k] * w[l];
2365 }
2366 }
2367 out[i][j] = acc;
2368 }
2369 }
2370 out
2371 }
2372
2373 /// Returns the number of bit-equality comparisons performed (`n·K·K·2`), so
2374 /// the caller can assert the intended workload actually ran: a generic
2375 /// (turbofish) helper call hides its internal assertions, so the count is
2376 /// surfaced and checked at the call site.
2377 fn check_bit_identical<const K: usize>(seed: u64, n: usize) -> usize {
2378 let mut r = Rng(seed);
2379 let mut checks = 0usize;
2380 for _ in 0..n {
2381 let t = rand_sym4::<K>(&mut r);
2382 let dir: [f64; K] = std::array::from_fn(|_| r.s());
2383 let u: [f64; K] = std::array::from_fn(|_| r.s());
2384 let w: [f64; K] = std::array::from_fn(|_| r.s());
2385 let t3_sym = t.third_contracted(&dir);
2386 let t3_full = third_full(&t, &dir);
2387 let t4_sym = t.fourth_contracted(&u, &w);
2388 let t4_full = fourth_full(&t, &u, &w);
2389 for a in 0..K {
2390 for b in 0..K {
2391 assert_eq!(
2392 t3_sym[a][b].to_bits(),
2393 t3_full[a][b].to_bits(),
2394 "third K={K} [{a}][{b}]"
2395 );
2396 assert_eq!(
2397 t4_sym[a][b].to_bits(),
2398 t4_full[a][b].to_bits(),
2399 "fourth K={K} [{a}][{b}]"
2400 );
2401 checks += 2;
2402 }
2403 }
2404 }
2405 checks
2406 }
2407
2408 /// The output-symmetric contraction is BIT-IDENTICAL to the full nest across
2409 /// `K ∈ {2,3,4,9}` (so no fingerprint re-baseline is owed — accuracy and bits
2410 /// are unchanged; this is a pure speed-only optimization).
2411 #[test]
2412 fn contraction_symmetry_is_bit_identical_to_full_nest() {
2413 let checks = check_bit_identical::<2>(0x0000_0002_C0FF_EE01, 1000)
2414 + check_bit_identical::<3>(0x0000_0003_C0FF_EE01, 800)
2415 + check_bit_identical::<4>(0x0000_0004_C0FF_EE01, 600)
2416 + check_bit_identical::<9>(0x0000_0009_C0FF_EE01, 300);
2417 // Guards against the loops silently not running (e.g. a zeroed count):
2418 // 1000·2²·2 + 800·3²·2 + 600·4²·2 + 300·9²·2.
2419 assert_eq!(checks, 8000 + 14400 + 19200 + 48600);
2420 }
2421
2422 /// The output-symmetric contraction at `K = 9` does strictly fewer inner
2423 /// contractions than the full nest (~2x), so it must be the faster arm.
2424 /// The bit-identity test above is the correctness gate; this is the speed
2425 /// contract, and it opens only in the release profile (`SpeedGate::open`
2426 /// documents why -- this gate once PASSED on a quiet node and FAILED at
2427 /// 1.62x on a loaded one, because its two arms were timed in separate
2428 /// windows in a fixed order; the paired harness times them adjacent, in a
2429 /// randomised order, and reports its own resolution).
2430 #[test]
2431 fn contraction_symmetry_speedup_is_reported() {
2432 use crate::paired_timing::{paired_interleaved, SpeedGate};
2433
2434 const K: usize = 9;
2435 let mut r = Rng(0xC0FF_EE99_1234_5678);
2436 let towers: Vec<Tower4<K>> = (0..512).map(|_| rand_sym4::<K>(&mut r)).collect();
2437 let dir: [f64; K] = std::array::from_fn(|_| r.s());
2438 let u: [f64; K] = std::array::from_fn(|_| r.s());
2439 let w: [f64; K] = std::array::from_fn(|_| r.s());
2440
2441 if cfg!(debug_assertions) {
2442 return;
2443 }
2444 let mut gate = SpeedGate::open("CONTRACTION-SYMMETRY-932");
2445 // One arm call contracts every tower once; the nudge perturbs both
2446 // directions so neither contraction is loop-invariant across calls.
2447 let timing = paired_interleaved(
2448 15,
2449 20,
2450 0x9320_5E11,
2451 |nudge| {
2452 let mut dir = dir;
2453 dir[0] += nudge;
2454 let mut u = u;
2455 u[0] += nudge;
2456 let mut sink = 0.0f64;
2457 for t in &towers {
2458 let o3 = t.third_contracted(&dir);
2459 let o4 = t.fourth_contracted(&u, &w);
2460 sink += o3[0][K - 1] + o4[0][K - 1];
2461 }
2462 sink
2463 },
2464 |nudge| {
2465 let mut dir = dir;
2466 dir[0] += nudge;
2467 let mut u = u;
2468 u[0] += nudge;
2469 let mut sink = 0.0f64;
2470 for t in &towers {
2471 let o3 = third_full(t, &dir);
2472 let o4 = fourth_full(t, &u, &w);
2473 sink += o3[0][K - 1] + o4[0][K - 1];
2474 }
2475 sink
2476 },
2477 );
2478 gate.faster(
2479 &format!("K={K} towers={}", towers.len()),
2480 &timing,
2481 "symmetric",
2482 "full_nest",
2483 );
2484 gate.finish();
2485 }
2486}