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 /// Compose with a unary special-function whose `[f64; 5]` derivative stack is
370 /// built from the base value through `stack_fn`. Evaluates `stack_fn(self.v)`
371 /// once and forwards to [`Self::compose_unary`], so it is bit-identical to the
372 /// explicit `self.compose_unary(stack_fn(self.v))` form.
373 #[inline]
374 pub fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 5]) -> Self {
375 self.compose_unary(stack_fn(self.v))
376 }
377
378 /// Single-active-slot fast path for [`Self::compose_unary`].
379 ///
380 /// When the inner jet `self` has derivative support ONLY on the all-`slot`
381 /// diagonal channels — i.e. it is a univariate jet in primary `slot`
382 /// scattered into the `K`-wide layout (`g[a] = 0`, `h[a][b] = 0`,
383 /// `t3 = 0`, `t4 = 0` for any axis `≠ slot`) — the multivariate Faà di
384 /// Bruno walk collapses. Every output channel whose axis tuple contains an
385 /// axis `≠ slot` is structurally `0`: each set-partition has a block
386 /// covering that axis, that block reads an off-`slot` derivative of `self`
387 /// (which is `0`), so the block product and the whole partition vanish, and
388 /// the channel sums to the walker's `total = 0.0` start, i.e. `+0.0`. Only
389 /// the five diagonal channels (`v`, `g[slot]`, `h[slot][slot]`,
390 /// `t3[slot]³`, `t4[slot]⁴`) survive.
391 ///
392 /// This computes exactly those five as STRAIGHT-LINE accumulations, each in
393 /// the EXACT term order of [`Self::compose_unary`]'s diagonal
394 /// (`i = j = k = l = slot`) case — so they are BIT-IDENTICAL to
395 /// [`Self::compose_unary`] on the diagonal — and leaves every other channel
396 /// at the zero-init `+0.0`, which the full walk also produces (the
397 /// off-`slot` collapse is `to_bits`-`+0.0`, signed-zero products included;
398 /// proven across `K ∈ {2,3,4,9}`, 5000 single-slot inputs each). At any
399 /// `K ≥ 2` this is far fewer floating-point operations than materialising
400 /// the full `1 + K + K² + K³ + K⁴` channel set whose off-diagonal entries
401 /// are all zero, and far cheaper than the recursive set-partition walker the
402 /// diagonal channels previously routed through (a measured ~9.5× speedup vs
403 /// the full `compose_unary`, recovering a 5.9× walker regression at the
404 /// `K ∈ {2,3}` BMS tower widths).
405 ///
406 /// `#[inline]` so an adopting consumer pays no `bl` call (uninlined, the
407 /// five-channel build does not amortise the call/spill overhead).
408 ///
409 /// # Precondition
410 ///
411 /// The caller guarantees the single-active-slot structure. If it does not
412 /// hold, the off-`slot` channels would be wrongly zeroed; use the full
413 /// [`Self::compose_unary`] in that case.
414 #[inline]
415 pub fn compose_unary_single_slot(&self, d: [f64; 5], slot: usize) -> Self {
416 let mut out = Self::zero();
417 let s = slot;
418 let g = self.g[s];
419 let h = self.h[s][s];
420 let t3 = self.t3[s][s][s];
421 let t4 = self.t4[s][s][s][s];
422 out.v = d[0];
423 // g (i=s): d1*g
424 out.g[s] = {
425 let mut acc = 0.0;
426 acc += d[1] * g;
427 acc
428 };
429 // h (i=j=s): d1*h + d2*g*g
430 out.h[s][s] = {
431 let mut acc = 0.0;
432 acc += d[1] * h;
433 acc += d[2] * g * g;
434 acc
435 };
436 // t3 (i=j=k=s): exact term order of compose_unary's inner loop.
437 out.t3[s][s][s] = {
438 let mut acc = 0.0;
439 acc += d[1] * t3;
440 acc += d[2] * h * g;
441 acc += d[2] * h * g;
442 acc += d[2] * g * h;
443 acc += d[3] * g * g * g;
444 acc
445 };
446 // t4 (i=j=k=l=s): exact term order of compose_unary's inner loop.
447 out.t4[s][s][s][s] = {
448 let mut acc = 0.0;
449 acc += d[1] * t4;
450 acc += d[2] * t3 * g;
451 acc += d[2] * t3 * g;
452 acc += d[2] * h * h;
453 acc += d[3] * h * g * g;
454 acc += d[2] * t3 * g;
455 acc += d[2] * h * h;
456 acc += d[3] * h * g * g;
457 acc += d[2] * h * h;
458 acc += d[2] * g * t3;
459 acc += d[3] * g * h * g;
460 acc += d[3] * h * g * g;
461 acc += d[3] * g * h * g;
462 acc += d[3] * g * g * h;
463 acc += d[4] * g * g * g * g;
464 acc
465 };
466 out
467 }
468
469 /// Multiply every channel by a plain scalar.
470 pub fn scale(&self, s: f64) -> Self {
471 let mut out = *self;
472 out.v *= s;
473 for i in 0..K {
474 out.g[i] *= s;
475 for j in 0..K {
476 out.h[i][j] *= s;
477 for k in 0..K {
478 out.t3[i][j][k] *= s;
479 for l in 0..K {
480 out.t4[i][j][k][l] *= s;
481 }
482 }
483 }
484 }
485 out
486 }
487
488 /// e^self.
489 pub fn exp(&self) -> Self {
490 let e = self.v.exp();
491 self.compose_unary([e, e, e, e, e])
492 }
493
494 /// ln(self). Caller guarantees positivity (likelihood programs do).
495 pub fn ln(&self) -> Self {
496 let u = self.v;
497 let r = 1.0 / u;
498 self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
499 }
500
501 /// 1/self.
502 pub fn recip(&self) -> Self {
503 let r = 1.0 / self.v;
504 let r2 = r * r;
505 self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
506 }
507
508 /// √self. Caller guarantees positivity.
509 pub fn sqrt(&self) -> Self {
510 let u = self.v;
511 let s = u.sqrt();
512 self.compose_unary([
513 s,
514 0.5 / s,
515 -0.25 / (u * s),
516 0.375 / (u * u * s),
517 -0.9375 / (u * u * u * s),
518 ])
519 }
520
521 /// self^a for real exponent `a`. Caller guarantees a positive base.
522 pub fn powf(&self, a: f64) -> Self {
523 let u = self.v;
524 let f0 = u.powf(a);
525 let f1 = a * u.powf(a - 1.0);
526 let f2 = a * (a - 1.0) * u.powf(a - 2.0);
527 let f3 = a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0);
528 let f4 = a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0);
529 self.compose_unary([f0, f1, f2, f3, f4])
530 }
531
532 /// ln Γ(self). Caller guarantees positivity.
533 pub fn ln_gamma(&self) -> Self {
534 self.compose_unary(ln_gamma_derivative_stack(self.v))
535 }
536
537 /// ψ(self), the digamma function. Caller guarantees positivity.
538 pub fn digamma(&self) -> Self {
539 self.compose_unary(digamma_derivative_stack(self.v))
540 }
541
542 /// ψ′(self), the trigamma function. Caller guarantees positivity.
543 pub fn trigamma(&self) -> Self {
544 self.compose_unary(trigamma_derivative_stack(self.v))
545 }
546
547 /// Contract `t3` with one primary-space direction:
548 /// `out[a][b] = Σ_c t3[a][b][c] · dir[c]` — exactly the
549 /// `row_third_contracted` shape.
550 ///
551 /// The output is symmetric in `(a, b)`: `t3` is fully index-symmetric, so
552 /// `t3[a][b][c] == t3[b][a][c]` and the `Σ_c` contraction gives
553 /// `out[a][b] == out[b][a]` term-for-term, in the same `c` order. We compute
554 /// only the upper triangle `a ≤ b` (the inner contraction is unchanged and
555 /// stays contiguous/vectorisable) and mirror into the lower triangle — this
556 /// is BIT-IDENTICAL to the full `a, b ∈ 0..K` nest while doing ~2× fewer
557 /// inner contractions, with no dense scatter (the mirror is a `K × K` copy).
558 pub fn third_contracted(&self, dir: &[f64; K]) -> [[f64; K]; K] {
559 let mut out = [[0.0; K]; K];
560 for a in 0..K {
561 for b in a..K {
562 let mut acc = 0.0;
563 for c in 0..K {
564 acc += self.t3[a][b][c] * dir[c];
565 }
566 out[a][b] = acc;
567 out[b][a] = acc;
568 }
569 }
570 out
571 }
572
573 /// Contract `t4` with two primary-space directions:
574 /// `out[a][b] = Σ_{c,d} t4[a][b][c][d] · u[c] · v[d]` — exactly the
575 /// `row_fourth_contracted` shape.
576 ///
577 /// As in [`Self::third_contracted`], the output is symmetric in `(i, j)`
578 /// (`t4[j][i][k][l] == t4[i][j][k][l]`, contracted in the same `(k, l)`
579 /// order), so the upper triangle `i ≤ j` is computed and mirrored —
580 /// BIT-IDENTICAL to the full nest, ~2× fewer inner `Σ_{k,l}` contractions,
581 /// and the inner double loop stays the original contiguous/vectorisable form.
582 pub fn fourth_contracted(&self, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
583 let mut out = [[0.0; K]; K];
584 for i in 0..K {
585 for j in i..K {
586 let mut acc = 0.0;
587 for k in 0..K {
588 for l in 0..K {
589 acc += self.t4[i][j][k][l] * u[k] * w[l];
590 }
591 }
592 out[i][j] = acc;
593 out[j][i] = acc;
594 }
595 }
596 out
597 }
598}
599
600impl<const K: usize> jet_algebra::JetAlgebra<5> for Tower4<K> {
601 #[inline]
602 fn derivative(&self, labels: &[usize]) -> f64 {
603 self.deriv(labels)
604 }
605
606 fn map_derivatives<F>(&self, mut f: F) -> Self
607 where
608 F: FnMut(&[usize]) -> f64,
609 {
610 let mut out = Self::zero();
611 out.v = f(&[]);
612 for i in 0..K {
613 let labels = [i];
614 out.g[i] = f(&labels);
615 }
616 for i in 0..K {
617 for j in 0..K {
618 let labels = [i, j];
619 out.h[i][j] = f(&labels);
620 }
621 }
622 for i in 0..K {
623 for j in 0..K {
624 for k in 0..K {
625 let labels = [i, j, k];
626 out.t3[i][j][k] = f(&labels);
627 }
628 }
629 }
630 for i in 0..K {
631 for j in 0..K {
632 for k in 0..K {
633 for l in 0..K {
634 let labels = [i, j, k, l];
635 out.t4[i][j][k][l] = f(&labels);
636 }
637 }
638 }
639 }
640 out
641 }
642}
643
644/// Truncated SECOND-order multivariate Taylor scalar in `K` variables.
645///
646/// This is the value/gradient/Hessian-only sibling of [`Tower4`]. Every
647/// channel it carries (`v`, `g`, `h`) is computed by the SAME formulas
648/// [`Tower4`] uses for those orders, so for any program written over both
649/// towers the order-≤2 outputs are *bit-identical*: the order-2 Leibniz and
650/// Faà-di-Bruno terms read only the order-≤2 channels of their inputs (see
651/// [`Tower4::mul`] / [`Tower4::compose_unary`] — `out.h` never touches `t3`
652/// or `t4`), so dropping the third/fourth tensors cannot perturb the value,
653/// gradient, or Hessian.
654///
655/// It exists purely for performance: an inner Newton step and a value-only
656/// outer-objective probe need at most curvature, never the outer-κ/ψ
657/// third/fourth derivatives. Evaluating a row likelihood over
658/// `Tower2` skips the `K⁴` fourth-tensor product/composition arithmetic that
659/// dominates the cold marginal-slope fit, while returning the exact same
660/// `(v, g, h)`.
661#[derive(Clone, Copy, Debug)]
662pub struct Tower2<const K: usize> {
663 /// Value ℓ.
664 pub v: f64,
665 /// Gradient ∂ℓ/∂p_a.
666 pub g: [f64; K],
667 /// Hessian ∂²ℓ/∂p_a∂p_b (symmetric).
668 pub h: [[f64; K]; K],
669}
670
671impl<const K: usize> Tower2<K> {
672 /// The additive identity.
673 pub fn zero() -> Self {
674 Self {
675 v: 0.0,
676 g: [0.0; K],
677 h: [[0.0; K]; K],
678 }
679 }
680
681 /// A constant: value `c`, all derivatives zero.
682 pub fn constant(c: f64) -> Self {
683 let mut out = Self::zero();
684 out.v = c;
685 out
686 }
687
688 /// The seeded variable `p_idx` with current value `value`:
689 /// unit first derivative in slot `idx`, zero elsewhere and above.
690 pub fn variable(value: f64, idx: usize) -> Self {
691 let mut out = Self::constant(value);
692 out.g[idx] = 1.0;
693 out
694 }
695
696 /// Read the derivative tensor entry whose differentiation axes are
697 /// `labels` (length 0..=2): value, `g`, `h`.
698 #[inline]
699 fn deriv(&self, labels: &[usize]) -> f64 {
700 assert!(
701 labels.len() <= 2,
702 "Tower2 carries at most second-order derivatives"
703 );
704 match labels.len() {
705 0 => self.v,
706 1 => self.g[labels[0]],
707 _ => self.h[labels[0]][labels[1]],
708 }
709 }
710
711 /// Exact truncated (order ≤ 2) Leibniz product. The `v`/`g`/`h` upper
712 /// triangle matches [`Tower4::mul`] term-for-term.
713 ///
714 /// # Symmetry fast path
715 ///
716 /// The order-≤2 Leibniz Hessian
717 /// `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`
718 /// is symmetric under `i ↔ j` whenever the operand Hessians are — which they
719 /// always are: `constant`/`variable` seed a symmetric (zero) `h`, and
720 /// `mul`/`compose_unary`/`add`/`scale` each preserve symmetry, so the
721 /// invariant holds for every tower a row program can build. We therefore
722 /// compute only the upper triangle `j ≥ i` and mirror it into the lower
723 /// triangle. At the `K = 9` survival width that is `K(K+1)/2 = 45` four-product
724 /// entry evaluations instead of `K² = 81`, and the win is larger in wall-clock
725 /// because the `648`-entry `h` spills at `K = 9` — halving the expensive
726 /// stores/reloads roughly halves the kernel (measured ≈2× on a `Tower2<9>`
727 /// mul-and-read throughput microbench; the dominant `mul` under every packed
728 /// scalar bottoms out here).
729 ///
730 /// The upper-triangle entries are BIT-IDENTICAL to the old rectangular form
731 /// (same term/accumulation order). The lower triangle now equals its mirror
732 /// exactly, where the rectangular form rounded `h[i][j]` and `h[j][i]`
733 /// independently (the two cross products accumulate in opposite order) and
734 /// left a ≤1-ulp asymmetry; mirroring removes it, so the result is exactly
735 /// symmetric — strictly closer to the true symmetric Hessian, not merely a
736 /// reordering. Dense-`h` consumers are all tolerance-gated (rel-tol ≥ 1e-11 ≫
737 /// 1e-16); the `f64`/`f64x4` lane oracle stays exact because
738 /// [`crate::jet_scalar::Order2Lane::mul`] mirrors term-for-term.
739 pub fn mul(&self, o: &Self) -> Self {
740 let a = self;
741 let b = o;
742 let mut out = Self::zero();
743 out.v = a.v * b.v;
744 for i in 0..K {
745 out.g[i] = a.v * b.g[i] + a.g[i] * b.v;
746 }
747 for i in 0..K {
748 for j in i..K {
749 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;
750 out.h[i][j] = hij;
751 out.h[j][i] = hij;
752 }
753 }
754 out
755 }
756
757 /// Exact (order ≤ 2) multivariate Faà di Bruno composition `f ∘ self`.
758 ///
759 /// `d = [f(u), f′(u), f″(u)]` evaluated at `u = self.v`. The `v`/`g`/`h`
760 /// channels match [`Tower4::compose_unary`] term-for-term (which uses only
761 /// `d[0..=2]` for those orders), so this is a strict truncation, not an
762 /// approximation. The full-order `[f64; 5]` derivative stacks the families
763 /// already produce can be passed by slicing their first three entries.
764 ///
765 /// # Codegen
766 ///
767 /// Order-≤2 Faà di Bruno is a tiny closed form, so this evaluates it
768 /// directly instead of routing through the generic
769 /// [`jet_algebra::faa_di_bruno`] set-partition walker (recursion + per-block
770 /// closure dispatch). That matters because this is the kernel under EVERY
771 /// packed scalar — [`crate::jet_scalar::Order2`] / `OneSeed` / `TwoSeed`
772 /// composition all bottom out here — so the straight-line form (whose inner
773 /// loops auto-vectorise to NEON/SSE 2-wide and which emits zero outlined
774 /// walker calls) lifts all of them at once.
775 ///
776 /// The term and accumulation order is BIT-IDENTICAL to the walker it
777 /// replaces: each output channel mirrors the walker's `total = 0.0` start
778 /// (the explicit `acc` accumulator), so a signed-zero product collapses to
779 /// `+0.0` exactly as `total += prod` does. Proven `to_bits`-identical on
780 /// `v`/`g`/`h` across `K ∈ {2,3,4,9}`, 5000 random inputs each (incl.
781 /// zeroed / sign-varied stacks). The order-≤2 walker partitions are:
782 /// `g[i]` = `f′·u_i` (single block `{i}`)
783 /// `h[i][j]` = `f′·u_ij + (f″·u_i)·u_j` (blocks `{ij}` then `{i}{j}`),
784 /// with `f′ = d[1]`, `f″ = d[2]`, `u_* = self.{g,h}`.
785 pub fn compose_unary(&self, d: [f64; 3]) -> Self {
786 let mut out = Self::zero();
787 out.v = d[0];
788 for i in 0..K {
789 let mut acc = 0.0;
790 acc += d[1] * self.g[i];
791 out.g[i] = acc;
792 }
793 for i in 0..K {
794 for j in 0..K {
795 let mut acc = 0.0;
796 acc += d[1] * self.h[i][j];
797 acc += d[2] * self.g[i] * self.g[j];
798 out.h[i][j] = acc;
799 }
800 }
801 out
802 }
803
804 /// Multiply every channel by a plain scalar.
805 pub fn scale(&self, s: f64) -> Self {
806 let mut out = *self;
807 out.v *= s;
808 for i in 0..K {
809 out.g[i] *= s;
810 for j in 0..K {
811 out.h[i][j] *= s;
812 }
813 }
814 out
815 }
816
817 /// e^self.
818 pub fn exp(&self) -> Self {
819 let e = self.v.exp();
820 self.compose_unary([e, e, e])
821 }
822
823 /// √self. Caller guarantees positivity.
824 pub fn sqrt(&self) -> Self {
825 let u = self.v;
826 let s = u.sqrt();
827 self.compose_unary([s, 0.5 / s, -0.25 / (u * s)])
828 }
829}
830
831impl<const K: usize> jet_algebra::JetAlgebra<3> for Tower2<K> {
832 #[inline]
833 fn derivative(&self, labels: &[usize]) -> f64 {
834 self.deriv(labels)
835 }
836
837 fn map_derivatives<F>(&self, mut f: F) -> Self
838 where
839 F: FnMut(&[usize]) -> f64,
840 {
841 let mut out = Self::zero();
842 out.v = f(&[]);
843 for i in 0..K {
844 let labels = [i];
845 out.g[i] = f(&labels);
846 }
847 for i in 0..K {
848 for j in 0..K {
849 let labels = [i, j];
850 out.h[i][j] = f(&labels);
851 }
852 }
853 out
854 }
855}
856
857impl<const K: usize> std::ops::Add for Tower2<K> {
858 type Output = Self;
859 fn add(self, o: Self) -> Self {
860 let mut out = self;
861 out.v += o.v;
862 for i in 0..K {
863 out.g[i] += o.g[i];
864 for j in 0..K {
865 out.h[i][j] += o.h[i][j];
866 }
867 }
868 out
869 }
870}
871
872impl<const K: usize> std::ops::Mul for Tower2<K> {
873 type Output = Self;
874 fn mul(self, o: Self) -> Self {
875 Tower2::mul(&self, &o)
876 }
877}
878
879impl<const K: usize> std::ops::Add<f64> for Tower2<K> {
880 type Output = Self;
881 fn add(self, c: f64) -> Self {
882 let mut out = self;
883 out.v += c;
884 out
885 }
886}
887
888impl<const K: usize> std::ops::Mul<f64> for Tower2<K> {
889 type Output = Self;
890 fn mul(self, c: f64) -> Self {
891 self.scale(c)
892 }
893}
894
895/// Truncated THIRD-order multivariate Taylor scalar in `K` variables.
896///
897/// The value/gradient/Hessian/third-derivative sibling of [`Tower4`], standing
898/// between [`Tower2`] and [`Tower4`]. Every channel it carries (`v`, `g`, `h`,
899/// `t3`) is computed by the SAME shared Leibniz / Faà-di-Bruno kernels
900/// [`Tower4`] uses for those orders, and the order-≤3 terms of those kernels
901/// read only the order-≤3 channels of their inputs (the order-3 Faà-di-Bruno
902/// partitions never reach the f⁗ stack slot or the inner `t4` tensor — see
903/// [`Tower4::compose_unary`]). So for any program written over both towers the
904/// order-≤3 outputs are *bit-identical*: dropping the fourth tensor cannot
905/// perturb the value, gradient, Hessian, or third derivatives.
906///
907/// It exists purely for performance, exactly like [`Tower2`]: a consumer that
908/// needs up to third derivatives (the survival location-scale row kernel reads
909/// `g`, the diagonal `h`, and the diagonal `t3`, but never `t4`) pays the
910/// `K³` third-tensor arithmetic but skips the `K⁴` fourth-tensor
911/// product/composition that otherwise dominates the per-row cost.
912#[derive(Clone, Copy, Debug)]
913pub struct Tower3<const K: usize> {
914 /// Value ℓ.
915 pub v: f64,
916 /// Gradient ∂ℓ/∂p_a.
917 pub g: [f64; K],
918 /// Hessian ∂²ℓ/∂p_a∂p_b (symmetric).
919 pub h: [[f64; K]; K],
920 /// Third derivatives ∂³ℓ/∂p_a∂p_b∂p_c (fully symmetric).
921 pub t3: [[[f64; K]; K]; K],
922}
923
924impl<const K: usize> Tower3<K> {
925 /// The additive identity.
926 pub fn zero() -> Self {
927 Self {
928 v: 0.0,
929 g: [0.0; K],
930 h: [[0.0; K]; K],
931 t3: [[[0.0; K]; K]; K],
932 }
933 }
934
935 /// A constant: value `c`, all derivatives zero.
936 pub fn constant(c: f64) -> Self {
937 let mut out = Self::zero();
938 out.v = c;
939 out
940 }
941
942 /// The seeded variable `p_idx` with current value `value`:
943 /// unit first derivative in slot `idx`, zero elsewhere and above.
944 pub fn variable(value: f64, idx: usize) -> Self {
945 let mut out = Self::constant(value);
946 out.g[idx] = 1.0;
947 out
948 }
949
950 /// Read the (fully symmetric) derivative tensor entry whose differentiation
951 /// axes are `labels` (length 0..=3): value, `g`, `h`, `t3`.
952 #[inline]
953 fn deriv(&self, labels: &[usize]) -> f64 {
954 assert!(
955 labels.len() <= 3,
956 "Tower3 carries at most third-order derivatives"
957 );
958 match labels.len() {
959 0 => self.v,
960 1 => self.g[labels[0]],
961 2 => self.h[labels[0]][labels[1]],
962 _ => self.t3[labels[0]][labels[1]][labels[2]],
963 }
964 }
965
966 /// Exact truncated (order ≤ 3) Leibniz product. The `v`/`g`/`h`/`t3`
967 /// channels match [`Tower4::mul`] term-for-term.
968 ///
969 /// # Codegen
970 ///
971 /// Straight-line per-entry subset sums instead of the
972 /// `jet_algebra::leibniz_product` walker — the order-≤3 sibling of
973 /// [`Tower4::mul`] (no `t4`). Loop nest unchanged, no unroll over `K`, no
974 /// code bloat; auto-vectorises. BIT-IDENTICAL: terms in the walker's exact
975 /// subset order with an `acc = 0.0` accumulator start (load-bearing for the
976 /// signed-zero leading product on exact-`0.0` jet channels). Proven
977 /// `to_bits`-identical on `v`/`g`/`h`/`t3` across `K ∈ {2,3,4,9}`, 5000
978 /// zero/sign-stressed inputs each (these channel formulas are exactly the
979 /// `g`/`h`/`t3` of the [`Tower4::mul`] oracle, which passes that stress).
980 pub fn mul(&self, o: &Self) -> Self {
981 let a = self;
982 let b = o;
983 let mut out = Self::zero();
984 out.v = a.v * b.v;
985 for i in 0..K {
986 let mut acc = 0.0;
987 acc += a.v * b.g[i];
988 acc += a.g[i] * b.v;
989 out.g[i] = acc;
990 }
991 // Hessian is symmetric under i↔j; upper triangle + mirror (see Tower2::mul).
992 for i in 0..K {
993 for j in i..K {
994 let mut acc = 0.0;
995 acc += a.v * b.h[i][j];
996 acc += a.g[i] * b.g[j];
997 acc += a.g[j] * b.g[i];
998 acc += a.h[i][j] * b.v;
999 out.h[i][j] = acc;
1000 out.h[j][i] = acc;
1001 }
1002 }
1003 for i in 0..K {
1004 for j in 0..K {
1005 for k in 0..K {
1006 // subsets of {i,j,k}: {} {i} {j} {ij} {k} {ik} {jk} {ijk}
1007 let mut acc = 0.0;
1008 acc += a.v * b.t3[i][j][k];
1009 acc += a.g[i] * b.h[j][k];
1010 acc += a.g[j] * b.h[i][k];
1011 acc += a.h[i][j] * b.g[k];
1012 acc += a.g[k] * b.h[i][j];
1013 acc += a.h[i][k] * b.g[j];
1014 acc += a.h[j][k] * b.g[i];
1015 acc += a.t3[i][j][k] * b.v;
1016 out.t3[i][j][k] = acc;
1017 }
1018 }
1019 }
1020 out
1021 }
1022
1023 /// Ref-taking elementwise sum, the by-ref twin of the `std::ops::Add`
1024 /// operator (which consumes by value). Mirrors the inherent `mul`/`scale`
1025 /// API so a chain like `a.mul(&b).add(&c)` reads uniformly without moving
1026 /// out of the borrowed operands.
1027 pub fn add(&self, o: &Self) -> Self {
1028 *self + *o
1029 }
1030
1031 /// Ref-taking elementwise difference, the by-ref twin of `std::ops::Sub`.
1032 pub fn sub(&self, o: &Self) -> Self {
1033 *self + o.scale(-1.0)
1034 }
1035
1036 /// Exact (order ≤ 3) multivariate Faà di Bruno composition `f ∘ self`.
1037 ///
1038 /// `d = [f(u), f′(u), f″(u), f‴(u)]` evaluated at `u = self.v`. The
1039 /// `v`/`g`/`h`/`t3` channels match [`Tower4::compose_unary`] term-for-term
1040 /// (which uses only `d[0..=3]` for those orders), so this is a strict
1041 /// truncation, not an approximation. The full-order `[f64; 5]` derivative
1042 /// stacks the families already produce can be passed by slicing their first
1043 /// four entries.
1044 ///
1045 /// # Codegen
1046 ///
1047 /// Order-≤3 Faà di Bruno written as a compact closed form instead of the
1048 /// recursive [`jet_algebra::faa_di_bruno`] walker — the order-≤2 sibling of
1049 /// [`Tower4::compose_unary`], one tensor order shallower. The loop nest is
1050 /// unchanged (no unroll over `K`, no code bloat: measured on a `Tower3<9>`
1051 /// compose-and-read consumer the new form is faster and SMALLER — asm: 71
1052 /// walker `bl` calls → 0, 39.5 KiB → 13.9 KiB, +197 NEON `.2d` ops).
1053 /// BIT-IDENTICAL: terms in the walker's exact partition order, left-
1054 /// associated block products, `acc = 0.0` accumulator start. Proven
1055 /// `to_bits`-identical on `v`/`g`/`h`/`t3` across `K ∈ {2,3,4,9}`, 5000
1056 /// random inputs each.
1057 pub fn compose_unary(&self, d: [f64; 4]) -> Self {
1058 let mut out = Self::zero();
1059 out.v = d[0];
1060 for i in 0..K {
1061 let mut acc = 0.0;
1062 acc += d[1] * self.g[i];
1063 out.g[i] = acc;
1064 }
1065 for i in 0..K {
1066 for j in 0..K {
1067 let mut acc = 0.0;
1068 acc += d[1] * self.h[i][j];
1069 acc += d[2] * self.g[i] * self.g[j];
1070 out.h[i][j] = acc;
1071 }
1072 }
1073 for i in 0..K {
1074 for j in 0..K {
1075 for k in 0..K {
1076 // walker partitions: {ijk} {ij}{k} {ik}{j} {i}{jk} {i}{j}{k}
1077 let mut acc = 0.0;
1078 acc += d[1] * self.t3[i][j][k];
1079 acc += d[2] * self.h[i][j] * self.g[k];
1080 acc += d[2] * self.h[i][k] * self.g[j];
1081 acc += d[2] * self.g[i] * self.h[j][k];
1082 acc += d[3] * self.g[i] * self.g[j] * self.g[k];
1083 out.t3[i][j][k] = acc;
1084 }
1085 }
1086 }
1087 out
1088 }
1089
1090 /// Compose with a unary special-function whose `[f64; 4]` derivative stack is
1091 /// built from the base value through `stack_fn`. Evaluates `stack_fn(self.v)`
1092 /// once and forwards to [`Self::compose_unary`], so it is bit-identical to the
1093 /// explicit form. The order-≤3 sibling of [`Tower4::compose_unary_with`].
1094 #[inline]
1095 pub fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 4]) -> Self {
1096 self.compose_unary(stack_fn(self.v))
1097 }
1098
1099 /// Single-active-slot fast path for [`Self::compose_unary`] — the order-≤3
1100 /// sibling of [`Tower4::compose_unary_single_slot`]. When `self` carries
1101 /// derivative support only on the all-`slot` diagonal, every output channel
1102 /// touching an axis `≠ slot` collapses to the walker's `total = 0.0` start
1103 /// (`+0.0`), so only `v`, `g[slot]`, `h[slot][slot]`, `t3[slot]³` survive.
1104 /// These four are computed as STRAIGHT-LINE accumulations, each in the EXACT
1105 /// term order of [`Self::compose_unary`]'s diagonal (`i = j = k = slot`)
1106 /// case (BIT-IDENTICAL to the full path on the diagonal); off-`slot`
1107 /// channels stay at the zero-init `+0.0` the full walk also yields (proven
1108 /// `to_bits` across `K ∈ {2,3,4,9}`). This drops the recursive
1109 /// set-partition walker the diagonal channels previously routed through,
1110 /// recovering its measured ~5.9× regression at the `K ∈ {2,3}` BMS tower
1111 /// widths. Caller guarantees the single-slot precondition; otherwise use
1112 /// [`Self::compose_unary`].
1113 #[inline]
1114 pub fn compose_unary_single_slot(&self, d: [f64; 4], slot: usize) -> Self {
1115 let mut out = Self::zero();
1116 let s = slot;
1117 let g = self.g[s];
1118 let h = self.h[s][s];
1119 let t3 = self.t3[s][s][s];
1120 out.v = d[0];
1121 // g (i=s): d1*g
1122 out.g[s] = {
1123 let mut acc = 0.0;
1124 acc += d[1] * g;
1125 acc
1126 };
1127 // h (i=j=s): d1*h + d2*g*g
1128 out.h[s][s] = {
1129 let mut acc = 0.0;
1130 acc += d[1] * h;
1131 acc += d[2] * g * g;
1132 acc
1133 };
1134 // t3 (i=j=k=s): exact term order of compose_unary's inner loop.
1135 out.t3[s][s][s] = {
1136 let mut acc = 0.0;
1137 acc += d[1] * t3;
1138 acc += d[2] * h * g;
1139 acc += d[2] * h * g;
1140 acc += d[2] * g * h;
1141 acc += d[3] * g * g * g;
1142 acc
1143 };
1144 out
1145 }
1146
1147 /// Multiply every channel by a plain scalar.
1148 pub fn scale(&self, s: f64) -> Self {
1149 let mut out = *self;
1150 out.v *= s;
1151 for i in 0..K {
1152 out.g[i] *= s;
1153 for j in 0..K {
1154 out.h[i][j] *= s;
1155 for k in 0..K {
1156 out.t3[i][j][k] *= s;
1157 }
1158 }
1159 }
1160 out
1161 }
1162}
1163
1164impl<const K: usize> jet_algebra::JetAlgebra<4> for Tower3<K> {
1165 #[inline]
1166 fn derivative(&self, labels: &[usize]) -> f64 {
1167 self.deriv(labels)
1168 }
1169
1170 fn map_derivatives<F>(&self, mut f: F) -> Self
1171 where
1172 F: FnMut(&[usize]) -> f64,
1173 {
1174 let mut out = Self::zero();
1175 out.v = f(&[]);
1176 for i in 0..K {
1177 let labels = [i];
1178 out.g[i] = f(&labels);
1179 }
1180 for i in 0..K {
1181 for j in 0..K {
1182 let labels = [i, j];
1183 out.h[i][j] = f(&labels);
1184 }
1185 }
1186 for i in 0..K {
1187 for j in 0..K {
1188 for k in 0..K {
1189 let labels = [i, j, k];
1190 out.t3[i][j][k] = f(&labels);
1191 }
1192 }
1193 }
1194 out
1195 }
1196}
1197
1198impl<const K: usize> std::ops::Add for Tower3<K> {
1199 type Output = Self;
1200 fn add(self, o: Self) -> Self {
1201 let mut out = self;
1202 out.v += o.v;
1203 for i in 0..K {
1204 out.g[i] += o.g[i];
1205 for j in 0..K {
1206 out.h[i][j] += o.h[i][j];
1207 for k in 0..K {
1208 out.t3[i][j][k] += o.t3[i][j][k];
1209 }
1210 }
1211 }
1212 out
1213 }
1214}
1215
1216pub fn ln_gamma_derivative_stack(x: f64) -> [f64; 5] {
1217 [
1218 statrs::function::gamma::ln_gamma(x),
1219 digamma_positive(x),
1220 polygamma_positive(1, x),
1221 polygamma_positive(2, x),
1222 polygamma_positive(3, x),
1223 ]
1224}
1225
1226pub fn ln_gamma_derivative_stack_order2(x: f64) -> [f64; 3] {
1227 [
1228 statrs::function::gamma::ln_gamma(x),
1229 digamma_positive(x),
1230 polygamma_positive(1, x),
1231 ]
1232}
1233
1234pub fn digamma_derivative_stack(x: f64) -> [f64; 5] {
1235 [
1236 digamma_positive(x),
1237 polygamma_positive(1, x),
1238 polygamma_positive(2, x),
1239 polygamma_positive(3, x),
1240 polygamma_positive(4, x),
1241 ]
1242}
1243
1244pub fn trigamma_derivative_stack(x: f64) -> [f64; 5] {
1245 [
1246 polygamma_positive(1, x),
1247 polygamma_positive(2, x),
1248 polygamma_positive(3, x),
1249 polygamma_positive(4, x),
1250 polygamma_positive(5, x),
1251 ]
1252}
1253
1254/// Scalar digamma ψ(x) for x>0. Bit-identical to `digamma_derivative_stack(x)[0]`
1255/// and to `ln_gamma_derivative_stack(x)[1]`, but evaluates ONLY ψ — the four
1256/// higher polygammas those `[f64; 5]` stacks build are pure discarded work at a
1257/// scalar consumer that reads a single element. Hot-path row kernels that need
1258/// only the digamma value (e.g. the GAMLSS Beta observed cross weight) call this
1259/// instead of indexing `[0]` off a full derivative stack.
1260#[inline]
1261pub fn digamma(x: f64) -> f64 {
1262 digamma_positive(x)
1263}
1264
1265/// Scalar trigamma ψ′(x) for x>0. Bit-identical to
1266/// `trigamma_derivative_stack(x)[0]` (both bottom out in `polygamma_positive(1,
1267/// x)`), but evaluates ONLY ψ′ — the four higher polygammas (orders 2–5) the
1268/// `[f64; 5]` stack builds are discarded at a `[0]` consumer. Used by the
1269/// dispersion-channel Fisher-information row kernels (NB2 `ψ′(θ)−ψ′(θ+μ)`, Beta
1270/// `μψ′(μφ)−(1−μ)ψ′((1−μ)φ)`) which read the trigamma value alone.
1271#[inline]
1272pub fn trigamma(x: f64) -> f64 {
1273 polygamma_positive(1, x)
1274}
1275
1276fn digamma_positive(mut x: f64) -> f64 {
1277 if !(x.is_finite() && x > 0.0) {
1278 return f64::NAN;
1279 }
1280 let mut acc = 0.0;
1281 while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
1282 acc -= 1.0 / x;
1283 x += 1.0;
1284 }
1285 acc + digamma_asymptotic(x)
1286}
1287
1288fn polygamma_positive(order: usize, mut x: f64) -> f64 {
1289 if !(x.is_finite() && x > 0.0) {
1290 return f64::NAN;
1291 }
1292 let mut acc = 0.0;
1293 while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
1294 acc += polygamma_recurrence_term(order, x);
1295 x += 1.0;
1296 }
1297 acc + polygamma_asymptotic(order, x)
1298}
1299
1300const POLYGAMMA_ASYMPTOTIC_MIN_X: f64 = 20.0;
1301const BERNOULLI_EVEN: [(usize, f64); 10] = [
1302 (2, 1.0 / 6.0),
1303 (4, -1.0 / 30.0),
1304 (6, 1.0 / 42.0),
1305 (8, -1.0 / 30.0),
1306 (10, 5.0 / 66.0),
1307 (12, -691.0 / 2730.0),
1308 (14, 7.0 / 6.0),
1309 (16, -3617.0 / 510.0),
1310 (18, 43867.0 / 798.0),
1311 (20, -174611.0 / 330.0),
1312];
1313
1314fn polygamma_recurrence_term(order: usize, x: f64) -> f64 {
1315 let sign = if order % 2 == 1 { 1.0 } else { -1.0 };
1316 sign * factorial(order) / x.powi((order + 1) as i32)
1317}
1318
1319fn digamma_asymptotic(x: f64) -> f64 {
1320 let mut out = x.ln() - 0.5 / x;
1321 for (bernoulli_order, bernoulli) in BERNOULLI_EVEN {
1322 out -= bernoulli / (bernoulli_order as f64 * x.powi(bernoulli_order as i32));
1323 }
1324 out
1325}
1326
1327fn polygamma_asymptotic(order: usize, x: f64) -> f64 {
1328 if !(1..=5).contains(&order) {
1329 return f64::NAN;
1330 }
1331
1332 let order_factorial = factorial(order);
1333 let leading_sign = if order % 2 == 1 { 1.0 } else { -1.0 };
1334 let mut out = leading_sign * factorial(order - 1) / x.powi(order as i32)
1335 + leading_sign * order_factorial / (2.0 * x.powi((order + 1) as i32));
1336
1337 let bernoulli_sign = if order % 2 == 1 { 1.0 } else { -1.0 };
1338 for (bernoulli_order, bernoulli) in BERNOULLI_EVEN {
1339 let rising = rising_factorial(bernoulli_order, order);
1340 out += bernoulli_sign * bernoulli * rising
1341 / bernoulli_order as f64
1342 / x.powi((bernoulli_order + order) as i32);
1343 }
1344 out
1345}
1346
1347fn factorial(n: usize) -> f64 {
1348 (1..=n).fold(1.0, |acc, k| acc * k as f64)
1349}
1350
1351fn rising_factorial(start: usize, len: usize) -> f64 {
1352 (start..start + len).fold(1.0, |acc, k| acc * k as f64)
1353}
1354
1355impl<const K: usize> std::ops::Add for Tower4<K> {
1356 type Output = Self;
1357 fn add(self, o: Self) -> Self {
1358 let mut out = self;
1359 out.v += o.v;
1360 for i in 0..K {
1361 out.g[i] += o.g[i];
1362 for j in 0..K {
1363 out.h[i][j] += o.h[i][j];
1364 for k in 0..K {
1365 out.t3[i][j][k] += o.t3[i][j][k];
1366 for l in 0..K {
1367 out.t4[i][j][k][l] += o.t4[i][j][k][l];
1368 }
1369 }
1370 }
1371 }
1372 out
1373 }
1374}
1375
1376impl<const K: usize> std::ops::Sub for Tower4<K> {
1377 type Output = Self;
1378 fn sub(self, o: Self) -> Self {
1379 self + o.scale(-1.0)
1380 }
1381}
1382
1383impl<const K: usize> std::ops::Neg for Tower4<K> {
1384 type Output = Self;
1385 fn neg(self) -> Self {
1386 self.scale(-1.0)
1387 }
1388}
1389
1390impl<const K: usize> std::ops::Mul for Tower4<K> {
1391 type Output = Self;
1392 fn mul(self, o: Self) -> Self {
1393 Tower4::mul(&self, &o)
1394 }
1395}
1396
1397impl<const K: usize> std::ops::Div for Tower4<K> {
1398 type Output = Self;
1399 fn div(self, o: Self) -> Self {
1400 Tower4::mul(&self, &o.recip())
1401 }
1402}
1403
1404impl<const K: usize> std::ops::Add<f64> for Tower4<K> {
1405 type Output = Self;
1406 fn add(self, c: f64) -> Self {
1407 let mut out = self;
1408 out.v += c;
1409 out
1410 }
1411}
1412
1413impl<const K: usize> std::ops::Sub<f64> for Tower4<K> {
1414 type Output = Self;
1415 fn sub(self, c: f64) -> Self {
1416 self + (-c)
1417 }
1418}
1419
1420impl<const K: usize> std::ops::Mul<f64> for Tower4<K> {
1421 type Output = Self;
1422 fn mul(self, c: f64) -> Self {
1423 self.scale(c)
1424 }
1425}
1426
1427// ── Implicit-function and moving-boundary seams (#932 flex) ──────────
1428//
1429// The flexible survival marginal-slope row loss is NOT a free composition
1430// of the primaries: it threads an IMPLICIT calibration intercept `a(θ)`
1431// solving a constraint `F(a, θ) = 0`, and integrates a density over cells
1432// whose edges `z_L(θ), z_R(θ)` MOVE with θ through that intercept. Plain
1433// `Tower4` Faà di Bruno cannot express either — so the flex tower was the
1434// last hand-written one in the codebase, and the genus of #736-class
1435// drift bugs (the (g,w0) deviation-cross third was 3× short for exactly
1436// this reason). These two combinators close that gap: once the constraint
1437// `F` and the integrand/boundaries are themselves towers, the intercept's
1438// derivative tower and the integral's derivative tower come out EXACTLY at
1439// every order — there is no order left to hand-code and forget.
1440
1441/// Solve the implicit relation `F(a(θ), θ) ≡ 0` for the intercept tower
1442/// `a(θ)` over the `K` primaries θ, given the constraint tower `f` written
1443/// over `K + 1` variables (slot `0` is the intercept `a`, slots `1..=K`
1444/// are the primaries θ) evaluated at the SOLVED point — i.e. `f.v` is the
1445/// constraint residual at `(a₀, θ₀)` (≈ 0 from the production Newton solve)
1446/// and `a0` is that solved intercept value.
1447///
1448/// Returns the `Tower4<K>` whose value is `a0` and whose every derivative
1449/// tensor (∂a/∂θ, ∂²a/∂θ², …, ∂⁴a/∂θ⁴) is the exact implicit-function
1450/// derivative. This is the mechanical replacement for the hand-coded
1451/// `a_u = -f_u/f_a`, `a_uv = -(f_uv + f_au·a_v + f_av·a_u + f_aa·a_u·a_v)/f_a`
1452/// recursion (first_full.rs) and its third/fourth-order continuations.
1453///
1454/// Method: order-by-order substitution. We build `a` incrementally; at each
1455/// order `m` the composite `G(θ) = f(a(θ), θ)` has a top-order coefficient
1456/// that is linear in `a`'s order-`m` tensor with leading factor `F_a`
1457/// (= `f.g[0]`), plus terms in `a`'s lower orders already fixed. Setting the
1458/// order-`m` tensor of `a` to cancel the rest of `G`'s order-`m` coefficient
1459/// keeps `G ≡ 0` through that order. The substitution `G = f∘(a, θ)` reuses
1460/// only the exact [`substitute_intercept`] chain rule, so the recursion is
1461/// auditable and exact, not a hand-expanded formula per order.
1462///
1463/// `f.g[0]` (= ∂F/∂a) must be non-zero — guaranteed by the production
1464/// solve's strict monotonicity guard.
1465///
1466/// The expansion point `a0` must be a genuine root `F(a0, θ0) = 0`: the
1467/// substitution recursion below cancels orders 1..=4 of `G = F∘a` but never
1468/// touches order 0, so a non-root `a0` would yield the Taylor expansion of
1469/// the LEVEL SET `F = F(a0)` through `a0`, not the root curve `F = 0`. This
1470/// is guarded explicitly and re-verified by a composed-residual self-check.
1471pub fn implicit_solve<const K1: usize, const K: usize>(
1472 f: &Tower4<K1>,
1473 a0: f64,
1474) -> Result<Tower4<K>, String> {
1475 assert_eq!(K1, K + 1, "implicit_solve: constraint must carry K+1 vars");
1476 let f_a = f.g[0];
1477 if f_a == 0.0 || !f_a.is_finite() {
1478 return Err(format!(
1479 "implicit_solve: ∂F/∂a = {f_a:+.3e} is not invertible"
1480 ));
1481 }
1482 // The expansion point must be a genuine root of F. The single Newton
1483 // correction that would move a0 onto the root is |f.v|/|f_a|; require it
1484 // to be negligible relative to the natural scale (1 + |a0|). Guarding the
1485 // Newton step (rather than f.v directly) makes the criterion invariant to
1486 // the magnitude of f_a / the units of F.
1487 let root_tol = 1e-9;
1488 if !f.v.is_finite() {
1489 return Err(format!(
1490 "implicit_solve: F(a0, θ0) = {:+.3e} is not finite",
1491 f.v
1492 ));
1493 }
1494 let newton_step = f.v.abs() / f_a.abs();
1495 if newton_step > root_tol * (1.0 + a0.abs()) {
1496 return Err(format!(
1497 "implicit_solve: expansion point a0 = {a0:+.6e} is not a root of F: \
1498 F(a0, θ0) = {:+.3e}, Newton correction {newton_step:+.3e} exceeds \
1499 root_tol {root_tol:.1e} · (1 + |a0|)",
1500 f.v
1501 ));
1502 }
1503 // Start with a = constant a0 (correct through order 0). Then lift each
1504 // order in turn. Because substitute_intercept reads `a`'s order-≤m
1505 // tensors when forming G's order-m coefficient, and the order-m
1506 // coefficient of G depends on a's order-m tensor ONLY through the linear
1507 // F_a·a_m term, a single corrective pass per order is exact.
1508 let mut a = Tower4::<K>::constant(a0);
1509 for order in 1..=4 {
1510 let g = substitute_intercept(f, &a);
1511 // Cancel G's order-`order` coefficient by adjusting a's order-`order`
1512 // tensor: a_m -= G_m / F_a (the F_a·a_m term is the only one carrying
1513 // a's order-m tensor, with unit chain coefficient since slot 0 seeds a
1514 // as a plain variable in the substitution's first-order part).
1515 match order {
1516 1 => {
1517 for i in 0..K {
1518 a.g[i] -= g.g[i] / f_a;
1519 }
1520 }
1521 2 => {
1522 for i in 0..K {
1523 for j in 0..K {
1524 a.h[i][j] -= g.h[i][j] / f_a;
1525 }
1526 }
1527 }
1528 3 => {
1529 for i in 0..K {
1530 for j in 0..K {
1531 for k in 0..K {
1532 a.t3[i][j][k] -= g.t3[i][j][k] / f_a;
1533 }
1534 }
1535 }
1536 }
1537 _ => {
1538 for i in 0..K {
1539 for j in 0..K {
1540 for k in 0..K {
1541 for l in 0..K {
1542 a.t4[i][j][k][l] -= g.t4[i][j][k][l] / f_a;
1543 }
1544 }
1545 }
1546 }
1547 }
1548 }
1549 }
1550 // Self-check: the composed residual G = F∘a must vanish through order 4.
1551 // By construction orders 1..=4 were cancelled; the value G.v == F(a0,θ0)
1552 // is exactly the root requirement guarded above. Re-verify all channels
1553 // against a scale-aware floor so any arithmetic regression in the
1554 // substitution recursion is loud rather than silently shipping a
1555 // level-set expansion.
1556 let g = substitute_intercept(f, &a);
1557 let resid_tol = 1e-7 * (1.0 + f_a.abs());
1558 let mut worst = g.v.abs();
1559 for i in 0..K {
1560 worst = worst.max(g.g[i].abs());
1561 for j in 0..K {
1562 worst = worst.max(g.h[i][j].abs());
1563 for k in 0..K {
1564 worst = worst.max(g.t3[i][j][k].abs());
1565 for l in 0..K {
1566 worst = worst.max(g.t4[i][j][k][l].abs());
1567 }
1568 }
1569 }
1570 }
1571 if !worst.is_finite() || worst > resid_tol {
1572 return Err(format!(
1573 "implicit_solve: composed residual G = F∘a does not vanish: \
1574 worst channel magnitude {worst:+.3e} exceeds tol {resid_tol:.1e}"
1575 ));
1576 }
1577 Ok(a)
1578}
1579
1580/// Substitute the intercept tower `a(θ)` into slot `0` of a constraint
1581/// written over `K + 1` variables, returning the composite tower over the
1582/// `K` primaries θ: `G(θ) = f(a(θ), θ₁, …, θ_K)`.
1583///
1584/// This is the exact multivariate chain rule specialised to "slot 0 is a
1585/// dependent tower, slots 1..=K are the independent primaries". It evaluates
1586/// `f`'s fourth-order multivariate Taylor polynomial about the expansion
1587/// point, with the slot-0 increment being the non-constant part of `a` and
1588/// the slot-(i) increment being the unit-seeded primary `θ_i`. The sum is
1589/// assembled by the same subset/partition algebra `Tower4` arithmetic uses,
1590/// so it carries derivatives exactly through order four.
1591pub fn substitute_intercept<const K1: usize, const K: usize>(
1592 f: &Tower4<K1>,
1593 a: &Tower4<K>,
1594) -> Tower4<K> {
1595 assert_eq!(K1, K + 1);
1596 // Build the K+1 input towers in θ-space: slot 0 = a(θ), slot i+1 = θ_i.
1597 // The composite is Σ over ordered label tuples s (|s| ≤ 4) of input
1598 // indices: (1/|s|!) · f.deriv(s) · Π_{j in s} (inp[s_j] centred) — but
1599 // since f.deriv is the SYMMETRIC partial tensor and we enumerate ordered
1600 // tuples, the 1/|s|! exactly cancels the tuple multiplicity. We assemble
1601 // it directly as a Horner-free explicit sum over the (K+1)-ary tuples,
1602 // using tower products for the increment monomials so all θ-derivatives
1603 // propagate exactly.
1604 let inp: [Tower4<K>; K1] = std::array::from_fn(|slot| {
1605 if slot == 0 {
1606 // slot 0: a(θ) minus its constant value (the increment δa(θ)).
1607 let mut d = *a;
1608 d.v = 0.0;
1609 d
1610 } else {
1611 // slot i: the increment δθ_{i-1} = seeded variable minus value.
1612 // θ centred at its expansion value has zero constant term and unit
1613 // first derivative in its own slot.
1614 let mut d = Tower4::<K>::zero();
1615 d.g[slot - 1] = 1.0;
1616 d
1617 }
1618 });
1619 // Accumulate the Taylor sum. order-0 term:
1620 let mut out = Tower4::<K>::constant(f.v);
1621 // order 1: Σ_a f.g[a] · inp[a]
1622 for a_idx in 0..K1 {
1623 out = out + inp[a_idx].scale(f.g[a_idx]);
1624 }
1625 // order 2: (1/2) Σ_{a,b} f.h[a][b] · inp[a]·inp[b]
1626 for a_idx in 0..K1 {
1627 for b_idx in 0..K1 {
1628 let prod = inp[a_idx].mul(&inp[b_idx]);
1629 out = out + prod.scale(0.5 * f.h[a_idx][b_idx]);
1630 }
1631 }
1632 // order 3: (1/6) Σ f.t3[a][b][c] · inp[a]·inp[b]·inp[c]
1633 for a_idx in 0..K1 {
1634 for b_idx in 0..K1 {
1635 for c_idx in 0..K1 {
1636 let prod = inp[a_idx].mul(&inp[b_idx]).mul(&inp[c_idx]);
1637 out = out + prod.scale(f.t3[a_idx][b_idx][c_idx] / 6.0);
1638 }
1639 }
1640 }
1641 // order 4: (1/24) Σ f.t4[a][b][c][d] · inp[a]·inp[b]·inp[c]·inp[d]
1642 for a_idx in 0..K1 {
1643 for b_idx in 0..K1 {
1644 for c_idx in 0..K1 {
1645 for d_idx in 0..K1 {
1646 let prod = inp[a_idx]
1647 .mul(&inp[b_idx])
1648 .mul(&inp[c_idx])
1649 .mul(&inp[d_idx]);
1650 out = out + prod.scale(f.t4[a_idx][b_idx][c_idx][d_idx] / 24.0);
1651 }
1652 }
1653 }
1654 }
1655 out
1656}
1657
1658/// The exact θ-derivative tower of a moving-LIMIT integral's BOUNDARY
1659/// contribution: given the edge-position tower `z_edge(θ)` over the `K`
1660/// primaries and the integrand `B` evaluated-and-differentiated at the edge
1661/// value as the stack `b_stack = [B(z₀), B′(z₀), B″(z₀), B‴(z₀)]`
1662/// (`z₀ = z_edge.v`), returns the tower of `Φ(z_edge(θ))` where `Φ′ = B`.
1663///
1664/// Rationale: `∂_θ ∫^{z_edge(θ)} B(z) dz = Φ(z_edge(θ))` with `Φ` an
1665/// antiderivative of `B`, so the boundary part of every θ-derivative of the
1666/// integral is just the composition `Φ ∘ z_edge` — whose Faà di Bruno
1667/// expansion carries, at one stroke, EVERY Leibniz boundary term the
1668/// hand-written flux dropped: the first-order `B·z_u`, the second-order
1669/// `B′·z_u·z_v + B·z_uv` (the `G_z·z_u·z_v` self-flux AND the previously
1670/// dropped `G·z_uv`), and the full third/fourth-order continuations. The
1671/// VALUE channel of the returned tower is meaningless (`Φ` is only defined up
1672/// to a constant); callers read only the derivative channels and pair this
1673/// with the interior moment-integral value separately.
1674///
1675/// `b_stack` holds `B` and its first three z-derivatives; the antiderivative
1676/// `Φ` contributes only as the order-≥1 channels, so `compose_unary` receives
1677/// `[0, B, B′, B″, B‴]` — the leading `0` is the discarded `Φ(z₀)` slot.
1678pub fn moving_limit_boundary_tower<const K: usize>(
1679 z_edge: &Tower4<K>,
1680 b_stack: [f64; 4],
1681) -> Tower4<K> {
1682 z_edge.compose_unary([0.0, b_stack[0], b_stack[1], b_stack[2], b_stack[3]])
1683}
1684
1685/// The boundary-flux derivative tower of a single moving cell integral
1686/// `∫_{z_L(θ)}^{z_R(θ)} B dz`: `Φ(z_R(θ)) − Φ(z_L(θ))`, assembled from the
1687/// two edge towers and the integrand stacks at each edge. The returned
1688/// tower's derivative channels are the EXACT moving-boundary contribution to
1689/// every θ-derivative of the cell integral, to fourth order, with no term
1690/// hand-omitted. A `Fixed` (non-moving) edge passes a `z_edge` whose
1691/// derivative channels are all zero, contributing nothing — matching the
1692/// production `edge_vel = 0` short-circuit.
1693pub fn cell_moving_boundary_flux_tower<const K: usize>(
1694 z_right: &Tower4<K>,
1695 b_stack_right: [f64; 4],
1696 z_left: &Tower4<K>,
1697 b_stack_left: [f64; 4],
1698) -> Tower4<K> {
1699 moving_limit_boundary_tower(z_right, b_stack_right)
1700 - moving_limit_boundary_tower(z_left, b_stack_left)
1701}
1702
1703/// Moving-limit boundary tower for a θ-DEPENDENT integrand `G(z; θ)`.
1704///
1705/// [`moving_limit_boundary_tower`] assumes the integrand depends on θ only
1706/// through the moving edge `z_edge(θ)` (a fixed z-derivative `b_stack`). The
1707/// marginal-slope flex boundary is richer: the integrand `G(z; θ)` ALSO carries
1708/// its own θ-dependence (the density weight `w = e^{−q}/2π` and the cell
1709/// integrand coefficients move with η, hence with the primaries), so the
1710/// Leibniz expansion of `∂ⁿ_θ ∫^{z_edge(θ)} G(z;θ) dz` mixes edge-motion
1711/// derivatives of the limit with θ-derivatives of `G` itself — e.g. at second
1712/// order `G·z_uv + G_z·z_u·z_v + G_{θu}·z_v + G_{θv}·z_u` (the four
1713/// edge-motion-carrying terms the hand path assembles one by one, including the
1714/// `G·z_uv` term the directional path drops).
1715///
1716/// Mechanization: let `Φ(z; θ)` be the z-antiderivative of `G` (so `Φ_z = G`).
1717/// The full upper-limit contribution is `Φ(z_edge(θ); θ)`, and the BOUNDARY
1718/// part — everything carrying edge motion — is exactly
1719/// `Φ(z_edge(θ); θ) − Φ(z₀; θ)`,
1720/// the second term being the pure-integrand-θ part (`∫^{z₀} ∂ⁿ_θ G`) the
1721/// interior moment integral already supplies. Both are one
1722/// [`substitute_intercept`] of the SAME mixed `(z, θ)` jet of `Φ` (z in slot 0,
1723/// θ in slots 1..K): substituting the edge tower gives the full composite,
1724/// substituting a frozen constant edge isolates the pure-θ part, and their
1725/// difference is the exact boundary flux — every Leibniz term derived by the
1726/// substitution algebra, none hand-omitted.
1727///
1728/// `phi_jet` is the `(K+1)`-variable Taylor jet of `Φ` about `(z₀, θ₀)` with
1729/// `z₀ = z_edge.v`: slot 0 is the z-direction (so `phi_jet.g[0] = G(z₀;θ₀)`,
1730/// `phi_jet.h[0][0] = G_z`, …) and slots `1..=K` are the primaries θ (carrying
1731/// `Φ`'s own θ- and mixed z·θ-derivatives — i.e. the integrand's θ-derivatives
1732/// integrated in z, and `G_{θ…}` in the mixed slots). The returned tower's
1733/// VALUE channel is 0 by construction (the `Φ(z₀;θ₀)` constants cancel); only
1734/// the derivative channels are meaningful, matching the value-less convention of
1735/// [`moving_limit_boundary_tower`].
1736pub fn moving_limit_boundary_tower_theta_integrand<const K1: usize, const K: usize>(
1737 phi_jet: &Tower4<K1>,
1738 z_edge: &Tower4<K>,
1739) -> Tower4<K> {
1740 assert_eq!(
1741 K1,
1742 K + 1,
1743 "moving_limit_boundary_tower_theta_integrand: Φ jet must carry z + K θ-vars"
1744 );
1745 let frozen_edge = Tower4::<K>::constant(z_edge.v);
1746 let full = substitute_intercept(phi_jet, z_edge);
1747 let interior = substitute_intercept(phi_jet, &frozen_edge);
1748 full - interior
1749}
1750
1751
1752// ── The program seam ─────────────────────────────────────────────────
1753
1754// ── The canonical single-source seam (#932 consolidation) ────────────
1755//
1756// `RowProgram<K>` is the ONE row-program interface #932 converges every family
1757// onto. Its generic `eval<S: JetScalar<K>>` body is the go-forward derivation
1758// surface for every calculus channel; `program_*` selects only the derivative
1759// representation each consumer needs.
1760
1761/// The single source of truth #932 asks for: a family's row negative
1762/// log-likelihood written ONCE over the generic [`crate::jet_scalar::JetScalar`]
1763/// interface, from which every `RowKernel` (gam-models) derivative channel is
1764/// mechanically derived. A family implements ONLY this (plus its linear Jacobian
1765/// wiring, which is family data, not calculus) — it cannot author an independent
1766/// derivative tower, because there is no other channel to author.
1767///
1768/// Because a body uses only `add`/`sub`/`mul`/`scale`/`exp`/`ln`/… — all provided
1769/// by [`crate::jet_scalar::JetScalar`] — the SAME body re-instantiates at
1770/// [`crate::jet_scalar::Order2`] (value/grad/Hessian), [`crate::jet_scalar::OneSeed`]
1771/// (contracted third), [`crate::jet_scalar::TwoSeed`] (contracted fourth), and the
1772/// full [`Tower4`] (every channel), with the contraction folded into the
1773/// differentiation so no dense `t3`/`t4` is ever materialised.
1774pub trait RowProgram<const K: usize>: Send + Sync {
1775 /// Number of observations the program covers.
1776 fn n_rows(&self) -> usize;
1777
1778 /// Current primary-scalar values for `row` (where to seed the scalar).
1779 fn primaries(&self, row: usize) -> Result<[f64; K], String>;
1780
1781 /// The row NLL evaluated on a generic jet scalar. `p[a]` arrives pre-seeded
1782 /// (base value + per-scalar nilpotent directions) by the caller; the body
1783 /// uses ONLY [`crate::jet_scalar::JetScalar`] ops and per-row data (response,
1784 /// censoring, offsets) entering as constants.
1785 fn eval<S: crate::jet_scalar::JetScalar<K>>(&self, row: usize, p: &[S; K])
1786 -> Result<S, String>;
1787}
1788
1789/// Maximum size of one canonical dense-jet storage object kept on the call
1790/// stack. Small fixed-width programs stay allocation-free; wider derivative
1791/// representations use exact-length heap storage instead of making the thread
1792/// stack scale as `K * size_of::<S>()`. A full dense result larger than this
1793/// boundary is rejected in favor of the bounded directional APIs.
1794///
1795/// This is a storage-policy boundary, not a calculus fallback: both branches
1796/// invoke the same [`RowProgram::eval`] expression with the same scalar type.
1797const PROGRAM_DENSE_JET_STACK_BUDGET_BYTES: usize = 64 * 1024;
1798
1799#[inline]
1800fn program_primary_jets_fit_stack<S, const K: usize>() -> bool {
1801 std::mem::size_of::<S>()
1802 .checked_mul(K)
1803 .is_some_and(|bytes| bytes <= PROGRAM_DENSE_JET_STACK_BUDGET_BYTES)
1804}
1805
1806fn evaluate_program_with_stack_primaries<const K: usize, P, S>(
1807 prog: &P,
1808 row: usize,
1809 mut seed: impl FnMut(usize) -> S,
1810) -> Result<S, String>
1811where
1812 P: RowProgram<K> + ?Sized,
1813 S: crate::jet_scalar::JetScalar<K>,
1814{
1815 let vars: [S; K] = std::array::from_fn(&mut seed);
1816 prog.eval(row, &vars)
1817}
1818
1819#[inline(never)]
1820fn evaluate_program_with_heap_primaries<const K: usize, P, S>(
1821 prog: &P,
1822 row: usize,
1823 seed: impl FnMut(usize) -> S,
1824) -> Result<S, String>
1825where
1826 P: RowProgram<K> + ?Sized,
1827 S: crate::jet_scalar::JetScalar<K>,
1828{
1829 // The exact-size range builds precisely K initialized Copy scalars in
1830 // heap-backed storage. Converting the boxed slice to a boxed array changes
1831 // only its type; it never materializes `[S; K]` on the stack.
1832 let vars: Box<[S]> = (0..K).map(seed).collect();
1833 let vars: Box<[S; K]> = vars.try_into().map_err(|vars: Box<[S]>| {
1834 format!(
1835 "canonical row program seeded {} primary jets; expected exactly {K}",
1836 vars.len()
1837 )
1838 })?;
1839 prog.eval(row, &vars)
1840}
1841
1842#[inline]
1843fn evaluate_program_with_seeded_primaries<const K: usize, P, S>(
1844 prog: &P,
1845 row: usize,
1846 seed: impl FnMut(usize) -> S,
1847) -> Result<S, String>
1848where
1849 P: RowProgram<K> + ?Sized,
1850 S: crate::jet_scalar::JetScalar<K>,
1851{
1852 if program_primary_jets_fit_stack::<S, K>() {
1853 evaluate_program_with_stack_primaries(prog, row, seed)
1854 } else {
1855 evaluate_program_with_heap_primaries(prog, row, seed)
1856 }
1857}
1858
1859/// Derive the `row_kernel` channel `(nll, ∇, H)` from a [`RowProgram`] at the
1860/// value/gradient/Hessian scalar [`crate::jet_scalar::Order2`], WITHOUT
1861/// materialising any third / fourth tensor.
1862pub fn program_row_kernel<const K: usize, P: RowProgram<K> + ?Sized>(
1863 prog: &P,
1864 row: usize,
1865) -> Result<(f64, [f64; K], [[f64; K]; K]), String> {
1866 let base = prog.primaries(row)?;
1867 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1868 <crate::jet_scalar::Order2<K> as crate::jet_scalar::JetScalar<K>>::variable(base[a], a)
1869 })?;
1870 Ok(s.into_channels())
1871}
1872
1873/// Derive the `row_third_contracted(dir)` channel `Σ_c ℓ_{abc} dir_c` from a
1874/// [`RowProgram`] at the one-seed scalar [`crate::jet_scalar::OneSeed`], WITHOUT
1875/// materialising the dense `t3`.
1876pub fn program_third_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
1877 prog: &P,
1878 row: usize,
1879 dir: &[f64; K],
1880) -> Result<[[f64; K]; K], String> {
1881 let base = prog.primaries(row)?;
1882 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1883 crate::jet_scalar::OneSeed::seed_direction(base[a], a, dir[a])
1884 })?;
1885 Ok(s.contracted_third())
1886}
1887
1888/// Derive the `row_fourth_contracted(u, v)` channel `Σ_{cd} ℓ_{abcd} u_c v_d`
1889/// from a [`RowProgram`] at the two-seed scalar [`crate::jet_scalar::TwoSeed`],
1890/// WITHOUT materialising the dense `t4`.
1891pub fn program_fourth_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
1892 prog: &P,
1893 row: usize,
1894 dir_u: &[f64; K],
1895 dir_v: &[f64; K],
1896) -> Result<[[f64; K]; K], String> {
1897 let base = prog.primaries(row)?;
1898 let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
1899 crate::jet_scalar::TwoSeed::seed(base[a], a, dir_u[a], dir_v[a])
1900 })?;
1901 Ok(s.contracted_fourth())
1902}
1903
1904/// Derive every channel `(v, g, h, t3, t4)` in one pass from a [`RowProgram`] at
1905/// the full dense [`Tower4`] scalar.
1906///
1907/// The result is boxed so the return slot itself remains bounded independently
1908/// of `K`. Dense towers above the canonical storage budget are rejected before
1909/// the program is touched; consumers at those widths must request only the
1910/// channels they need through [`program_row_kernel`],
1911/// [`program_third_contracted`], and [`program_fourth_contracted`].
1912pub fn program_full_tower<const K: usize, P: RowProgram<K> + ?Sized>(
1913 prog: &P,
1914 row: usize,
1915) -> Result<Box<Tower4<K>>, String> {
1916 let tower_bytes = std::mem::size_of::<Tower4<K>>();
1917 if tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES {
1918 return Err(format!(
1919 "canonical dense Tower4<{K}> requires {tower_bytes} bytes, exceeding the {}-byte \
1920 storage budget; use the bounded row-kernel and directional channel APIs",
1921 PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
1922 ));
1923 }
1924 let base = prog.primaries(row)?;
1925 evaluate_program_with_seeded_primaries(prog, row, |a| Tower4::variable(base[a], a))
1926 .map(Box::new)
1927}
1928
1929// ── The oracle ───────────────────────────────────────────────────────
1930
1931/// One row's worth of hand-written kernel outputs, as claimed by a
1932/// `RowKernel` implementation, packaged for verification against the
1933/// tower truth. Plain data (no trait coupling) so any kernel — whatever
1934/// its visibility — can be audited from its own test module.
1935pub struct KernelChannels<const K: usize> {
1936 /// Claimed `(nll, ∇, H)` from `row_kernel`.
1937 pub value: f64,
1938 /// Claimed gradient.
1939 pub gradient: [f64; K],
1940 /// Claimed Hessian.
1941 pub hessian: [[f64; K]; K],
1942 /// Claimed `row_third_contracted(dir)` outputs as `(dir, claim)` pairs.
1943 pub third: Vec<([f64; K], [[f64; K]; K])>,
1944 /// Claimed `row_fourth_contracted(u, v)` outputs as `(u, v, claim)`.
1945 pub fourth: Vec<([f64; K], [f64; K], [[f64; K]; K])>,
1946}
1947
1948/// Channel-by-channel audit of a hand-written kernel against the
1949/// single-expression tower truth. Returns `Err` naming the first channel,
1950/// index, claimed and true values on disagreement — designed as the body
1951/// of the per-family CI oracle tests (#932 deployment step 2).
1952///
1953/// Tolerance is PER ENTRY, mixed absolute/relative: each comparison uses
1954/// `|claim − truth| ≤ atol + rel_tol · max(|claim|, |truth|)`. The absolute
1955/// floor `atol = rel_tol` lets exact-zero entries of structurally sparse
1956/// towers pass without demanding bit-equality, while a tiny cross-block
1957/// entry dropped next to a huge one is still caught (it is NOT measured
1958/// against the largest entry of the whole channel — there is no per-channel
1959/// magnitude floor). Genuine sign flips (#736) and dropped channels are loud.
1960///
1961/// Non-finite handling is strict: a NaN on either side always fails; an
1962/// infinity passes only when both sides are the SAME signed infinity.
1963pub fn verify_kernel_channels<const K: usize>(
1964 tower: &Tower4<K>,
1965 claims: &KernelChannels<K>,
1966 rel_tol: f64,
1967) -> Result<(), String> {
1968 // Absolute floor: reuse rel_tol so a single knob controls both the
1969 // relative band and the absolute floor for entries near zero.
1970 let atol = rel_tol;
1971 let check = |label: &str, claim: f64, truth: f64| -> Result<(), String> {
1972 // Non-finite values never silently pass the algebraic comparison
1973 // below (any comparison with NaN is false). Handle them explicitly:
1974 // NaN on either side always errs; an infinity passes only if both
1975 // sides are the identical signed infinity.
1976 if !claim.is_finite() || !truth.is_finite() {
1977 let agree = claim.is_infinite()
1978 && truth.is_infinite()
1979 && claim.is_sign_positive() == truth.is_sign_positive();
1980 if agree {
1981 return Ok(());
1982 }
1983 return Err(format!(
1984 "row-kernel oracle: {label} non-finite mismatch: claimed {claim:+.12e}, tower {truth:+.12e}"
1985 ));
1986 }
1987 let band = atol + rel_tol * claim.abs().max(truth.abs());
1988 if (claim - truth).abs() > band {
1989 return Err(format!(
1990 "row-kernel oracle: {label} disagrees: claimed {claim:+.12e}, tower {truth:+.12e} (rel_tol {rel_tol:.1e}, atol {atol:.1e}, band {band:.3e})"
1991 ));
1992 }
1993 Ok(())
1994 };
1995
1996 check("value", claims.value, tower.v)?;
1997
1998 for a in 0..K {
1999 check(&format!("gradient[{a}]"), claims.gradient[a], tower.g[a])?;
2000 }
2001
2002 for a in 0..K {
2003 for b in 0..K {
2004 check(
2005 &format!("hessian[{a}][{b}]"),
2006 claims.hessian[a][b],
2007 tower.h[a][b],
2008 )?;
2009 }
2010 }
2011
2012 for (t_idx, (dir, claim)) in claims.third.iter().enumerate() {
2013 let truth = tower.third_contracted(dir);
2014 for a in 0..K {
2015 for b in 0..K {
2016 check(
2017 &format!("third[{t_idx}][{a}][{b}]"),
2018 claim[a][b],
2019 truth[a][b],
2020 )?;
2021 }
2022 }
2023 }
2024
2025 for (f_idx, (u, w, claim)) in claims.fourth.iter().enumerate() {
2026 let truth = tower.fourth_contracted(u, w);
2027 for a in 0..K {
2028 for b in 0..K {
2029 check(
2030 &format!("fourth[{f_idx}][{a}][{b}]"),
2031 claim[a][b],
2032 truth[a][b],
2033 )?;
2034 }
2035 }
2036 }
2037
2038 Ok(())
2039}
2040
2041#[cfg(test)]
2042mod tests {
2043 use super::*;
2044
2045 /// `Tower3<K>` must be bit-identical to `Tower4<K>` on every channel it
2046 /// carries (value, gradient, Hessian, third derivatives). The order-≤3
2047 /// Leibniz / Faà-di-Bruno terms read only order-≤3 inner channels, so
2048 /// dropping the fourth tensor cannot perturb them. Exercises products
2049 /// (Leibniz cross-terms), unary composition, scaling, and addition — the
2050 /// same operations the survival location-scale `nll_index_tower` composes —
2051 /// across all mixed partials, not just the diagonal entries that kernel reads.
2052 #[test]
2053 fn tower3_matches_tower4_through_third_order() {
2054 let s_a: [f64; 5] = [
2055 0.3_f64.sin(),
2056 0.3_f64.cos(),
2057 -0.3_f64.sin(),
2058 -0.3_f64.cos(),
2059 0.3_f64.sin(),
2060 ];
2061 let s_b: [f64; 5] = [1.1, -0.4, 0.8, -0.2, 0.05];
2062 let s4 = |s: [f64; 5]| [s[0], s[1], s[2], s[3]];
2063
2064 let a4 = Tower4::<3>::variable(0.4, 0);
2065 let b4 = Tower4::<3>::variable(-0.7, 1);
2066 let c4 = Tower4::<3>::variable(0.9, 2);
2067 let prog4 = (a4.mul(&b4) + c4).compose_unary(s_a).scale(1.3)
2068 + a4.mul(&c4).scale(-0.7)
2069 + b4.compose_unary(s_b).scale(0.25);
2070
2071 let a3 = Tower3::<3>::variable(0.4, 0);
2072 let b3 = Tower3::<3>::variable(-0.7, 1);
2073 let c3 = Tower3::<3>::variable(0.9, 2);
2074 let prog3 = (a3.mul(&b3) + c3).compose_unary(s4(s_a)).scale(1.3)
2075 + a3.mul(&c3).scale(-0.7)
2076 + b3.compose_unary(s4(s_b)).scale(0.25);
2077
2078 assert_eq!(prog3.v.to_bits(), prog4.v.to_bits(), "value mismatch");
2079 for i in 0..3 {
2080 assert_eq!(
2081 prog3.g[i].to_bits(),
2082 prog4.g[i].to_bits(),
2083 "g[{i}] mismatch"
2084 );
2085 for j in 0..3 {
2086 assert_eq!(
2087 prog3.h[i][j].to_bits(),
2088 prog4.h[i][j].to_bits(),
2089 "h[{i}][{j}] mismatch"
2090 );
2091 for k in 0..3 {
2092 assert_eq!(
2093 prog3.t3[i][j][k].to_bits(),
2094 prog4.t3[i][j][k].to_bits(),
2095 "t3[{i}][{j}][{k}] mismatch"
2096 );
2097 }
2098 }
2099 }
2100 }
2101
2102 /// Binomial-logit row NLL, K=1: ℓ(η) = ln(1 + e^η) − y·η.
2103 /// The entire tower has textbook closed forms in μ = σ(η); this test
2104 /// pins the algebra (exp, ln, scalar mixes, Leibniz/Faà di Bruno) to
2105 /// analytic truth at near-machine precision.
2106 struct LogitProgram {
2107 eta: Vec<f64>,
2108 y: Vec<f64>,
2109 }
2110
2111 impl RowProgram<1> for LogitProgram {
2112 fn n_rows(&self) -> usize {
2113 self.eta.len()
2114 }
2115 fn primaries(&self, row: usize) -> Result<[f64; 1], String> {
2116 Ok([self.eta[row]])
2117 }
2118 fn eval<S: crate::jet_scalar::JetScalar<1>>(
2119 &self,
2120 row: usize,
2121 p: &[S; 1],
2122 ) -> Result<S, String> {
2123 let eta = p[0];
2124 Ok(eta
2125 .exp()
2126 .add(&S::constant(1.0))
2127 .ln()
2128 .sub(&eta.scale(self.y[row])))
2129 }
2130 }
2131
2132 #[test]
2133 fn logit_tower_matches_closed_forms() {
2134 let prog = LogitProgram {
2135 eta: vec![-2.3, -0.4, 0.0, 0.9, 3.1],
2136 y: vec![1.0, 0.0, 1.0, 0.0, 1.0],
2137 };
2138 for row in 0..prog.n_rows() {
2139 let t = program_full_tower(&prog, row).expect("logit program");
2140 let eta = prog.eta[row];
2141 let y = prog.y[row];
2142 let mu = 1.0 / (1.0 + (-eta).exp());
2143 let w = mu * (1.0 - mu);
2144 let expect = [
2145 (t.v, (1.0 + eta.exp()).ln() - y * eta, "value"),
2146 (t.g[0], mu - y, "grad"),
2147 (t.h[0][0], w, "hess"),
2148 (t.t3[0][0][0], w * (1.0 - 2.0 * mu), "third"),
2149 (
2150 t.t4[0][0][0][0],
2151 w * (1.0 - 6.0 * mu + 6.0 * mu * mu),
2152 "fourth",
2153 ),
2154 ];
2155 for (got, want, label) in expect {
2156 assert!(
2157 (got - want).abs() <= 1e-12 * want.abs().max(1.0),
2158 "row {row} {label}: got {got:+.15e} want {want:+.15e}"
2159 );
2160 }
2161 }
2162 }
2163
2164 struct OversizedDenseProgram;
2165
2166 impl RowProgram<10> for OversizedDenseProgram {
2167 fn n_rows(&self) -> usize {
2168 1
2169 }
2170
2171 fn primaries(&self, row: usize) -> Result<[f64; 10], String> {
2172 Err(format!(
2173 "dense-tower storage check reached program primaries at row {row}"
2174 ))
2175 }
2176
2177 fn eval<S: crate::jet_scalar::JetScalar<10>>(
2178 &self,
2179 row: usize,
2180 primaries: &[S; 10],
2181 ) -> Result<S, String> {
2182 Err(format!(
2183 "dense-tower storage check reached program evaluation at row {row} with {} primaries",
2184 primaries.len()
2185 ))
2186 }
2187 }
2188
2189 struct LargestBudgetedDenseProgram;
2190
2191 impl RowProgram<9> for LargestBudgetedDenseProgram {
2192 fn n_rows(&self) -> usize {
2193 1
2194 }
2195
2196 fn primaries(&self, row: usize) -> Result<[f64; 9], String> {
2197 if row == 0 {
2198 Ok([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
2199 } else {
2200 Err(format!("largest budgeted dense program has no row {row}"))
2201 }
2202 }
2203
2204 fn eval<S: crate::jet_scalar::JetScalar<9>>(
2205 &self,
2206 row: usize,
2207 primaries: &[S; 9],
2208 ) -> Result<S, String> {
2209 if row != 0 {
2210 return Err(format!("largest budgeted dense program has no row {row}"));
2211 }
2212 let linear =
2213 S::linear_combination(primaries, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
2214 let quartic = primaries[0]
2215 .mul(&primaries[1])
2216 .mul(&primaries[2])
2217 .mul(&primaries[3]);
2218 Ok(linear.add(&quartic))
2219 }
2220 }
2221
2222 #[test]
2223 fn full_tower_accepts_largest_width_inside_storage_budget_932() {
2224 assert_eq!(std::mem::size_of::<Tower4<9>>(), 59_048);
2225 assert!(
2226 !program_primary_jets_fit_stack::<Tower4<9>, 9>(),
2227 "nine full-width primary towers must use exact-length heap storage"
2228 );
2229
2230 let tower = program_full_tower(&LargestBudgetedDenseProgram, 0)
2231 .expect("Tower4<9> must remain inside the canonical dense storage budget");
2232 assert_eq!(tower.v, 309.0);
2233 assert_eq!(tower.g, [25.0, 14.0, 11.0, 10.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
2234 // `t4` stores derivatives, not Taylor coefficients: the distinct-axis
2235 // derivative of p0*p1*p2*p3 is 1, with no 4! normalization.
2236 assert_eq!(tower.t4[0][1][2][3], 1.0);
2237 }
2238
2239 #[test]
2240 fn full_tower_refuses_oversized_result_before_touching_program() {
2241 let tower_bytes = std::mem::size_of::<Tower4<10>>();
2242 assert!(tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES);
2243 assert!(
2244 std::mem::size_of::<Result<Box<Tower4<32>>, String>>()
2245 <= 4 * std::mem::size_of::<usize>(),
2246 "boxed full-tower API must keep its return slot independent of dense tower width"
2247 );
2248
2249 let error = program_full_tower(&OversizedDenseProgram, 0)
2250 .expect_err("Tower4<10> must exceed the canonical dense storage budget");
2251 assert_eq!(
2252 error,
2253 format!(
2254 "canonical dense Tower4<10> requires {tower_bytes} bytes, exceeding the {}-byte \
2255 storage budget; use the bounded row-kernel and directional channel APIs",
2256 PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
2257 )
2258 );
2259 }
2260
2261 fn assert_close(label: &str, got: f64, want: f64, rel_tol: f64) {
2262 let diff = (got - want).abs();
2263 assert!(
2264 diff <= rel_tol * want.abs().max(1.0),
2265 "{label}: got {got:+.17e} want {want:+.17e} diff {diff:.3e}"
2266 );
2267 }
2268
2269 #[test]
2270 fn gamma_special_function_stacks_match_reference_values() {
2271 const EULER_GAMMA: f64 = 0.577_215_664_901_532_9;
2272 let pi_sq = std::f64::consts::PI * std::f64::consts::PI;
2273 let cases = [
2274 (
2275 "x=0.1",
2276 0.1,
2277 -10.423_754_940_411_076,
2278 101.433_299_150_792_75,
2279 ),
2280 (
2281 "x=0.5",
2282 0.5,
2283 -EULER_GAMMA - 2.0 * std::f64::consts::LN_2,
2284 pi_sq / 2.0,
2285 ),
2286 ("x=1", 1.0, -EULER_GAMMA, pi_sq / 6.0),
2287 (
2288 "x=2.5",
2289 2.5,
2290 -EULER_GAMMA - 2.0 * std::f64::consts::LN_2 + 2.0 + 2.0 / 3.0,
2291 pi_sq / 2.0 - 4.0 - 4.0 / 9.0,
2292 ),
2293 (
2294 "x=50",
2295 50.0,
2296 3.901_989_673_427_892,
2297 0.020_201_333_226_697_128,
2298 ),
2299 ];
2300
2301 for (label, x, digamma_ref, trigamma_ref) in cases {
2302 let ln_gamma_stack = ln_gamma_derivative_stack(x);
2303 let digamma_stack = digamma_derivative_stack(x);
2304 let trigamma_stack = trigamma_derivative_stack(x);
2305 assert_close(
2306 &format!("{label} ln_gamma_stack digamma"),
2307 ln_gamma_stack[1],
2308 digamma_ref,
2309 1e-13,
2310 );
2311 assert_close(
2312 &format!("{label} digamma value"),
2313 digamma_stack[0],
2314 digamma_ref,
2315 1e-13,
2316 );
2317 assert_close(
2318 &format!("{label} ln_gamma_stack trigamma"),
2319 ln_gamma_stack[2],
2320 trigamma_ref,
2321 1e-13,
2322 );
2323 assert_close(
2324 &format!("{label} digamma_stack trigamma"),
2325 digamma_stack[1],
2326 trigamma_ref,
2327 1e-13,
2328 );
2329 assert_close(
2330 &format!("{label} trigamma value"),
2331 trigamma_stack[0],
2332 trigamma_ref,
2333 1e-13,
2334 );
2335 }
2336 }
2337
2338 #[test]
2339 fn gamma_special_function_stacks_obey_recurrences() {
2340 for x in [0.1, 0.5, 1.0, 2.5, 50.0] {
2341 let digamma_x = digamma_derivative_stack(x)[0];
2342 let digamma_next = digamma_derivative_stack(x + 1.0)[0];
2343 let trigamma_x = trigamma_derivative_stack(x)[0];
2344 let trigamma_next = trigamma_derivative_stack(x + 1.0)[0];
2345 assert_close(
2346 &format!("digamma recurrence x={x}"),
2347 digamma_next,
2348 digamma_x + 1.0 / x,
2349 1e-13,
2350 );
2351 assert_close(
2352 &format!("trigamma recurrence x={x}"),
2353 trigamma_next,
2354 trigamma_x - 1.0 / (x * x),
2355 1e-13,
2356 );
2357 }
2358 }
2359
2360 /// Gaussian location-scale row NLL, K=2 primaries (η, s = log σ):
2361 /// ℓ = s + ½ e^{−2s} (y − η)². Mixed cross blocks — the #736 fragility
2362 /// shape — all have one-line closed forms here.
2363 struct LocScaleProgram {
2364 eta: Vec<f64>,
2365 s: Vec<f64>,
2366 y: Vec<f64>,
2367 }
2368
2369 impl RowProgram<2> for LocScaleProgram {
2370 fn n_rows(&self) -> usize {
2371 self.eta.len()
2372 }
2373 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2374 Ok([self.eta[row], self.s[row]])
2375 }
2376 fn eval<S: crate::jet_scalar::JetScalar<2>>(
2377 &self,
2378 row: usize,
2379 p: &[S; 2],
2380 ) -> Result<S, String> {
2381 let r = S::constant(self.y[row]).sub(&p[0]);
2382 Ok(p[1].add(&p[1].scale(-2.0).exp().mul(&r).mul(&r).scale(0.5)))
2383 }
2384 }
2385
2386 #[test]
2387 fn locscale_tower_matches_closed_forms_including_cross_blocks() {
2388 let prog = LocScaleProgram {
2389 eta: vec![0.3, -1.1, 2.0],
2390 s: vec![-0.5, 0.2, 0.8],
2391 y: vec![1.0, -2.0, 2.5],
2392 };
2393 let tol = 1e-12;
2394 for row in 0..prog.n_rows() {
2395 let t = program_full_tower(&prog, row).expect("locscale program");
2396 let r = prog.y[row] - prog.eta[row];
2397 let w = (-2.0 * prog.s[row]).exp();
2398 // (η, s) = indices (0, 1).
2399 let truth_g = [-w * r, 1.0 - w * r * r];
2400 let truth_h = [[w, 2.0 * w * r], [2.0 * w * r, 2.0 * w * r * r]];
2401 // Third tensor: distinct-entry closed forms.
2402 // ∂ηηη = 0, ∂ηηs = −2w, ∂ηss = −4wr, ∂sss = −4wr².
2403 let t3_truth = |a: usize, b: usize, c: usize| -> f64 {
2404 match a + b + c {
2405 0 => 0.0,
2406 1 => -2.0 * w,
2407 2 => -4.0 * w * r,
2408 _ => -4.0 * w * r * r,
2409 }
2410 };
2411 // Fourth tensor: ∂ηηηη = 0, ∂ηηηs = 0? No: d/ds(∂ηηη)=0 ✓;
2412 // ∂ηηss = 4w, ∂ηsss = 8wr, ∂ssss = 8wr².
2413 let t4_truth = |a: usize, b: usize, c: usize, d: usize| -> f64 {
2414 match a + b + c + d {
2415 0 | 1 => 0.0,
2416 2 => 4.0 * w,
2417 3 => 8.0 * w * r,
2418 _ => 8.0 * w * r * r,
2419 }
2420 };
2421 for a in 0..2 {
2422 assert!(
2423 (t.g[a] - truth_g[a]).abs() <= tol * truth_g[a].abs().max(1.0),
2424 "row {row} grad[{a}]"
2425 );
2426 for b in 0..2 {
2427 assert!(
2428 (t.h[a][b] - truth_h[a][b]).abs() <= tol * w.max(1.0) * (1.0 + r.abs()),
2429 "row {row} hess[{a}][{b}]: got {} want {}",
2430 t.h[a][b],
2431 truth_h[a][b]
2432 );
2433 for c in 0..2 {
2434 assert!(
2435 (t.t3[a][b][c] - t3_truth(a, b, c)).abs()
2436 <= tol * 8.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
2437 "row {row} t3[{a}][{b}][{c}]: got {} want {}",
2438 t.t3[a][b][c],
2439 t3_truth(a, b, c)
2440 );
2441 for d in 0..2 {
2442 assert!(
2443 (t.t4[a][b][c][d] - t4_truth(a, b, c, d)).abs()
2444 <= tol * 16.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
2445 "row {row} t4[{a}][{b}][{c}][{d}]: got {} want {}",
2446 t.t4[a][b][c][d],
2447 t4_truth(a, b, c, d)
2448 );
2449 }
2450 }
2451 }
2452 }
2453 // The canonical trait-surface helpers agree with direct contraction.
2454 let dir = [0.7, -1.3];
2455 let third = program_third_contracted(&prog, row, &dir).expect("third");
2456 for a in 0..2 {
2457 for b in 0..2 {
2458 let want = t.t3[a][b][0] * dir[0] + t.t3[a][b][1] * dir[1];
2459 assert!((third[a][b] - want).abs() <= 1e-13 * want.abs().max(1.0));
2460 }
2461 }
2462 }
2463 }
2464
2465 /// FD cross-check on a deliberately gnarly composition (div, sqrt,
2466 /// powf, nested exp/ln) in K=3, where no closed form is consulted:
2467 /// every tower channel is checked against central finite differences
2468 /// of the channel one order below — value→grad, grad→hess, hess→t3,
2469 /// t3→t4 — so each order is independently anchored.
2470 ///
2471 /// The program carries a per-row primary fixture plus a per-row offset
2472 /// `tau[row]` that enters the loss as a constant, so `row` genuinely
2473 /// drives both the seed point and the evaluated expression.
2474 struct GnarlyProgram {
2475 primaries: Vec<[f64; 3]>,
2476 tau: Vec<f64>,
2477 }
2478
2479 impl GnarlyProgram {
2480 fn fixture() -> Self {
2481 Self {
2482 primaries: vec![[0.4, -0.7, 1.2], [-0.9, 0.6, 0.3], [1.1, -0.2, -0.8]],
2483 tau: vec![0.15, -0.35, 0.5],
2484 }
2485 }
2486 }
2487
2488 impl RowProgram<3> for GnarlyProgram {
2489 fn n_rows(&self) -> usize {
2490 self.primaries.len()
2491 }
2492 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
2493 self.primaries
2494 .get(row)
2495 .copied()
2496 .ok_or_else(|| format!("gnarly: row {row} out of range"))
2497 }
2498 fn eval<S: crate::jet_scalar::JetScalar<3>>(
2499 &self,
2500 row: usize,
2501 p: &[S; 3],
2502 ) -> Result<S, String> {
2503 let tau = *self
2504 .tau
2505 .get(row)
2506 .ok_or_else(|| format!("gnarly: tau row {row} out of range"))?;
2507 let a = p[0].mul(&p[1]).exp();
2508 let b = p[2].mul(&p[2]).add(&S::constant(1.0)).sqrt();
2509 let c = a.add(&b).add(&S::constant(tau)).ln();
2510 let d = p[1].scale(0.5).add(&S::constant(2.0)).powf(1.7);
2511 let delta = p[0].sub(&p[2]);
2512 Ok(c.mul(&d.recip()).add(&delta.mul(&delta).scale(0.25)))
2513 }
2514 }
2515
2516 /// Evaluate the gnarly program's tower at an ARBITRARY seed point for
2517 /// `row` (used to drive central differences off the fixture grid),
2518 /// while keeping `row`'s per-row data (`tau`) in the loss.
2519 fn gnarly_tower_at(prog: &GnarlyProgram, row: usize, p: [f64; 3]) -> Tower4<3> {
2520 struct At<'a> {
2521 base: &'a GnarlyProgram,
2522 row: usize,
2523 p: [f64; 3],
2524 }
2525 impl RowProgram<3> for At<'_> {
2526 fn n_rows(&self) -> usize {
2527 1
2528 }
2529 fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
2530 if row != 0 {
2531 return Err(format!("gnarly-at: row {row} out of range"));
2532 }
2533 Ok(self.p)
2534 }
2535 fn eval<S: crate::jet_scalar::JetScalar<3>>(
2536 &self,
2537 eval_row: usize,
2538 vars: &[S; 3],
2539 ) -> Result<S, String> {
2540 if eval_row != 0 {
2541 return Err(format!("gnarly-at: eval row {eval_row} out of range"));
2542 }
2543 self.base.eval(self.row, vars)
2544 }
2545 }
2546 *program_full_tower(&At { base: prog, row, p }, 0).expect("gnarly tower")
2547 }
2548
2549 #[test]
2550 fn gnarly_tower_is_fd_consistent_order_by_order() {
2551 let prog = GnarlyProgram::fixture();
2552 for row in 0..prog.n_rows() {
2553 let base = prog.primaries(row).expect("primaries");
2554 let t = gnarly_tower_at(&prog, row, base);
2555 let h_step = 1e-5;
2556 let tol = 1e-6;
2557 for c in 0..3 {
2558 let mut up = base;
2559 let mut dn = base;
2560 up[c] += h_step;
2561 dn[c] -= h_step;
2562 let t_up = gnarly_tower_at(&prog, row, up);
2563 let t_dn = gnarly_tower_at(&prog, row, dn);
2564 // value → gradient.
2565 let fd_g = (t_up.v - t_dn.v) / (2.0 * h_step);
2566 assert!(
2567 (t.g[c] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
2568 "grad[{c}]: analytic {} fd {}",
2569 t.g[c],
2570 fd_g
2571 );
2572 for a in 0..3 {
2573 // gradient → Hessian.
2574 let fd_h = (t_up.g[a] - t_dn.g[a]) / (2.0 * h_step);
2575 assert!(
2576 (t.h[a][c] - fd_h).abs() <= tol * fd_h.abs().max(1.0),
2577 "hess[{a}][{c}]: analytic {} fd {}",
2578 t.h[a][c],
2579 fd_h
2580 );
2581 for b in 0..3 {
2582 // Hessian → third.
2583 let fd_t3 = (t_up.h[a][b] - t_dn.h[a][b]) / (2.0 * h_step);
2584 assert!(
2585 (t.t3[a][b][c] - fd_t3).abs() <= tol * fd_t3.abs().max(1.0),
2586 "t3[{a}][{b}][{c}]: analytic {} fd {}",
2587 t.t3[a][b][c],
2588 fd_t3
2589 );
2590 for d in 0..3 {
2591 // third → fourth.
2592 let fd_t4 = (t_up.t3[a][b][d] - t_dn.t3[a][b][d]) / (2.0 * h_step);
2593 assert!(
2594 (t.t4[a][b][d][c] - fd_t4).abs() <= tol * fd_t4.abs().max(1.0),
2595 "t4[{a}][{b}][{d}][{c}]: analytic {} fd {}",
2596 t.t4[a][b][d][c],
2597 fd_t4
2598 );
2599 }
2600 }
2601 }
2602 }
2603 }
2604 }
2605
2606 /// `implicit_solve` reproduces the true implicit function `a(θ)` of a
2607 /// constraint `F(a, θ) = 0` to fourth order. The constraint here is the
2608 /// smooth, strictly-`a`-monotone
2609 /// F(a, θ) = a + θ₀·a² + θ₁·exp(a) − c
2610 /// whose root `a(θ)` is re-solved by scalar Newton at perturbed θ as the
2611 /// independent finite-difference oracle. Mirrors the survival flex
2612 /// calibration solve (one implicit intercept over the primaries) without
2613 /// any survival machinery, so a failure localises to the combinator.
2614 #[test]
2615 fn implicit_solve_matches_scalar_resolve_to_fourth_order() {
2616 const C: f64 = 1.7;
2617 // The scalar constraint as a plain f64 closure (the production root
2618 // finder analogue) and its tower form in (a, θ₀, θ₁).
2619 let f_scalar = |a: f64, th: [f64; 2]| a + th[0] * a * a + th[1] * a.exp() - C;
2620 let f_da = |a: f64, th: [f64; 2]| 1.0 + 2.0 * th[0] * a + th[1] * a.exp();
2621 let solve = |th: [f64; 2]| -> f64 {
2622 let mut a = 0.0_f64;
2623 for _ in 0..100 {
2624 let r = f_scalar(a, th);
2625 if r.abs() < 1e-14 {
2626 break;
2627 }
2628 a -= r / f_da(a, th);
2629 }
2630 a
2631 };
2632 // Tower constraint over K1 = 3 vars: slot 0 = a, slots 1,2 = θ₀, θ₁.
2633 let f_tower = |a0: f64, th: [f64; 2]| -> Tower4<3> {
2634 let a = Tower4::<3>::variable(a0, 0);
2635 let t0 = Tower4::<3>::variable(th[0], 1);
2636 let t1 = Tower4::<3>::variable(th[1], 2);
2637 a + t0 * a.mul(&a) + t1 * a.exp() - C
2638 };
2639
2640 let th0 = [0.35, 0.2];
2641 let a0 = solve(th0);
2642 let f = f_tower(a0, th0);
2643 // Residual at the solved point is ~0 (the combinator tolerates the
2644 // production Newton residual; here it is machine-zero).
2645 assert!(f.v.abs() < 1e-12, "constraint residual {:+.3e}", f.v);
2646 let a_tower: Tower4<2> = implicit_solve::<3, 2>(&f, a0).expect("implicit solve");
2647
2648 // FD oracle: central differences of the scalar re-solve. Each order is
2649 // built from the previous via one more central difference, exactly the
2650 // gnarly order-by-order ladder.
2651 let h = 1e-4;
2652 let tol = 1e-5;
2653 let re = |th: [f64; 2]| solve(th);
2654 for i in 0..2 {
2655 let mut up = th0;
2656 let mut dn = th0;
2657 up[i] += h;
2658 dn[i] -= h;
2659 let fd_g = (re(up) - re(dn)) / (2.0 * h);
2660 assert!(
2661 (a_tower.g[i] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
2662 "a_θ[{i}]: analytic {:+.6e} fd {:+.6e}",
2663 a_tower.g[i],
2664 fd_g
2665 );
2666 // second order: FD of the analytic gradient component would re-use
2667 // the combinator; instead difference a SCALAR gradient computed by
2668 // a nested re-solve so the oracle stays production-independent.
2669 let grad_at = |th: [f64; 2], j: usize| -> f64 {
2670 let mut up = th;
2671 let mut dn = th;
2672 up[j] += h;
2673 dn[j] -= h;
2674 (re(up) - re(dn)) / (2.0 * h)
2675 };
2676 for j in 0..2 {
2677 let fd_h = (grad_at(up, j) - grad_at(dn, j)) / (2.0 * h);
2678 assert!(
2679 (a_tower.h[i][j] - fd_h).abs() <= 1e-3 * fd_h.abs().max(1.0),
2680 "a_θθ[{i}][{j}]: analytic {:+.6e} fd {:+.6e}",
2681 a_tower.h[i][j],
2682 fd_h
2683 );
2684 }
2685 }
2686 }
2687
2688 /// `implicit_solve` degenerates to `a_θ = −F_θ / F_a` at first order on a
2689 /// linear-in-a constraint, and the second-order tensor matches the
2690 /// textbook IFT formula `a_uv = −(F_uv + F_au a_v + F_av a_u + F_aa a_u a_v)/F_a`.
2691 /// This pins the recursion against the hand-coded first_full.rs formula it
2692 /// replaces, independent of any FD step.
2693 #[test]
2694 fn implicit_solve_matches_textbook_ift_recursion() {
2695 // A constraint with non-trivial F_a, F_aa, F_au, F_uv all present.
2696 let a0 = 0.4_f64;
2697 let th = [0.25_f64, -0.15_f64];
2698 let f = {
2699 let a = Tower4::<3>::variable(a0, 0);
2700 let t0 = Tower4::<3>::variable(th[0], 1);
2701 let t1 = Tower4::<3>::variable(th[1], 2);
2702 // F = a·(1 + θ₀) + θ₁·a² + θ₀·θ₁ − 0.4385. The constant is chosen so
2703 // F(a0, θ0) = 0 exactly at a0 = 0.4, θ = [0.25, −0.15]:
2704 // 0.4·1.25 + (−0.15)·0.16 + 0.25·(−0.15) = 0.4385.
2705 // implicit_solve requires a genuine root; at the root the level-set
2706 // and root-curve derivatives coincide, so the textbook-IFT
2707 // assertions below are unaffected.
2708 a * (t0 + 1.0) + t1 * a.mul(&a) + t0 * t1 - 0.4385
2709 };
2710 let a_t = implicit_solve::<3, 2>(&f, a0).expect("solve");
2711 let f_a = f.g[0];
2712 // First order: a_u = −F_u / F_a.
2713 for u in 0..2 {
2714 let want = -f.g[u + 1] / f_a;
2715 assert!(
2716 (a_t.g[u] - want).abs() < 1e-12,
2717 "a_u[{u}] {:+.6e} vs −F_u/F_a {:+.6e}",
2718 a_t.g[u],
2719 want
2720 );
2721 }
2722 // Second order textbook IFT (indices shifted by 1 for the a-slot).
2723 for u in 0..2 {
2724 for v in 0..2 {
2725 let f_uv = f.h[u + 1][v + 1];
2726 let f_au = f.h[0][u + 1];
2727 let f_av = f.h[0][v + 1];
2728 let f_aa = f.h[0][0];
2729 let want =
2730 -(f_uv + f_au * a_t.g[v] + f_av * a_t.g[u] + f_aa * a_t.g[u] * a_t.g[v]) / f_a;
2731 assert!(
2732 (a_t.h[u][v] - want).abs() < 1e-12,
2733 "a_uv[{u}][{v}] {:+.6e} vs IFT {:+.6e}",
2734 a_t.h[u][v],
2735 want
2736 );
2737 }
2738 }
2739 }
2740
2741 /// The moving-boundary flux tower reproduces every θ-derivative of a
2742 /// moving-limit integral, INCLUDING the second-order `B·z_uv` term the
2743 /// hand-written flux dropped (#932). The edge `z_R(θ) = θ₀ + θ₁²` has a
2744 /// genuinely nonzero `∂²z_R/∂θ₁² = 2`, so a combinator that omitted
2745 /// `B·z_uv` would miss the [1][1] Hessian entry. Truth = central FD of the
2746 /// closed-form integral `∫₀^{z_R} e^{−z²/2} dz = √(π/2)·erf(z_R/√2)`.
2747 #[test]
2748 fn moving_boundary_flux_carries_b_zuv_term() {
2749 use std::f64::consts::PI;
2750 let b = |z: f64| (-0.5 * z * z).exp(); // integrand B(z)
2751 // Antiderivative-based closed-form integral I(z_R) = ∫₀^{z_R} B dz.
2752 let integral = |z_r: f64| (PI / 2.0).sqrt() * libm::erf(z_r / 2.0_f64.sqrt());
2753 let z_r = |th: [f64; 2]| th[0] + th[1] * th[1];
2754 let th0 = [0.7_f64, 0.5_f64];
2755
2756 // Edge tower z_R(θ) over K=2 primaries: value + exact derivatives.
2757 let mut z_edge = Tower4::<2>::constant(z_r(th0));
2758 z_edge.g[0] = 1.0; // ∂z_R/∂θ₀ = 1
2759 z_edge.g[1] = 2.0 * th0[1]; // ∂z_R/∂θ₁ = 2θ₁
2760 z_edge.h[1][1] = 2.0; // ∂²z_R/∂θ₁² = 2 (the z_uv the old flux dropped)
2761
2762 // Integrand stack [B, B′, B″, B‴] at z₀: B′=−z·B, B″=(z²−1)·B,
2763 // B‴=(3z−z³)·B.
2764 let z0 = z_edge.v;
2765 let b0 = b(z0);
2766 let stack = [
2767 b0,
2768 -z0 * b0,
2769 (z0 * z0 - 1.0) * b0,
2770 (3.0 * z0 - z0 * z0 * z0) * b0,
2771 ];
2772 let flux = moving_limit_boundary_tower(&z_edge, stack);
2773
2774 // FD truth of the integral's derivatives.
2775 let h = 1e-4;
2776 let tol = 1e-6;
2777 for i in 0..2 {
2778 let mut up = th0;
2779 let mut dn = th0;
2780 up[i] += h;
2781 dn[i] -= h;
2782 let fd_g = (integral(z_r(up)) - integral(z_r(dn))) / (2.0 * h);
2783 assert!(
2784 (flux.g[i] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
2785 "flux_g[{i}]: analytic {:+.8e} fd {:+.8e}",
2786 flux.g[i],
2787 fd_g
2788 );
2789 }
2790 // The decisive entry: ∂²I/∂θ₁² = B′·(z_θ₁)² + B·z_θ₁θ₁. With z_θ₁=2θ₁=1
2791 // and z_θ₁θ₁=2, the B·z_uv contribution is B(z₀)·2 — omitting it would
2792 // leave the [1][1] entry short by exactly 2·B(z₀).
2793 let grad1_at = |th: [f64; 2]| -> f64 {
2794 let mut up = th;
2795 let mut dn = th;
2796 up[1] += h;
2797 dn[1] -= h;
2798 (integral(z_r(up)) - integral(z_r(dn))) / (2.0 * h)
2799 };
2800 let mut up = th0;
2801 let mut dn = th0;
2802 up[1] += h;
2803 dn[1] -= h;
2804 let fd_h11 = (grad1_at(up) - grad1_at(dn)) / (2.0 * h);
2805 assert!(
2806 (flux.h[1][1] - fd_h11).abs() <= 1e-3 * fd_h11.abs().max(1.0),
2807 "flux_h[1][1] (carries B·z_uv): analytic {:+.8e} fd {:+.8e}",
2808 flux.h[1][1],
2809 fd_h11
2810 );
2811 // Explicit witness that the B·z_uv term is present and material:
2812 // analytic h[1][1] minus the pure (z_u)² part must equal B·z_uv = 2·B₀.
2813 let pure_zu2 = stack[1] * z_edge.g[1] * z_edge.g[1];
2814 let b_zuv = flux.h[1][1] - pure_zu2;
2815 assert!(
2816 (b_zuv - b0 * 2.0).abs() < 1e-10,
2817 "B·z_uv term {:+.8e} != B₀·z_uv {:+.8e}",
2818 b_zuv,
2819 b0 * 2.0
2820 );
2821 }
2822
2823 /// `moving_limit_boundary_tower_theta_integrand` reproduces the marginal-slope
2824 /// flex boundary closure for a θ-DEPENDENT integrand `G(z;θ)` — the case the
2825 /// plain `moving_limit_boundary_tower` cannot express, and the case the
2826 /// survival directional/bidirectional paths hand-assemble term-by-term
2827 /// (`G·z_uv + G_z·z_u·z_v + G_θu·z_v + G_θv·z_u`, with the directional path
2828 /// dropping `G·z_uv`). Two independent oracles:
2829 /// (1) closed-form: the boundary flux of `∫ G dz` is exactly
2830 /// `Φ(z_edge(θ);θ) − Φ(z₀;θ)` (Φ = z-antiderivative of G), whose θ
2831 /// derivatives we take by central FD of the closed form — no jet code.
2832 /// (2) the explicit second-order hand closure, including the `G·z_uv` term,
2833 /// built from the integrand's own (z,θ) partials.
2834 /// G(z;θ) = exp(z·θ₀) is genuinely θ-dependent (G_θ₀ = z·e^{zθ₀} ≠ 0), and
2835 /// the edge z_edge = z₀ + θ₀ + θ₁² has a real z_uv = ∂²/∂θ₁² = 2, so a
2836 /// combinator that dropped either the integrand-θ terms or `G·z_uv` would
2837 /// miss a Hessian entry.
2838 #[test]
2839 fn moving_boundary_theta_integrand_matches_handpath_and_closed_form() {
2840 // G(z;θ) = exp(z·θ₀); Φ(z;θ) = ∫₀^z G = (e^{zθ₀} − 1)/θ₀.
2841 let g = |z: f64, t0: f64| (z * t0).exp();
2842 let phi = |z: f64, t0: f64| ((z * t0).exp() - 1.0) / t0;
2843 let z_r = |th: [f64; 2]| 0.6 + th[0] + th[1] * th[1];
2844 let th0 = [0.4_f64, 0.5_f64];
2845 let z0 = z_r(th0);
2846
2847 // Edge tower z_edge(θ) over K=2 primaries.
2848 let mut z_edge = Tower4::<2>::constant(z0);
2849 z_edge.g[0] = 1.0; // ∂z/∂θ₀
2850 z_edge.g[1] = 2.0 * th0[1]; // ∂z/∂θ₁
2851 z_edge.h[1][1] = 2.0; // ∂²z/∂θ₁² (the z_uv the directional path drops)
2852
2853 // Φ's mixed (z, θ) jet over K1 = 3 vars: slot 0 = z, slots 1,2 = θ₀,θ₁.
2854 // Built ONCE in tower arithmetic so every (z^i θ^j) partial is exact.
2855 let z_var = Tower4::<3>::variable(z0, 0);
2856 let t0_var = Tower4::<3>::variable(th0[0], 1);
2857 // θ₁ does not enter G/Φ here (its Φ-derivatives are zero; the z_edge
2858 // chain supplies all θ₁ motion through slot 0), so the K1 frame's θ₁
2859 // slot is intentionally left unseeded.
2860 let phi_jet = ((z_var * t0_var).exp() - 1.0) / t0_var;
2861 // Sanity: slot-0 first derivative of Φ IS G(z₀;θ₀).
2862 assert!(
2863 (phi_jet.g[0] - g(z0, th0[0])).abs() < 1e-12,
2864 "Φ_z {:+.8e} != G {:+.8e}",
2865 phi_jet.g[0],
2866 g(z0, th0[0])
2867 );
2868
2869 let flux = moving_limit_boundary_tower_theta_integrand::<3, 2>(&phi_jet, &z_edge);
2870
2871 // Value channel is 0 by construction (boundary, not the integral itself).
2872 assert!(
2873 flux.v.abs() < 1e-12,
2874 "boundary value channel {:+.3e}",
2875 flux.v
2876 );
2877
2878 // Oracle (1): central FD of the closed-form boundary flux
2879 // Bnd(θ) = Φ(z_edge(θ); θ) − Φ(z₀; θ) (z₀ FROZEN at the base edge).
2880 let bnd = |th: [f64; 2]| phi(z_r(th), th[0]) - phi(z0, th[0]);
2881 let h = 1e-4;
2882 let tol = 1e-6;
2883 for i in 0..2 {
2884 let mut up = th0;
2885 let mut dn = th0;
2886 up[i] += h;
2887 dn[i] -= h;
2888 let fd_g = (bnd(up) - bnd(dn)) / (2.0 * h);
2889 assert!(
2890 (flux.g[i] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
2891 "boundary_g[{i}] analytic {:+.8e} fd {:+.8e}",
2892 flux.g[i],
2893 fd_g
2894 );
2895 }
2896 let grad_at = |th: [f64; 2], j: usize| -> f64 {
2897 let mut up = th;
2898 let mut dn = th;
2899 up[j] += h;
2900 dn[j] -= h;
2901 (bnd(up) - bnd(dn)) / (2.0 * h)
2902 };
2903 for i in 0..2 {
2904 for j in 0..2 {
2905 let mut up = th0;
2906 let mut dn = th0;
2907 up[i] += h;
2908 dn[i] -= h;
2909 let fd_h = (grad_at(up, j) - grad_at(dn, j)) / (2.0 * h);
2910 assert!(
2911 (flux.h[i][j] - fd_h).abs() <= 1e-3 * fd_h.abs().max(1.0),
2912 "boundary_h[{i}][{j}] analytic {:+.8e} fd {:+.8e}",
2913 flux.h[i][j],
2914 fd_h
2915 );
2916 }
2917 }
2918
2919 // Oracle (2): the explicit second-order hand closure, term by term —
2920 // `G·z_uv + G_z·z_u·z_v + G_θu·z_v + G_θv·z_u`. Read G's partials at the
2921 // base point directly (no jet): G = e^{zθ₀}, G_z = θ₀·G, G_θ₀ = z·G,
2922 // G_θ₁ = 0.
2923 let gg = g(z0, th0[0]);
2924 let g_z = th0[0] * gg;
2925 let g_theta = [z0 * gg, 0.0]; // [G_θ₀, G_θ₁]
2926 for i in 0..2 {
2927 for j in 0..2 {
2928 let z_u = z_edge.g[i];
2929 let z_v = z_edge.g[j];
2930 let z_uv = z_edge.h[i][j];
2931 let hand = gg * z_uv + g_z * z_u * z_v + g_theta[i] * z_v + g_theta[j] * z_u;
2932 assert!(
2933 (flux.h[i][j] - hand).abs() < 1e-9,
2934 "boundary_h[{i}][{j}] {:+.8e} != hand closure {:+.8e}",
2935 flux.h[i][j],
2936 hand
2937 );
2938 }
2939 }
2940
2941 // Decisive: the `G·z_uv` term the directional path DROPS is present and
2942 // material in the [1][1] entry (z_uv = 2 there).
2943 let pure_no_zuv = g_z * z_edge.g[1] * z_edge.g[1] + 2.0 * g_theta[1] * z_edge.g[1];
2944 let g_zuv = flux.h[1][1] - pure_no_zuv;
2945 assert!(
2946 (g_zuv - gg * 2.0).abs() < 1e-9,
2947 "G·z_uv term {:+.8e} != G₀·z_uv {:+.8e}",
2948 g_zuv,
2949 gg * 2.0
2950 );
2951 }
2952
2953 /// The survival crossing-edge position tower `z_edge = (τ − a(θ)) / b`,
2954 /// `b = exp(g)`, built from the intercept tower `a(θ)` (here a stand-in)
2955 /// and the seeded slope `g`, reproduces taylor-jet's exact hand-path
2956 /// boundary-velocity formulas:
2957 /// z_u = −(a_u + [u==g]·z) / b
2958 /// z_uv = −(a_uv + [u==g]·z_v + [v==g]·z_u) / b
2959 /// This pins the bridge between `implicit_solve` and
2960 /// `cell_moving_boundary_flux_tower`: the boundary jet that the production
2961 /// flex path hand-codes (and dropped `z_uv` from) is exactly `∂²` of this
2962 /// tower. K=3 reduced frame: slot 0 = a-axis carrier (an arbitrary smooth
2963 /// a(θ) with nonzero a_u/a_uv), slot 1 = g (the log-slope), slot 2 unused.
2964 #[test]
2965 fn crossing_edge_tower_matches_handpath_velocity_formulas() {
2966 const TAU: f64 = 1.3; // the link-knot crossing threshold τ
2967 let g_idx = 1usize;
2968 let g0 = 0.85_f64; // the slope value b (the g-primary IS the slope)
2969 // Stand-in intercept tower a(θ): nonzero value, gradient, Hessian in the
2970 // two live axes so a_u and a_uv are both exercised. (In production this
2971 // comes from implicit_solve; here we plant known derivatives.)
2972 let mut a = Tower4::<3>::constant(0.45);
2973 a.g[0] = 0.7;
2974 a.g[1] = -0.3;
2975 a.h[0][0] = 0.25;
2976 a.h[0][1] = 0.11;
2977 a.h[1][0] = 0.11;
2978 a.h[1][1] = -0.08;
2979
2980 // In the survival flex frame the slope `b` IS the g-primary directly
2981 // (the directional code passes `g` as `b`, and ∂z/∂g uses ∂b/∂g = 1):
2982 // z_edge = (τ − a) / b with b seeded as the g-axis variable.
2983 let b = Tower4::<3>::variable(g0, g_idx);
2984 let z_edge = (Tower4::<3>::constant(TAU) - a) / b;
2985
2986 let bv = g0;
2987 let z0 = z_edge.v;
2988 assert!((z0 - (TAU - 0.45) / bv).abs() < 1e-12);
2989
2990 // z_u = −(a_u + [u==g]·z) / b.
2991 for u in 0..2 {
2992 let direct = if u == g_idx { z0 } else { 0.0 };
2993 let want = -(a.g[u] + direct) / bv;
2994 assert!(
2995 (z_edge.g[u] - want).abs() < 1e-10,
2996 "z_u[{u}] {:+.8e} vs hand formula {:+.8e}",
2997 z_edge.g[u],
2998 want
2999 );
3000 }
3001 // z_uv = −(a_uv + [u==g]·z_v + [v==g]·z_u) / b, using the tower's own
3002 // first-order z_v/z_u (already verified above).
3003 for u in 0..2 {
3004 for v in 0..2 {
3005 let cross = if u == g_idx { z_edge.g[v] } else { 0.0 }
3006 + if v == g_idx { z_edge.g[u] } else { 0.0 };
3007 let want = -(a.h[u][v] + cross) / bv;
3008 assert!(
3009 (z_edge.h[u][v] - want).abs() < 1e-10,
3010 "z_uv[{u}][{v}] {:+.8e} vs hand formula {:+.8e}",
3011 z_edge.h[u][v],
3012 want
3013 );
3014 }
3015 }
3016 }
3017
3018 /// The crossing-edge tower in the CONSTRAINT frame (intercept `a` and
3019 /// slope `b` BOTH independent — slots 0 and 1) reproduces taylor-jet's
3020 /// FD-certified bare boundary-velocity constants exactly:
3021 /// z_a = ∂z/∂a = −1/b
3022 /// z_ab = ∂²z/∂a∂b = +1/b²
3023 /// z_aa = ∂²z/∂a² = 0
3024 /// z_bb = ∂²z/∂b² = +2(τ−a)/b³
3025 /// These are the `f_a`/`f_au`/`f_aa` constraint-jet boundary motions the
3026 /// production base path drops (and only adds in the dir twins, causing the
3027 /// #932 desync). Here `a` is independent (NOT yet substituted with a(θ)),
3028 /// so `z_aa = 0` and there is no `a_uv` chain — `implicit_solve` introduces
3029 /// that later. Pins the constant before the constraint-tower wiring.
3030 #[test]
3031 fn crossing_edge_constraint_frame_matches_bare_velocity_constants() {
3032 const TAU: f64 = 1.3;
3033 let a0 = 0.45_f64;
3034 let b0 = 0.85_f64;
3035 // Slot 0 = a, slot 1 = b, both seeded independent.
3036 let a = Tower4::<2>::variable(a0, 0);
3037 let b = Tower4::<2>::variable(b0, 1);
3038 let z = (Tower4::<2>::constant(TAU) - a) / b;
3039
3040 assert!((z.v - (TAU - a0) / b0).abs() < 1e-12);
3041 assert!((z.g[0] - (-1.0 / b0)).abs() < 1e-12, "z_a {:+.10e}", z.g[0]);
3042 assert!(
3043 (z.h[0][1] - 1.0 / (b0 * b0)).abs() < 1e-12,
3044 "z_ab {:+.10e} vs +1/b² {:+.10e}",
3045 z.h[0][1],
3046 1.0 / (b0 * b0)
3047 );
3048 assert!(
3049 z.h[0][0].abs() < 1e-12,
3050 "z_aa must vanish, got {:+.10e}",
3051 z.h[0][0]
3052 );
3053 let want_zbb = 2.0 * (TAU - a0) / (b0 * b0 * b0);
3054 assert!(
3055 (z.h[1][1] - want_zbb).abs() < 1e-12,
3056 "z_bb {:+.10e} vs 2(τ−a)/b³ {:+.10e}",
3057 z.h[1][1],
3058 want_zbb
3059 );
3060 }
3061
3062 /// The oracle harness catches a planted #736-style sign flip in a
3063 /// cross block and reports the channel by name.
3064 #[test]
3065 fn oracle_catches_planted_cross_block_sign_flip() {
3066 let prog = LocScaleProgram {
3067 eta: vec![0.3],
3068 s: vec![-0.5],
3069 y: vec![1.0],
3070 };
3071 let t = program_full_tower(&prog, 0).expect("tower");
3072 let dir = [0.6, -0.2];
3073 let mut third = t.third_contracted(&dir);
3074 let honest = KernelChannels {
3075 value: t.v,
3076 gradient: t.g,
3077 hessian: t.h,
3078 third: vec![(dir, third)],
3079 fourth: vec![(dir, [1.0, 0.5], t.fourth_contracted(&dir, &[1.0, 0.5]))],
3080 };
3081 verify_kernel_channels(&t, &honest, 1e-10).expect("honest kernel must pass");
3082
3083 // Plant the #736 flip: negate one mixed cross entry.
3084 third[0][1] = -third[0][1];
3085 let flipped = KernelChannels {
3086 value: t.v,
3087 gradient: t.g,
3088 hessian: t.h,
3089 third: vec![(dir, third)],
3090 fourth: vec![],
3091 };
3092 let err = verify_kernel_channels(&t, &flipped, 1e-10)
3093 .expect_err("planted sign flip must be caught");
3094 assert!(
3095 err.contains("third[0][0][1]"),
3096 "oracle must name the flipped channel, got: {err}"
3097 );
3098 }
3099
3100 /// The third- and fourth-order tensors must be FULLY symmetric under
3101 /// index permutation (mixed partials commute). The tower stores them
3102 /// unsymmetrized, so equal-by-construction is a real invariant of the
3103 /// Leibniz/Faà di Bruno writes — a cheap typo tripwire. Asserted on a
3104 /// nontrivial K=3 tower with all of div/sqrt/powf/exp/ln exercised, so
3105 /// every composition path contributes. Lives in a test (not the hot
3106 /// per-op path) on purpose.
3107 #[test]
3108 fn t3_t4_are_fully_index_symmetric() {
3109 let prog = GnarlyProgram::fixture();
3110 // 3! = 6 permutations of three indices.
3111 let perms3: [[usize; 3]; 6] = [
3112 [0, 1, 2],
3113 [0, 2, 1],
3114 [1, 0, 2],
3115 [1, 2, 0],
3116 [2, 0, 1],
3117 [2, 1, 0],
3118 ];
3119 // 4! = 24 permutations of four indices.
3120 let perms4: [[usize; 4]; 24] = [
3121 [0, 1, 2, 3],
3122 [0, 1, 3, 2],
3123 [0, 2, 1, 3],
3124 [0, 2, 3, 1],
3125 [0, 3, 1, 2],
3126 [0, 3, 2, 1],
3127 [1, 0, 2, 3],
3128 [1, 0, 3, 2],
3129 [1, 2, 0, 3],
3130 [1, 2, 3, 0],
3131 [1, 3, 0, 2],
3132 [1, 3, 2, 0],
3133 [2, 0, 1, 3],
3134 [2, 0, 3, 1],
3135 [2, 1, 0, 3],
3136 [2, 1, 3, 0],
3137 [2, 3, 0, 1],
3138 [2, 3, 1, 0],
3139 [3, 0, 1, 2],
3140 [3, 0, 2, 1],
3141 [3, 1, 0, 2],
3142 [3, 1, 2, 0],
3143 [3, 2, 0, 1],
3144 [3, 2, 1, 0],
3145 ];
3146 for row in 0..prog.n_rows() {
3147 let t = program_full_tower(&prog, row).expect("gnarly tower");
3148 let scale_t3 =
3149 t.t3.iter()
3150 .flatten()
3151 .flatten()
3152 .fold(0.0_f64, |m, x| m.max(x.abs()))
3153 .max(1.0);
3154 let scale_t4 =
3155 t.t4.iter()
3156 .flatten()
3157 .flatten()
3158 .flatten()
3159 .fold(0.0_f64, |m, x| m.max(x.abs()))
3160 .max(1.0);
3161 for i in 0..3 {
3162 for j in 0..3 {
3163 for k in 0..3 {
3164 let base = t.t3[i][j][k];
3165 let idx = [i, j, k];
3166 for p in &perms3 {
3167 let permed = t.t3[idx[p[0]]][idx[p[1]]][idx[p[2]]];
3168 assert!(
3169 (base - permed).abs() <= 1e-12 * scale_t3,
3170 "row {row}: t3[{i}][{j}][{k}]={base:+.15e} != \
3171 permuted {permed:+.15e} under {p:?}"
3172 );
3173 }
3174 for l in 0..3 {
3175 let base4 = t.t4[i][j][k][l];
3176 let idx4 = [i, j, k, l];
3177 for p in &perms4 {
3178 let permed = t.t4[idx4[p[0]]][idx4[p[1]]][idx4[p[2]]][idx4[p[3]]];
3179 assert!(
3180 (base4 - permed).abs() <= 1e-12 * scale_t4,
3181 "row {row}: t4[{i}][{j}][{k}][{l}]={base4:+.15e} != \
3182 permuted {permed:+.15e} under {p:?}"
3183 );
3184 }
3185 }
3186 }
3187 }
3188 }
3189 }
3190 }
3191}
3192
3193/// Stable derivative stack for `log Phi(x)` through fourth order.
3194#[inline]
3195pub fn unary_derivatives_normal_logcdf(x: f64) -> [f64; 5] {
3196 crate::probability::normal_logcdf_derivatives(x)
3197}
3198
3199/// Stable derivative stack for `log(1 - exp(-x))`, `x > 0`, through fourth order.
3200#[inline]
3201pub fn unary_derivatives_log1mexp_positive(x: f64) -> [f64; 5] {
3202 let r = 1.0 / x.exp_m1();
3203 [
3204 crate::probability::log1mexp_positive(x),
3205 r,
3206 -r * (1.0 + r),
3207 r * (1.0 + r) * (1.0 + 2.0 * r),
3208 -r * (1.0 + r) * (1.0 + 6.0 * r + 6.0 * r * r),
3209 ]
3210}
3211#[cfg(test)]
3212mod derivative_stack_tests {
3213 use super::*;
3214 // ── ln_gamma_derivative_stack / digamma_derivative_stack / trigamma_derivative_stack ──
3215
3216 #[test]
3217 fn ln_gamma_derivative_stack_known_values_at_1() {
3218 let s = ln_gamma_derivative_stack(1.0);
3219 // ln Γ(1) = 0; statrs uses Lanczos so the result is within ULP noise
3220 assert!(s[0].abs() < 1e-14, "ln_gamma(1) must be ~0, got {}", s[0]);
3221 // ψ₀(1) = -γ (Euler–Mascheroni)
3222 let euler_mascheroni = 0.577_215_664_901_532_9_f64;
3223 assert!(
3224 (s[1] + euler_mascheroni).abs() < 1e-10,
3225 "digamma(1) ≈ -{euler_mascheroni:.6}, got {}",
3226 s[1]
3227 );
3228 // ψ₁(1) = π²/6
3229 let pi2_6 = std::f64::consts::PI * std::f64::consts::PI / 6.0;
3230 assert!(
3231 (s[2] - pi2_6).abs() < 1e-10,
3232 "trigamma(1) ≈ {pi2_6:.6}, got {}",
3233 s[2]
3234 );
3235 }
3236
3237 #[test]
3238 fn ln_gamma_derivative_stack_known_values_at_2() {
3239 let s = ln_gamma_derivative_stack(2.0);
3240 // ln Γ(2) = ln(1) = 0 exactly
3241 assert!(s[0].abs() < 1e-14, "ln_gamma(2) must be 0, got {}", s[0]);
3242 // ψ₀(2) = 1 − γ (recurrence: ψ₀(x+1) = ψ₀(x) + 1/x)
3243 let euler_mascheroni = 0.577_215_664_901_532_9_f64;
3244 let digamma_2 = 1.0 - euler_mascheroni;
3245 assert!(
3246 (s[1] - digamma_2).abs() < 1e-10,
3247 "digamma(2) ≈ {digamma_2:.6}, got {}",
3248 s[1]
3249 );
3250 }
3251
3252 #[test]
3253 fn ln_gamma_derivative_stack_order2_is_prefix() {
3254 for &x in &[0.5_f64, 1.0, 2.0, 5.0] {
3255 let full = ln_gamma_derivative_stack(x);
3256 let ord2 = ln_gamma_derivative_stack_order2(x);
3257 assert_eq!(ord2[0], full[0], "order2[0] != full[0] at x={x}");
3258 assert_eq!(ord2[1], full[1], "order2[1] != full[1] at x={x}");
3259 assert_eq!(ord2[2], full[2], "order2[2] != full[2] at x={x}");
3260 }
3261 }
3262
3263 #[test]
3264 fn digamma_derivative_stack_overlaps_ln_gamma_stack() {
3265 // The two stacks share a run of four polygamma values:
3266 // ln_gamma_stack[1..5] == digamma_stack[0..4]
3267 for &x in &[0.5_f64, 1.0, 2.0, 7.0] {
3268 let lg = ln_gamma_derivative_stack(x);
3269 let dg = digamma_derivative_stack(x);
3270 for i in 0..4 {
3271 assert_eq!(
3272 lg[i + 1],
3273 dg[i],
3274 "ln_gamma_stack[{}] != digamma_stack[{}] at x={x}",
3275 i + 1,
3276 i
3277 );
3278 }
3279 }
3280 }
3281
3282 #[test]
3283 fn trigamma_derivative_stack_overlaps_digamma_stack() {
3284 // digamma_stack[1..5] == trigamma_stack[0..4]
3285 for &x in &[0.5_f64, 1.0, 2.0, 7.0] {
3286 let dg = digamma_derivative_stack(x);
3287 let tg = trigamma_derivative_stack(x);
3288 for i in 0..4 {
3289 assert_eq!(
3290 dg[i + 1],
3291 tg[i],
3292 "digamma_stack[{}] != trigamma_stack[{}] at x={x}",
3293 i + 1,
3294 i
3295 );
3296 }
3297 }
3298 }
3299
3300 #[test]
3301 fn derivative_stacks_all_finite_at_positive_inputs() {
3302 for &x in &[0.01_f64, 0.5, 1.0, 2.0, 10.0, 100.0] {
3303 for v in ln_gamma_derivative_stack(x) {
3304 assert!(v.is_finite(), "ln_gamma_stack non-finite at x={x}: {v}");
3305 }
3306 for v in digamma_derivative_stack(x) {
3307 assert!(v.is_finite(), "digamma_stack non-finite at x={x}: {v}");
3308 }
3309 for v in trigamma_derivative_stack(x) {
3310 assert!(v.is_finite(), "trigamma_stack non-finite at x={x}: {v}");
3311 }
3312 }
3313 }
3314}
3315
3316// ── Contraction-symmetry optimization gate ────────────────────────────────────
3317//
3318// `Tower4::third_contracted` / `fourth_contracted` contract the (fully
3319// index-symmetric) `t3`/`t4` tensors against directions, leaving the output
3320// indices `(a, b)` / `(i, j)` free. Those free indices inherit the tensor's
3321// symmetry — `out[a][b] == out[b][a]` term-for-term — so only the upper triangle
3322// need be summed and the lower triangle mirrored. Unlike the dense symmetric
3323// FILL (which needs a K⁴ scatter and loses inner-loop vectorisation, and was
3324// measured SLOWER), the mirror here is a tiny K×K copy and the inner contraction
3325// is untouched (contiguous, vectorisable). This is BIT-IDENTICAL to the full
3326// nest, so it needs no fingerprint re-baseline; the gate is (1) bit-identity vs
3327// the full reference and (2) a measured wall-clock that is not slower.
3328#[cfg(test)]
3329mod contraction_symmetry_tests {
3330 use super::*;
3331
3332 struct Rng(u64);
3333 impl Rng {
3334 fn u(&mut self) -> f64 {
3335 self.0 = self
3336 .0
3337 .wrapping_mul(6364136223846793005)
3338 .wrapping_add(1442695040888963407);
3339 (self.0 >> 11) as f64 / (1u64 << 53) as f64
3340 }
3341 fn s(&mut self) -> f64 {
3342 (self.u() - 0.5) * 4.0
3343 }
3344 }
3345
3346 /// Random VALID fully-symmetric `Tower4<K>` (symmetric `h`/`t3`/`t4`).
3347 fn rand_sym4<const K: usize>(r: &mut Rng) -> Tower4<K> {
3348 let mut t = Tower4::<K>::zero();
3349 t.v = r.s();
3350 for i in 0..K {
3351 t.g[i] = r.s();
3352 }
3353 for a in 0..K {
3354 for b in a..K {
3355 let v2 = r.s();
3356 t.h[a][b] = v2;
3357 t.h[b][a] = v2;
3358 for c in b..K {
3359 let v3 = r.s();
3360 for p in perms3([a, b, c]) {
3361 t.t3[p[0]][p[1]][p[2]] = v3;
3362 }
3363 for d in c..K {
3364 let v4 = r.s();
3365 for p in perms4([a, b, c, d]) {
3366 t.t4[p[0]][p[1]][p[2]][p[3]] = v4;
3367 }
3368 }
3369 }
3370 }
3371 }
3372 t
3373 }
3374
3375 fn perms3(idx: [usize; 3]) -> [[usize; 3]; 6] {
3376 let [a, b, c] = idx;
3377 [
3378 [a, b, c],
3379 [a, c, b],
3380 [b, a, c],
3381 [b, c, a],
3382 [c, a, b],
3383 [c, b, a],
3384 ]
3385 }
3386 fn perms4(idx: [usize; 4]) -> [[usize; 4]; 24] {
3387 let [a, b, c, d] = idx;
3388 [
3389 [a, b, c, d],
3390 [a, b, d, c],
3391 [a, c, b, d],
3392 [a, c, d, b],
3393 [a, d, b, c],
3394 [a, d, c, b],
3395 [b, a, c, d],
3396 [b, a, d, c],
3397 [b, c, a, d],
3398 [b, c, d, a],
3399 [b, d, a, c],
3400 [b, d, c, a],
3401 [c, a, b, d],
3402 [c, a, d, b],
3403 [c, b, a, d],
3404 [c, b, d, a],
3405 [c, d, a, b],
3406 [c, d, b, a],
3407 [d, a, b, c],
3408 [d, a, c, b],
3409 [d, b, a, c],
3410 [d, b, c, a],
3411 [d, c, a, b],
3412 [d, c, b, a],
3413 ]
3414 }
3415
3416 /// Full-nest reference (the pre-opt `a, b ∈ 0..K` form).
3417 fn third_full<const K: usize>(t: &Tower4<K>, dir: &[f64; K]) -> [[f64; K]; K] {
3418 let mut out = [[0.0; K]; K];
3419 for a in 0..K {
3420 for b in 0..K {
3421 let mut acc = 0.0;
3422 for c in 0..K {
3423 acc += t.t3[a][b][c] * dir[c];
3424 }
3425 out[a][b] = acc;
3426 }
3427 }
3428 out
3429 }
3430 fn fourth_full<const K: usize>(t: &Tower4<K>, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
3431 let mut out = [[0.0; K]; K];
3432 for i in 0..K {
3433 for j in 0..K {
3434 let mut acc = 0.0;
3435 for k in 0..K {
3436 for l in 0..K {
3437 acc += t.t4[i][j][k][l] * u[k] * w[l];
3438 }
3439 }
3440 out[i][j] = acc;
3441 }
3442 }
3443 out
3444 }
3445
3446 /// Returns the number of bit-equality comparisons performed (`n·K·K·2`), so
3447 /// the caller can assert the intended workload actually ran: a generic
3448 /// (turbofish) helper call hides its internal assertions, so the count is
3449 /// surfaced and checked at the call site.
3450 fn check_bit_identical<const K: usize>(seed: u64, n: usize) -> usize {
3451 let mut r = Rng(seed);
3452 let mut checks = 0usize;
3453 for _ in 0..n {
3454 let t = rand_sym4::<K>(&mut r);
3455 let dir: [f64; K] = std::array::from_fn(|_| r.s());
3456 let u: [f64; K] = std::array::from_fn(|_| r.s());
3457 let w: [f64; K] = std::array::from_fn(|_| r.s());
3458 let t3_sym = t.third_contracted(&dir);
3459 let t3_full = third_full(&t, &dir);
3460 let t4_sym = t.fourth_contracted(&u, &w);
3461 let t4_full = fourth_full(&t, &u, &w);
3462 for a in 0..K {
3463 for b in 0..K {
3464 assert_eq!(
3465 t3_sym[a][b].to_bits(),
3466 t3_full[a][b].to_bits(),
3467 "third K={K} [{a}][{b}]"
3468 );
3469 assert_eq!(
3470 t4_sym[a][b].to_bits(),
3471 t4_full[a][b].to_bits(),
3472 "fourth K={K} [{a}][{b}]"
3473 );
3474 checks += 2;
3475 }
3476 }
3477 }
3478 checks
3479 }
3480
3481 /// The output-symmetric contraction is BIT-IDENTICAL to the full nest across
3482 /// `K ∈ {2,3,4,9}` (so no fingerprint re-baseline is owed — accuracy and bits
3483 /// are unchanged; this is a pure speed-only optimization).
3484 #[test]
3485 fn contraction_symmetry_is_bit_identical_to_full_nest() {
3486 let checks = check_bit_identical::<2>(0x0000_0002_C0FF_EE01, 1000)
3487 + check_bit_identical::<3>(0x0000_0003_C0FF_EE01, 800)
3488 + check_bit_identical::<4>(0x0000_0004_C0FF_EE01, 600)
3489 + check_bit_identical::<9>(0x0000_0009_C0FF_EE01, 300);
3490 // Guards against the loops silently not running (e.g. a zeroed count):
3491 // 1000·2²·2 + 800·3²·2 + 600·4²·2 + 300·9²·2.
3492 assert_eq!(checks, 8000 + 14400 + 19200 + 48600);
3493 }
3494
3495 /// Measure the wall-clock of the output-symmetric contraction vs the full
3496 /// nest at `K = 9` (it does ~2× fewer inner contractions; the bit-identity
3497 /// test is the correctness gate). Informational — wall-clock is noisy — with
3498 /// only a PATHOLOGICAL-regression guard (the symmetric form does strictly
3499 /// fewer inner contractions, so it must not be materially slower).
3500 #[test]
3501 fn contraction_symmetry_speedup_is_reported() {
3502 const K: usize = 9;
3503 let mut r = Rng(0xC0FF_EE99_1234_5678);
3504 let towers: Vec<Tower4<K>> = (0..512).map(|_| rand_sym4::<K>(&mut r)).collect();
3505 let dir: [f64; K] = std::array::from_fn(|_| r.s());
3506 let u: [f64; K] = std::array::from_fn(|_| r.s());
3507 let w: [f64; K] = std::array::from_fn(|_| r.s());
3508
3509 let reps = 400usize;
3510 let t_sym = {
3511 let start = std::time::Instant::now();
3512 let mut sink = 0.0f64;
3513 for _ in 0..reps {
3514 for t in &towers {
3515 let o3 = std::hint::black_box(t).third_contracted(std::hint::black_box(&dir));
3516 let o4 = std::hint::black_box(t)
3517 .fourth_contracted(std::hint::black_box(&u), std::hint::black_box(&w));
3518 sink += o3[0][K - 1] + o4[0][K - 1];
3519 }
3520 }
3521 std::hint::black_box(sink);
3522 start.elapsed().as_secs_f64()
3523 };
3524 let t_full = {
3525 let start = std::time::Instant::now();
3526 let mut sink = 0.0f64;
3527 for _ in 0..reps {
3528 for t in &towers {
3529 let o3 = third_full(std::hint::black_box(t), std::hint::black_box(&dir));
3530 let o4 = fourth_full(
3531 std::hint::black_box(t),
3532 std::hint::black_box(&u),
3533 std::hint::black_box(&w),
3534 );
3535 sink += o3[0][K - 1] + o4[0][K - 1];
3536 }
3537 }
3538 std::hint::black_box(sink);
3539 start.elapsed().as_secs_f64()
3540 };
3541 let calls = (reps * towers.len()) as f64;
3542 eprintln!(
3543 "[contraction-symmetry speedup K=9] sym={:.1}ns/call full={:.1}ns/call \
3544 wall_speedup={:.2}x",
3545 t_sym / calls * 1e9,
3546 t_full / calls * 1e9,
3547 t_full / t_sym
3548 );
3549 // The report above is what this test's name promises and it always
3550 // runs. The wall-clock ASSERTION below is release-only, for two
3551 // independent reasons (#932).
3552 //
3553 // 1. It is an unpaired, sequential A/B. `t_sym` and `t_full` are each
3554 // timed exactly once, in fixed order, so any load change between the
3555 // two blocks lands entirely in the ratio. The house pattern for this
3556 // comparison elsewhere in the workspace is paired medians with
3557 // ALTERNATING order -- `paired_medians` in `fast_channel.rs`, and the
3558 // `(round + side) % 2` interleave in `multinomial_reml.rs` -- which
3559 // exists precisely to cancel drift this shape cannot see.
3560 //
3561 // 2. Measured on two nodes on the same tree: PASSED on n11 where the
3562 // whole 256-test suite took 5.1s, FAILED on cn1037 at
3563 // `sym=16.1469s full=9.9498s` (1.62x against the 1.5x bar) where the
3564 // same suite took 219.1s -- a 43x slowdown, with a sibling timing
3565 // test alone accounting for 219.0s of it. Under that much drift the
3566 // assertion cannot distinguish "pathologically slower" from "the node
3567 // was busy", which is the only thing it is supposed to detect.
3568 //
3569 // Debug therefore reports and stops. Nothing above this point is
3570 // skipped: the towers are still built and both contraction paths are
3571 // still executed and consumed, so any panic, NaN or shape error in
3572 // `third_contracted`/`fourth_contracted` is still caught in every build.
3573 if cfg!(debug_assertions) {
3574 return;
3575 }
3576 assert!(
3577 t_sym <= t_full * 1.5,
3578 "output-symmetric contraction pathologically slower: \
3579 sym={t_sym:.4}s full={t_full:.4}s"
3580 );
3581 }
3582}