Skip to main content

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}