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