use crate::jet_algebra;
#[derive(Clone, Copy, Debug)]
pub struct Tower4<const K: usize> {
pub v: f64,
pub g: [f64; K],
pub h: [[f64; K]; K],
pub t3: [[[f64; K]; K]; K],
pub t4: [[[[f64; K]; K]; K]; K],
}
impl<const K: usize> Tower4<K> {
pub fn zero() -> Self {
Self {
v: 0.0,
g: [0.0; K],
h: [[0.0; K]; K],
t3: [[[0.0; K]; K]; K],
t4: [[[[0.0; K]; K]; K]; K],
}
}
pub fn constant(c: f64) -> Self {
let mut out = Self::zero();
out.v = c;
out
}
pub fn variable(value: f64, idx: usize) -> Self {
let mut out = Self::constant(value);
out.g[idx] = 1.0;
out
}
#[inline]
fn deriv(&self, labels: &[usize]) -> f64 {
assert!(
labels.len() <= 4,
"Tower4 carries at most fourth-order derivatives"
);
match labels.len() {
0 => self.v,
1 => self.g[labels[0]],
2 => self.h[labels[0]][labels[1]],
3 => self.t3[labels[0]][labels[1]][labels[2]],
_ => self.t4[labels[0]][labels[1]][labels[2]][labels[3]],
}
}
pub fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let mut out = Self::zero();
out.v = a.v * b.v;
for i in 0..K {
let mut acc = 0.0;
acc += a.v * b.g[i];
acc += a.g[i] * b.v;
out.g[i] = acc;
}
for i in 0..K {
for j in i..K {
let mut acc = 0.0;
acc += a.v * b.h[i][j];
acc += a.g[i] * b.g[j];
acc += a.g[j] * b.g[i];
acc += a.h[i][j] * b.v;
out.h[i][j] = acc;
out.h[j][i] = acc;
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let mut acc = 0.0;
acc += a.v * b.t3[i][j][k];
acc += a.g[i] * b.h[j][k];
acc += a.g[j] * b.h[i][k];
acc += a.h[i][j] * b.g[k];
acc += a.g[k] * b.h[i][j];
acc += a.h[i][k] * b.g[j];
acc += a.h[j][k] * b.g[i];
acc += a.t3[i][j][k] * b.v;
out.t3[i][j][k] = acc;
}
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
for l in 0..K {
let mut acc = 0.0;
acc += a.v * b.t4[i][j][k][l];
acc += a.g[i] * b.t3[j][k][l];
acc += a.g[j] * b.t3[i][k][l];
acc += a.h[i][j] * b.h[k][l];
acc += a.g[k] * b.t3[i][j][l];
acc += a.h[i][k] * b.h[j][l];
acc += a.h[j][k] * b.h[i][l];
acc += a.t3[i][j][k] * b.g[l];
acc += a.g[l] * b.t3[i][j][k];
acc += a.h[i][l] * b.h[j][k];
acc += a.h[j][l] * b.h[i][k];
acc += a.t3[i][j][l] * b.g[k];
acc += a.h[k][l] * b.h[i][j];
acc += a.t3[i][k][l] * b.g[j];
acc += a.t3[j][k][l] * b.g[i];
acc += a.t4[i][j][k][l] * b.v;
out.t4[i][j][k][l] = acc;
}
}
}
}
out
}
pub fn add(&self, o: &Self) -> Self {
*self + *o
}
pub fn sub(&self, o: &Self) -> Self {
*self + o.scale(-1.0)
}
pub fn compose_unary(&self, d: [f64; 5]) -> Self {
let mut out = Self::zero();
out.v = d[0];
for i in 0..K {
let mut acc = 0.0;
acc += d[1] * self.g[i];
out.g[i] = acc;
}
for i in 0..K {
for j in 0..K {
let mut acc = 0.0;
acc += d[1] * self.h[i][j];
acc += d[2] * self.g[i] * self.g[j];
out.h[i][j] = acc;
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let mut acc = 0.0;
acc += d[1] * self.t3[i][j][k];
acc += d[2] * self.h[i][j] * self.g[k];
acc += d[2] * self.h[i][k] * self.g[j];
acc += d[2] * self.g[i] * self.h[j][k];
acc += d[3] * self.g[i] * self.g[j] * self.g[k];
out.t3[i][j][k] = acc;
}
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
for l in 0..K {
let mut acc = 0.0;
acc += d[1] * self.t4[i][j][k][l];
acc += d[2] * self.t3[i][j][k] * self.g[l];
acc += d[2] * self.t3[i][j][l] * self.g[k];
acc += d[2] * self.h[i][j] * self.h[k][l];
acc += d[3] * self.h[i][j] * self.g[k] * self.g[l];
acc += d[2] * self.t3[i][k][l] * self.g[j];
acc += d[2] * self.h[i][k] * self.h[j][l];
acc += d[3] * self.h[i][k] * self.g[j] * self.g[l];
acc += d[2] * self.h[i][l] * self.h[j][k];
acc += d[2] * self.g[i] * self.t3[j][k][l];
acc += d[3] * self.g[i] * self.h[j][k] * self.g[l];
acc += d[3] * self.h[i][l] * self.g[j] * self.g[k];
acc += d[3] * self.g[i] * self.h[j][l] * self.g[k];
acc += d[3] * self.g[i] * self.g[j] * self.h[k][l];
acc += d[4] * self.g[i] * self.g[j] * self.g[k] * self.g[l];
out.t4[i][j][k][l] = acc;
}
}
}
}
out
}
pub fn scale(&self, s: f64) -> Self {
let mut out = *self;
out.v *= s;
for i in 0..K {
out.g[i] *= s;
for j in 0..K {
out.h[i][j] *= s;
for k in 0..K {
out.t3[i][j][k] *= s;
for l in 0..K {
out.t4[i][j][k][l] *= s;
}
}
}
}
out
}
pub fn exp(&self) -> Self {
let e = self.v.exp();
self.compose_unary([e, e, e, e, e])
}
pub fn ln(&self) -> Self {
let u = self.v;
let r = 1.0 / u;
self.compose_unary([u.ln(), r, -r * r, 2.0 * r * r * r, -6.0 * r * r * r * r])
}
pub fn recip(&self) -> Self {
let r = 1.0 / self.v;
let r2 = r * r;
self.compose_unary([r, -r2, 2.0 * r2 * r, -6.0 * r2 * r2, 24.0 * r2 * r2 * r])
}
pub fn sqrt(&self) -> Self {
let u = self.v;
let s = u.sqrt();
self.compose_unary([
s,
0.5 / s,
-0.25 / (u * s),
0.375 / (u * u * s),
-0.9375 / (u * u * u * s),
])
}
pub fn powf(&self, a: f64) -> Self {
let u = self.v;
let f0 = u.powf(a);
let f1 = a * u.powf(a - 1.0);
let f2 = a * (a - 1.0) * u.powf(a - 2.0);
let f3 = a * (a - 1.0) * (a - 2.0) * u.powf(a - 3.0);
let f4 = a * (a - 1.0) * (a - 2.0) * (a - 3.0) * u.powf(a - 4.0);
self.compose_unary([f0, f1, f2, f3, f4])
}
pub fn ln_gamma(&self) -> Self {
self.compose_unary(ln_gamma_derivative_stack(self.v))
}
pub fn third_contracted(&self, dir: &[f64; K]) -> [[f64; K]; K] {
let mut out = [[0.0; K]; K];
for a in 0..K {
for b in a..K {
let mut acc = 0.0;
for c in 0..K {
acc += self.t3[a][b][c] * dir[c];
}
out[a][b] = acc;
out[b][a] = acc;
}
}
out
}
pub fn fourth_contracted(&self, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
let mut out = [[0.0; K]; K];
for i in 0..K {
for j in i..K {
let mut acc = 0.0;
for k in 0..K {
for l in 0..K {
acc += self.t4[i][j][k][l] * u[k] * w[l];
}
}
out[i][j] = acc;
out[j][i] = acc;
}
}
out
}
}
impl<const K: usize> jet_algebra::JetAlgebra<5> for Tower4<K> {
#[inline]
fn derivative(&self, labels: &[usize]) -> f64 {
self.deriv(labels)
}
fn map_derivatives<F>(&self, mut f: F) -> Self
where
F: FnMut(&[usize]) -> f64,
{
let mut out = Self::zero();
out.v = f(&[]);
for i in 0..K {
let labels = [i];
out.g[i] = f(&labels);
}
for i in 0..K {
for j in 0..K {
let labels = [i, j];
out.h[i][j] = f(&labels);
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let labels = [i, j, k];
out.t3[i][j][k] = f(&labels);
}
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
for l in 0..K {
let labels = [i, j, k, l];
out.t4[i][j][k][l] = f(&labels);
}
}
}
}
out
}
}
#[derive(Clone, Copy, Debug)]
pub struct Tower2<const K: usize> {
pub v: f64,
pub g: [f64; K],
pub h: [[f64; K]; K],
}
impl<const K: usize> Tower2<K> {
pub fn zero() -> Self {
Self {
v: 0.0,
g: [0.0; K],
h: [[0.0; K]; K],
}
}
pub fn constant(c: f64) -> Self {
let mut out = Self::zero();
out.v = c;
out
}
pub fn variable(value: f64, idx: usize) -> Self {
let mut out = Self::constant(value);
out.g[idx] = 1.0;
out
}
#[inline]
fn deriv(&self, labels: &[usize]) -> f64 {
assert!(
labels.len() <= 2,
"Tower2 carries at most second-order derivatives"
);
match labels.len() {
0 => self.v,
1 => self.g[labels[0]],
_ => self.h[labels[0]][labels[1]],
}
}
pub fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let mut out = Self::zero();
out.v = a.v * b.v;
for i in 0..K {
out.g[i] = a.v * b.g[i] + a.g[i] * b.v;
}
for i in 0..K {
for j in i..K {
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;
out.h[i][j] = hij;
out.h[j][i] = hij;
}
}
out
}
pub fn compose_unary(&self, d: [f64; 3]) -> Self {
let mut out = Self::zero();
out.v = d[0];
for i in 0..K {
let mut acc = 0.0;
acc += d[1] * self.g[i];
out.g[i] = acc;
}
for i in 0..K {
for j in 0..K {
let mut acc = 0.0;
acc += d[1] * self.h[i][j];
acc += d[2] * self.g[i] * self.g[j];
out.h[i][j] = acc;
}
}
out
}
pub fn scale(&self, s: f64) -> Self {
let mut out = *self;
out.v *= s;
for i in 0..K {
out.g[i] *= s;
for j in 0..K {
out.h[i][j] *= s;
}
}
out
}
pub fn exp(&self) -> Self {
let e = self.v.exp();
self.compose_unary([e, e, e])
}
pub fn sqrt(&self) -> Self {
let u = self.v;
let s = u.sqrt();
self.compose_unary([s, 0.5 / s, -0.25 / (u * s)])
}
}
impl<const K: usize> jet_algebra::JetAlgebra<3> for Tower2<K> {
#[inline]
fn derivative(&self, labels: &[usize]) -> f64 {
self.deriv(labels)
}
fn map_derivatives<F>(&self, mut f: F) -> Self
where
F: FnMut(&[usize]) -> f64,
{
let mut out = Self::zero();
out.v = f(&[]);
for i in 0..K {
let labels = [i];
out.g[i] = f(&labels);
}
for i in 0..K {
for j in 0..K {
let labels = [i, j];
out.h[i][j] = f(&labels);
}
}
out
}
}
impl<const K: usize> std::ops::Add for Tower2<K> {
type Output = Self;
fn add(self, o: Self) -> Self {
let mut out = self;
out.v += o.v;
for i in 0..K {
out.g[i] += o.g[i];
for j in 0..K {
out.h[i][j] += o.h[i][j];
}
}
out
}
}
impl<const K: usize> std::ops::Mul for Tower2<K> {
type Output = Self;
fn mul(self, o: Self) -> Self {
Tower2::mul(&self, &o)
}
}
impl<const K: usize> std::ops::Add<f64> for Tower2<K> {
type Output = Self;
fn add(self, c: f64) -> Self {
let mut out = self;
out.v += c;
out
}
}
impl<const K: usize> std::ops::Mul<f64> for Tower2<K> {
type Output = Self;
fn mul(self, c: f64) -> Self {
self.scale(c)
}
}
#[derive(Clone, Copy, Debug)]
pub struct Tower3<const K: usize> {
pub v: f64,
pub g: [f64; K],
pub h: [[f64; K]; K],
pub t3: [[[f64; K]; K]; K],
}
impl<const K: usize> Tower3<K> {
pub fn zero() -> Self {
Self {
v: 0.0,
g: [0.0; K],
h: [[0.0; K]; K],
t3: [[[0.0; K]; K]; K],
}
}
pub fn constant(c: f64) -> Self {
let mut out = Self::zero();
out.v = c;
out
}
pub fn variable(value: f64, idx: usize) -> Self {
let mut out = Self::constant(value);
out.g[idx] = 1.0;
out
}
#[inline]
fn deriv(&self, labels: &[usize]) -> f64 {
assert!(
labels.len() <= 3,
"Tower3 carries at most third-order derivatives"
);
match labels.len() {
0 => self.v,
1 => self.g[labels[0]],
2 => self.h[labels[0]][labels[1]],
_ => self.t3[labels[0]][labels[1]][labels[2]],
}
}
pub fn mul(&self, o: &Self) -> Self {
let a = self;
let b = o;
let mut out = Self::zero();
out.v = a.v * b.v;
for i in 0..K {
let mut acc = 0.0;
acc += a.v * b.g[i];
acc += a.g[i] * b.v;
out.g[i] = acc;
}
for i in 0..K {
for j in i..K {
let mut acc = 0.0;
acc += a.v * b.h[i][j];
acc += a.g[i] * b.g[j];
acc += a.g[j] * b.g[i];
acc += a.h[i][j] * b.v;
out.h[i][j] = acc;
out.h[j][i] = acc;
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let mut acc = 0.0;
acc += a.v * b.t3[i][j][k];
acc += a.g[i] * b.h[j][k];
acc += a.g[j] * b.h[i][k];
acc += a.h[i][j] * b.g[k];
acc += a.g[k] * b.h[i][j];
acc += a.h[i][k] * b.g[j];
acc += a.h[j][k] * b.g[i];
acc += a.t3[i][j][k] * b.v;
out.t3[i][j][k] = acc;
}
}
}
out
}
pub fn add(&self, o: &Self) -> Self {
*self + *o
}
pub fn sub(&self, o: &Self) -> Self {
*self + o.scale(-1.0)
}
pub fn compose_unary(&self, d: [f64; 4]) -> Self {
let mut out = Self::zero();
out.v = d[0];
for i in 0..K {
let mut acc = 0.0;
acc += d[1] * self.g[i];
out.g[i] = acc;
}
for i in 0..K {
for j in 0..K {
let mut acc = 0.0;
acc += d[1] * self.h[i][j];
acc += d[2] * self.g[i] * self.g[j];
out.h[i][j] = acc;
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let mut acc = 0.0;
acc += d[1] * self.t3[i][j][k];
acc += d[2] * self.h[i][j] * self.g[k];
acc += d[2] * self.h[i][k] * self.g[j];
acc += d[2] * self.g[i] * self.h[j][k];
acc += d[3] * self.g[i] * self.g[j] * self.g[k];
out.t3[i][j][k] = acc;
}
}
}
out
}
#[inline]
pub fn compose_unary_with(&self, stack_fn: impl Fn(f64) -> [f64; 4]) -> Self {
self.compose_unary(stack_fn(self.v))
}
pub fn scale(&self, s: f64) -> Self {
let mut out = *self;
out.v *= s;
for i in 0..K {
out.g[i] *= s;
for j in 0..K {
out.h[i][j] *= s;
for k in 0..K {
out.t3[i][j][k] *= s;
}
}
}
out
}
}
impl<const K: usize> jet_algebra::JetAlgebra<4> for Tower3<K> {
#[inline]
fn derivative(&self, labels: &[usize]) -> f64 {
self.deriv(labels)
}
fn map_derivatives<F>(&self, mut f: F) -> Self
where
F: FnMut(&[usize]) -> f64,
{
let mut out = Self::zero();
out.v = f(&[]);
for i in 0..K {
let labels = [i];
out.g[i] = f(&labels);
}
for i in 0..K {
for j in 0..K {
let labels = [i, j];
out.h[i][j] = f(&labels);
}
}
for i in 0..K {
for j in 0..K {
for k in 0..K {
let labels = [i, j, k];
out.t3[i][j][k] = f(&labels);
}
}
}
out
}
}
impl<const K: usize> std::ops::Add for Tower3<K> {
type Output = Self;
fn add(self, o: Self) -> Self {
let mut out = self;
out.v += o.v;
for i in 0..K {
out.g[i] += o.g[i];
for j in 0..K {
out.h[i][j] += o.h[i][j];
for k in 0..K {
out.t3[i][j][k] += o.t3[i][j][k];
}
}
}
out
}
}
pub fn ln_gamma_derivative_stack(x: f64) -> [f64; 5] {
[
statrs::function::gamma::ln_gamma(x),
digamma_positive(x),
polygamma_positive(1, x),
polygamma_positive(2, x),
polygamma_positive(3, x),
]
}
pub fn ln_gamma_derivative_stack_order2(x: f64) -> [f64; 3] {
[
statrs::function::gamma::ln_gamma(x),
digamma_positive(x),
polygamma_positive(1, x),
]
}
pub fn digamma_derivative_stack(x: f64) -> [f64; 5] {
[
digamma_positive(x),
polygamma_positive(1, x),
polygamma_positive(2, x),
polygamma_positive(3, x),
polygamma_positive(4, x),
]
}
#[inline]
pub fn digamma(x: f64) -> f64 {
digamma_positive(x)
}
#[inline]
pub fn trigamma(x: f64) -> f64 {
polygamma_positive(1, x)
}
fn digamma_positive(mut x: f64) -> f64 {
if !(x.is_finite() && x > 0.0) {
return f64::NAN;
}
let mut acc = 0.0;
while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
acc -= 1.0 / x;
x += 1.0;
}
acc + digamma_asymptotic(x)
}
fn polygamma_positive(order: usize, mut x: f64) -> f64 {
if !(x.is_finite() && x > 0.0) {
return f64::NAN;
}
let mut acc = 0.0;
while x < POLYGAMMA_ASYMPTOTIC_MIN_X {
acc += polygamma_recurrence_term(order, x);
x += 1.0;
}
acc + polygamma_asymptotic(order, x)
}
const POLYGAMMA_ASYMPTOTIC_MIN_X: f64 = 20.0;
const BERNOULLI_EVEN: [(usize, f64); 10] = [
(2, 1.0 / 6.0),
(4, -1.0 / 30.0),
(6, 1.0 / 42.0),
(8, -1.0 / 30.0),
(10, 5.0 / 66.0),
(12, -691.0 / 2730.0),
(14, 7.0 / 6.0),
(16, -3617.0 / 510.0),
(18, 43867.0 / 798.0),
(20, -174611.0 / 330.0),
];
fn polygamma_recurrence_term(order: usize, x: f64) -> f64 {
let sign = if order % 2 == 1 { 1.0 } else { -1.0 };
sign * factorial(order) / x.powi((order + 1) as i32)
}
fn digamma_asymptotic(x: f64) -> f64 {
let mut out = x.ln() - 0.5 / x;
for (bernoulli_order, bernoulli) in BERNOULLI_EVEN {
out -= bernoulli / (bernoulli_order as f64 * x.powi(bernoulli_order as i32));
}
out
}
fn polygamma_asymptotic(order: usize, x: f64) -> f64 {
if !(1..=5).contains(&order) {
return f64::NAN;
}
let order_factorial = factorial(order);
let leading_sign = if order % 2 == 1 { 1.0 } else { -1.0 };
let mut out = leading_sign * factorial(order - 1) / x.powi(order as i32)
+ leading_sign * order_factorial / (2.0 * x.powi((order + 1) as i32));
let bernoulli_sign = if order % 2 == 1 { 1.0 } else { -1.0 };
for (bernoulli_order, bernoulli) in BERNOULLI_EVEN {
let rising = rising_factorial(bernoulli_order, order);
out += bernoulli_sign * bernoulli * rising
/ bernoulli_order as f64
/ x.powi((bernoulli_order + order) as i32);
}
out
}
fn factorial(n: usize) -> f64 {
(1..=n).fold(1.0, |acc, k| acc * k as f64)
}
fn rising_factorial(start: usize, len: usize) -> f64 {
(start..start + len).fold(1.0, |acc, k| acc * k as f64)
}
impl<const K: usize> std::ops::Add for Tower4<K> {
type Output = Self;
fn add(self, o: Self) -> Self {
let mut out = self;
out.v += o.v;
for i in 0..K {
out.g[i] += o.g[i];
for j in 0..K {
out.h[i][j] += o.h[i][j];
for k in 0..K {
out.t3[i][j][k] += o.t3[i][j][k];
for l in 0..K {
out.t4[i][j][k][l] += o.t4[i][j][k][l];
}
}
}
}
out
}
}
impl<const K: usize> std::ops::Sub for Tower4<K> {
type Output = Self;
fn sub(self, o: Self) -> Self {
self + o.scale(-1.0)
}
}
impl<const K: usize> std::ops::Neg for Tower4<K> {
type Output = Self;
fn neg(self) -> Self {
self.scale(-1.0)
}
}
impl<const K: usize> std::ops::Mul for Tower4<K> {
type Output = Self;
fn mul(self, o: Self) -> Self {
Tower4::mul(&self, &o)
}
}
impl<const K: usize> std::ops::Div for Tower4<K> {
type Output = Self;
fn div(self, o: Self) -> Self {
Tower4::mul(&self, &o.recip())
}
}
impl<const K: usize> std::ops::Add<f64> for Tower4<K> {
type Output = Self;
fn add(self, c: f64) -> Self {
let mut out = self;
out.v += c;
out
}
}
impl<const K: usize> std::ops::Sub<f64> for Tower4<K> {
type Output = Self;
fn sub(self, c: f64) -> Self {
self + (-c)
}
}
impl<const K: usize> std::ops::Mul<f64> for Tower4<K> {
type Output = Self;
fn mul(self, c: f64) -> Self {
self.scale(c)
}
}
pub trait RowProgram<const K: usize>: Send + Sync {
fn n_rows(&self) -> usize;
fn primaries(&self, row: usize) -> Result<[f64; K], String>;
fn eval<S: crate::jet_scalar::JetScalar<K>>(&self, row: usize, p: &[S; K])
-> Result<S, String>;
}
const PROGRAM_DENSE_JET_STACK_BUDGET_BYTES: usize = 64 * 1024;
#[inline]
fn program_primary_jets_fit_stack<S, const K: usize>() -> bool {
std::mem::size_of::<S>()
.checked_mul(K)
.is_some_and(|bytes| bytes <= PROGRAM_DENSE_JET_STACK_BUDGET_BYTES)
}
fn evaluate_program_with_stack_primaries<const K: usize, P, S>(
prog: &P,
row: usize,
mut seed: impl FnMut(usize) -> S,
) -> Result<S, String>
where
P: RowProgram<K> + ?Sized,
S: crate::jet_scalar::JetScalar<K>,
{
let vars: [S; K] = std::array::from_fn(&mut seed);
prog.eval(row, &vars)
}
#[inline(never)]
fn evaluate_program_with_heap_primaries<const K: usize, P, S>(
prog: &P,
row: usize,
seed: impl FnMut(usize) -> S,
) -> Result<S, String>
where
P: RowProgram<K> + ?Sized,
S: crate::jet_scalar::JetScalar<K>,
{
let vars: Box<[S]> = (0..K).map(seed).collect();
let vars: Box<[S; K]> = vars.try_into().map_err(|vars: Box<[S]>| {
format!(
"canonical row program seeded {} primary jets; expected exactly {K}",
vars.len()
)
})?;
prog.eval(row, &vars)
}
#[inline]
fn evaluate_program_with_seeded_primaries<const K: usize, P, S>(
prog: &P,
row: usize,
seed: impl FnMut(usize) -> S,
) -> Result<S, String>
where
P: RowProgram<K> + ?Sized,
S: crate::jet_scalar::JetScalar<K>,
{
if program_primary_jets_fit_stack::<S, K>() {
evaluate_program_with_stack_primaries(prog, row, seed)
} else {
evaluate_program_with_heap_primaries(prog, row, seed)
}
}
pub fn program_row_kernel<const K: usize, P: RowProgram<K> + ?Sized>(
prog: &P,
row: usize,
) -> Result<(f64, [f64; K], [[f64; K]; K]), String> {
let base = prog.primaries(row)?;
let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
<crate::jet_scalar::Order2<K> as crate::jet_scalar::JetScalar<K>>::variable(base[a], a)
})?;
Ok(s.into_channels())
}
pub fn program_third_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
prog: &P,
row: usize,
dir: &[f64; K],
) -> Result<[[f64; K]; K], String> {
let base = prog.primaries(row)?;
let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
crate::jet_scalar::OneSeed::seed_direction(base[a], a, dir[a])
})?;
Ok(s.contracted_third())
}
pub fn program_fourth_contracted<const K: usize, P: RowProgram<K> + ?Sized>(
prog: &P,
row: usize,
dir_u: &[f64; K],
dir_v: &[f64; K],
) -> Result<[[f64; K]; K], String> {
let base = prog.primaries(row)?;
let s = evaluate_program_with_seeded_primaries(prog, row, |a| {
crate::jet_scalar::TwoSeed::seed(base[a], a, dir_u[a], dir_v[a])
})?;
Ok(s.contracted_fourth())
}
pub fn program_full_tower<const K: usize, P: RowProgram<K> + ?Sized>(
prog: &P,
row: usize,
) -> Result<Box<Tower4<K>>, String> {
let tower_bytes = std::mem::size_of::<Tower4<K>>();
if tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES {
return Err(format!(
"canonical dense Tower4<{K}> requires {tower_bytes} bytes, exceeding the {}-byte \
storage budget; use the bounded row-kernel and directional channel APIs",
PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
));
}
let base = prog.primaries(row)?;
evaluate_program_with_seeded_primaries(prog, row, |a| Tower4::variable(base[a], a))
.map(Box::new)
}
pub struct KernelChannels<const K: usize> {
pub value: f64,
pub gradient: [f64; K],
pub hessian: [[f64; K]; K],
pub third: Vec<([f64; K], [[f64; K]; K])>,
pub fourth: Vec<([f64; K], [f64; K], [[f64; K]; K])>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tower3_matches_tower4_through_third_order() {
let s_a: [f64; 5] = [
0.3_f64.sin(),
0.3_f64.cos(),
-0.3_f64.sin(),
-0.3_f64.cos(),
0.3_f64.sin(),
];
let s_b: [f64; 5] = [1.1, -0.4, 0.8, -0.2, 0.05];
let s4 = |s: [f64; 5]| [s[0], s[1], s[2], s[3]];
let a4 = Tower4::<3>::variable(0.4, 0);
let b4 = Tower4::<3>::variable(-0.7, 1);
let c4 = Tower4::<3>::variable(0.9, 2);
let prog4 = (a4.mul(&b4) + c4).compose_unary(s_a).scale(1.3)
+ a4.mul(&c4).scale(-0.7)
+ b4.compose_unary(s_b).scale(0.25);
let a3 = Tower3::<3>::variable(0.4, 0);
let b3 = Tower3::<3>::variable(-0.7, 1);
let c3 = Tower3::<3>::variable(0.9, 2);
let prog3 = (a3.mul(&b3) + c3).compose_unary(s4(s_a)).scale(1.3)
+ a3.mul(&c3).scale(-0.7)
+ b3.compose_unary(s4(s_b)).scale(0.25);
assert_eq!(prog3.v.to_bits(), prog4.v.to_bits(), "value mismatch");
for i in 0..3 {
assert_eq!(
prog3.g[i].to_bits(),
prog4.g[i].to_bits(),
"g[{i}] mismatch"
);
for j in 0..3 {
assert_eq!(
prog3.h[i][j].to_bits(),
prog4.h[i][j].to_bits(),
"h[{i}][{j}] mismatch"
);
for k in 0..3 {
assert_eq!(
prog3.t3[i][j][k].to_bits(),
prog4.t3[i][j][k].to_bits(),
"t3[{i}][{j}][{k}] mismatch"
);
}
}
}
}
struct LogitProgram {
eta: Vec<f64>,
y: Vec<f64>,
}
impl RowProgram<1> for LogitProgram {
fn n_rows(&self) -> usize {
self.eta.len()
}
fn primaries(&self, row: usize) -> Result<[f64; 1], String> {
Ok([self.eta[row]])
}
fn eval<S: crate::jet_scalar::JetScalar<1>>(
&self,
row: usize,
p: &[S; 1],
) -> Result<S, String> {
let eta = p[0];
Ok(eta
.exp()
.add(&S::constant(1.0))
.ln()
.sub(&eta.scale(self.y[row])))
}
}
#[test]
fn logit_tower_matches_closed_forms() {
let prog = LogitProgram {
eta: vec![-2.3, -0.4, 0.0, 0.9, 3.1],
y: vec![1.0, 0.0, 1.0, 0.0, 1.0],
};
for row in 0..prog.n_rows() {
let t = program_full_tower(&prog, row).expect("logit program");
let eta = prog.eta[row];
let y = prog.y[row];
let mu = 1.0 / (1.0 + (-eta).exp());
let w = mu * (1.0 - mu);
let expect = [
(t.v, (1.0 + eta.exp()).ln() - y * eta, "value"),
(t.g[0], mu - y, "grad"),
(t.h[0][0], w, "hess"),
(t.t3[0][0][0], w * (1.0 - 2.0 * mu), "third"),
(
t.t4[0][0][0][0],
w * (1.0 - 6.0 * mu + 6.0 * mu * mu),
"fourth",
),
];
for (got, want, label) in expect {
assert!(
(got - want).abs() <= 1e-12 * want.abs().max(1.0),
"row {row} {label}: got {got:+.15e} want {want:+.15e}"
);
}
}
}
struct OversizedDenseProgram;
impl RowProgram<10> for OversizedDenseProgram {
fn n_rows(&self) -> usize {
1
}
fn primaries(&self, row: usize) -> Result<[f64; 10], String> {
Err(format!(
"dense-tower storage check reached program primaries at row {row}"
))
}
fn eval<S: crate::jet_scalar::JetScalar<10>>(
&self,
row: usize,
primaries: &[S; 10],
) -> Result<S, String> {
Err(format!(
"dense-tower storage check reached program evaluation at row {row} with {} primaries",
primaries.len()
))
}
}
struct LargestBudgetedDenseProgram;
impl RowProgram<9> for LargestBudgetedDenseProgram {
fn n_rows(&self) -> usize {
1
}
fn primaries(&self, row: usize) -> Result<[f64; 9], String> {
if row == 0 {
Ok([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
} else {
Err(format!("largest budgeted dense program has no row {row}"))
}
}
fn eval<S: crate::jet_scalar::JetScalar<9>>(
&self,
row: usize,
primaries: &[S; 9],
) -> Result<S, String> {
if row != 0 {
return Err(format!("largest budgeted dense program has no row {row}"));
}
let linear =
S::linear_combination(primaries, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
let quartic = primaries[0]
.mul(&primaries[1])
.mul(&primaries[2])
.mul(&primaries[3]);
Ok(linear.add(&quartic))
}
}
#[test]
fn full_tower_accepts_largest_width_inside_storage_budget_932() {
assert_eq!(std::mem::size_of::<Tower4<9>>(), 59_048);
assert!(
!program_primary_jets_fit_stack::<Tower4<9>, 9>(),
"nine full-width primary towers must use exact-length heap storage"
);
let tower = program_full_tower(&LargestBudgetedDenseProgram, 0)
.expect("Tower4<9> must remain inside the canonical dense storage budget");
assert_eq!(tower.v, 309.0);
assert_eq!(tower.g, [25.0, 14.0, 11.0, 10.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
assert_eq!(tower.t4[0][1][2][3], 1.0);
}
#[test]
fn full_tower_refuses_oversized_result_before_touching_program() {
let tower_bytes = std::mem::size_of::<Tower4<10>>();
assert!(tower_bytes > PROGRAM_DENSE_JET_STACK_BUDGET_BYTES);
assert!(
std::mem::size_of::<Result<Box<Tower4<32>>, String>>()
<= 4 * std::mem::size_of::<usize>(),
"boxed full-tower API must keep its return slot independent of dense tower width"
);
let error = program_full_tower(&OversizedDenseProgram, 0)
.expect_err("Tower4<10> must exceed the canonical dense storage budget");
assert_eq!(
error,
format!(
"canonical dense Tower4<10> requires {tower_bytes} bytes, exceeding the {}-byte \
storage budget; use the bounded row-kernel and directional channel APIs",
PROGRAM_DENSE_JET_STACK_BUDGET_BYTES
)
);
}
struct LocScaleProgram {
eta: Vec<f64>,
s: Vec<f64>,
y: Vec<f64>,
}
impl RowProgram<2> for LocScaleProgram {
fn n_rows(&self) -> usize {
self.eta.len()
}
fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
Ok([self.eta[row], self.s[row]])
}
fn eval<S: crate::jet_scalar::JetScalar<2>>(
&self,
row: usize,
p: &[S; 2],
) -> Result<S, String> {
let r = S::constant(self.y[row]).sub(&p[0]);
Ok(p[1].add(&p[1].scale(-2.0).exp().mul(&r).mul(&r).scale(0.5)))
}
}
#[test]
fn locscale_tower_matches_closed_forms_including_cross_blocks() {
let prog = LocScaleProgram {
eta: vec![0.3, -1.1, 2.0],
s: vec![-0.5, 0.2, 0.8],
y: vec![1.0, -2.0, 2.5],
};
let tol = 1e-12;
for row in 0..prog.n_rows() {
let t = program_full_tower(&prog, row).expect("locscale program");
let r = prog.y[row] - prog.eta[row];
let w = (-2.0 * prog.s[row]).exp();
let truth_g = [-w * r, 1.0 - w * r * r];
let truth_h = [[w, 2.0 * w * r], [2.0 * w * r, 2.0 * w * r * r]];
let t3_truth = |a: usize, b: usize, c: usize| -> f64 {
match a + b + c {
0 => 0.0,
1 => -2.0 * w,
2 => -4.0 * w * r,
_ => -4.0 * w * r * r,
}
};
let t4_truth = |a: usize, b: usize, c: usize, d: usize| -> f64 {
match a + b + c + d {
0 | 1 => 0.0,
2 => 4.0 * w,
3 => 8.0 * w * r,
_ => 8.0 * w * r * r,
}
};
for a in 0..2 {
assert!(
(t.g[a] - truth_g[a]).abs() <= tol * truth_g[a].abs().max(1.0),
"row {row} grad[{a}]"
);
for b in 0..2 {
assert!(
(t.h[a][b] - truth_h[a][b]).abs() <= tol * w.max(1.0) * (1.0 + r.abs()),
"row {row} hess[{a}][{b}]: got {} want {}",
t.h[a][b],
truth_h[a][b]
);
for c in 0..2 {
assert!(
(t.t3[a][b][c] - t3_truth(a, b, c)).abs()
<= tol * 8.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
"row {row} t3[{a}][{b}][{c}]: got {} want {}",
t.t3[a][b][c],
t3_truth(a, b, c)
);
for d in 0..2 {
assert!(
(t.t4[a][b][c][d] - t4_truth(a, b, c, d)).abs()
<= tol * 16.0 * w.max(1.0) * (1.0 + r.abs() + r * r),
"row {row} t4[{a}][{b}][{c}][{d}]: got {} want {}",
t.t4[a][b][c][d],
t4_truth(a, b, c, d)
);
}
}
}
}
let dir = [0.7, -1.3];
let third = program_third_contracted(&prog, row, &dir).expect("third");
for a in 0..2 {
for b in 0..2 {
let want = t.t3[a][b][0] * dir[0] + t.t3[a][b][1] * dir[1];
assert!((third[a][b] - want).abs() <= 1e-13 * want.abs().max(1.0));
}
}
}
}
struct GnarlyProgram {
primaries: Vec<[f64; 3]>,
tau: Vec<f64>,
}
impl GnarlyProgram {
fn fixture() -> Self {
Self {
primaries: vec![[0.4, -0.7, 1.2], [-0.9, 0.6, 0.3], [1.1, -0.2, -0.8]],
tau: vec![0.15, -0.35, 0.5],
}
}
}
impl RowProgram<3> for GnarlyProgram {
fn n_rows(&self) -> usize {
self.primaries.len()
}
fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
self.primaries
.get(row)
.copied()
.ok_or_else(|| format!("gnarly: row {row} out of range"))
}
fn eval<S: crate::jet_scalar::JetScalar<3>>(
&self,
row: usize,
p: &[S; 3],
) -> Result<S, String> {
let tau = *self
.tau
.get(row)
.ok_or_else(|| format!("gnarly: tau row {row} out of range"))?;
let a = p[0].mul(&p[1]).exp();
let b = p[2].mul(&p[2]).add(&S::constant(1.0)).sqrt();
let c = a.add(&b).add(&S::constant(tau)).ln();
let d = p[1].scale(0.5).add(&S::constant(2.0)).powf(1.7);
let delta = p[0].sub(&p[2]);
Ok(c.mul(&d.recip()).add(&delta.mul(&delta).scale(0.25)))
}
}
fn gnarly_tower_at(prog: &GnarlyProgram, row: usize, p: [f64; 3]) -> Tower4<3> {
struct At<'a> {
base: &'a GnarlyProgram,
row: usize,
p: [f64; 3],
}
impl RowProgram<3> for At<'_> {
fn n_rows(&self) -> usize {
1
}
fn primaries(&self, row: usize) -> Result<[f64; 3], String> {
if row != 0 {
return Err(format!("gnarly-at: row {row} out of range"));
}
Ok(self.p)
}
fn eval<S: crate::jet_scalar::JetScalar<3>>(
&self,
eval_row: usize,
vars: &[S; 3],
) -> Result<S, String> {
if eval_row != 0 {
return Err(format!("gnarly-at: eval row {eval_row} out of range"));
}
self.base.eval(self.row, vars)
}
}
*program_full_tower(&At { base: prog, row, p }, 0).expect("gnarly tower")
}
#[test]
fn gnarly_tower_is_fd_consistent_order_by_order() {
let prog = GnarlyProgram::fixture();
for row in 0..prog.n_rows() {
let base = prog.primaries(row).expect("primaries");
let t = gnarly_tower_at(&prog, row, base);
let h_step = 1e-5;
let tol = 1e-6;
for c in 0..3 {
let mut up = base;
let mut dn = base;
up[c] += h_step;
dn[c] -= h_step;
let t_up = gnarly_tower_at(&prog, row, up);
let t_dn = gnarly_tower_at(&prog, row, dn);
let fd_g = (t_up.v - t_dn.v) / (2.0 * h_step);
assert!(
(t.g[c] - fd_g).abs() <= tol * fd_g.abs().max(1.0),
"grad[{c}]: analytic {} fd {}",
t.g[c],
fd_g
);
for a in 0..3 {
let fd_h = (t_up.g[a] - t_dn.g[a]) / (2.0 * h_step);
assert!(
(t.h[a][c] - fd_h).abs() <= tol * fd_h.abs().max(1.0),
"hess[{a}][{c}]: analytic {} fd {}",
t.h[a][c],
fd_h
);
for b in 0..3 {
let fd_t3 = (t_up.h[a][b] - t_dn.h[a][b]) / (2.0 * h_step);
assert!(
(t.t3[a][b][c] - fd_t3).abs() <= tol * fd_t3.abs().max(1.0),
"t3[{a}][{b}][{c}]: analytic {} fd {}",
t.t3[a][b][c],
fd_t3
);
for d in 0..3 {
let fd_t4 = (t_up.t3[a][b][d] - t_dn.t3[a][b][d]) / (2.0 * h_step);
assert!(
(t.t4[a][b][d][c] - fd_t4).abs() <= tol * fd_t4.abs().max(1.0),
"t4[{a}][{b}][{d}][{c}]: analytic {} fd {}",
t.t4[a][b][d][c],
fd_t4
);
}
}
}
}
}
}
#[test]
fn crossing_edge_tower_matches_handpath_velocity_formulas() {
const TAU: f64 = 1.3; let g_idx = 1usize;
let g0 = 0.85_f64; let mut a = Tower4::<3>::constant(0.45);
a.g[0] = 0.7;
a.g[1] = -0.3;
a.h[0][0] = 0.25;
a.h[0][1] = 0.11;
a.h[1][0] = 0.11;
a.h[1][1] = -0.08;
let b = Tower4::<3>::variable(g0, g_idx);
let z_edge = (Tower4::<3>::constant(TAU) - a) / b;
let bv = g0;
let z0 = z_edge.v;
assert!((z0 - (TAU - 0.45) / bv).abs() < 1e-12);
for u in 0..2 {
let direct = if u == g_idx { z0 } else { 0.0 };
let want = -(a.g[u] + direct) / bv;
assert!(
(z_edge.g[u] - want).abs() < 1e-10,
"z_u[{u}] {:+.8e} vs hand formula {:+.8e}",
z_edge.g[u],
want
);
}
for u in 0..2 {
for v in 0..2 {
let cross = if u == g_idx { z_edge.g[v] } else { 0.0 }
+ if v == g_idx { z_edge.g[u] } else { 0.0 };
let want = -(a.h[u][v] + cross) / bv;
assert!(
(z_edge.h[u][v] - want).abs() < 1e-10,
"z_uv[{u}][{v}] {:+.8e} vs hand formula {:+.8e}",
z_edge.h[u][v],
want
);
}
}
}
#[test]
fn crossing_edge_constraint_frame_matches_bare_velocity_constants() {
const TAU: f64 = 1.3;
let a0 = 0.45_f64;
let b0 = 0.85_f64;
let a = Tower4::<2>::variable(a0, 0);
let b = Tower4::<2>::variable(b0, 1);
let z = (Tower4::<2>::constant(TAU) - a) / b;
assert!((z.v - (TAU - a0) / b0).abs() < 1e-12);
assert!((z.g[0] - (-1.0 / b0)).abs() < 1e-12, "z_a {:+.10e}", z.g[0]);
assert!(
(z.h[0][1] - 1.0 / (b0 * b0)).abs() < 1e-12,
"z_ab {:+.10e} vs +1/b² {:+.10e}",
z.h[0][1],
1.0 / (b0 * b0)
);
assert!(
z.h[0][0].abs() < 1e-12,
"z_aa must vanish, got {:+.10e}",
z.h[0][0]
);
let want_zbb = 2.0 * (TAU - a0) / (b0 * b0 * b0);
assert!(
(z.h[1][1] - want_zbb).abs() < 1e-12,
"z_bb {:+.10e} vs 2(τ−a)/b³ {:+.10e}",
z.h[1][1],
want_zbb
);
}
#[test]
fn t3_t4_are_fully_index_symmetric() {
let prog = GnarlyProgram::fixture();
let perms3: [[usize; 3]; 6] = [
[0, 1, 2],
[0, 2, 1],
[1, 0, 2],
[1, 2, 0],
[2, 0, 1],
[2, 1, 0],
];
let perms4: [[usize; 4]; 24] = [
[0, 1, 2, 3],
[0, 1, 3, 2],
[0, 2, 1, 3],
[0, 2, 3, 1],
[0, 3, 1, 2],
[0, 3, 2, 1],
[1, 0, 2, 3],
[1, 0, 3, 2],
[1, 2, 0, 3],
[1, 2, 3, 0],
[1, 3, 0, 2],
[1, 3, 2, 0],
[2, 0, 1, 3],
[2, 0, 3, 1],
[2, 1, 0, 3],
[2, 1, 3, 0],
[2, 3, 0, 1],
[2, 3, 1, 0],
[3, 0, 1, 2],
[3, 0, 2, 1],
[3, 1, 0, 2],
[3, 1, 2, 0],
[3, 2, 0, 1],
[3, 2, 1, 0],
];
for row in 0..prog.n_rows() {
let t = program_full_tower(&prog, row).expect("gnarly tower");
let scale_t3 =
t.t3.iter()
.flatten()
.flatten()
.fold(0.0_f64, |m, x| m.max(x.abs()))
.max(1.0);
let scale_t4 =
t.t4.iter()
.flatten()
.flatten()
.flatten()
.fold(0.0_f64, |m, x| m.max(x.abs()))
.max(1.0);
for i in 0..3 {
for j in 0..3 {
for k in 0..3 {
let base = t.t3[i][j][k];
let idx = [i, j, k];
for p in &perms3 {
let permed = t.t3[idx[p[0]]][idx[p[1]]][idx[p[2]]];
assert!(
(base - permed).abs() <= 1e-12 * scale_t3,
"row {row}: t3[{i}][{j}][{k}]={base:+.15e} != \
permuted {permed:+.15e} under {p:?}"
);
}
for l in 0..3 {
let base4 = t.t4[i][j][k][l];
let idx4 = [i, j, k, l];
for p in &perms4 {
let permed = t.t4[idx4[p[0]]][idx4[p[1]]][idx4[p[2]]][idx4[p[3]]];
assert!(
(base4 - permed).abs() <= 1e-12 * scale_t4,
"row {row}: t4[{i}][{j}][{k}][{l}]={base4:+.15e} != \
permuted {permed:+.15e} under {p:?}"
);
}
}
}
}
}
}
}
}
#[cfg(test)]
mod derivative_stack_tests {
use super::*;
#[test]
fn ln_gamma_derivative_stack_known_values_at_1() {
let s = ln_gamma_derivative_stack(1.0);
assert!(s[0].abs() < 1e-14, "ln_gamma(1) must be ~0, got {}", s[0]);
let euler_mascheroni = 0.577_215_664_901_532_9_f64;
assert!(
(s[1] + euler_mascheroni).abs() < 1e-10,
"digamma(1) ≈ -{euler_mascheroni:.6}, got {}",
s[1]
);
let pi2_6 = std::f64::consts::PI * std::f64::consts::PI / 6.0;
assert!(
(s[2] - pi2_6).abs() < 1e-10,
"trigamma(1) ≈ {pi2_6:.6}, got {}",
s[2]
);
}
#[test]
fn ln_gamma_derivative_stack_known_values_at_2() {
let s = ln_gamma_derivative_stack(2.0);
assert!(s[0].abs() < 1e-14, "ln_gamma(2) must be 0, got {}", s[0]);
let euler_mascheroni = 0.577_215_664_901_532_9_f64;
let digamma_2 = 1.0 - euler_mascheroni;
assert!(
(s[1] - digamma_2).abs() < 1e-10,
"digamma(2) ≈ {digamma_2:.6}, got {}",
s[1]
);
}
#[test]
fn ln_gamma_derivative_stack_order2_is_prefix() {
for &x in &[0.5_f64, 1.0, 2.0, 5.0] {
let full = ln_gamma_derivative_stack(x);
let ord2 = ln_gamma_derivative_stack_order2(x);
assert_eq!(ord2[0], full[0], "order2[0] != full[0] at x={x}");
assert_eq!(ord2[1], full[1], "order2[1] != full[1] at x={x}");
assert_eq!(ord2[2], full[2], "order2[2] != full[2] at x={x}");
}
}
#[test]
fn digamma_derivative_stack_overlaps_ln_gamma_stack() {
for &x in &[0.5_f64, 1.0, 2.0, 7.0] {
let lg = ln_gamma_derivative_stack(x);
let dg = digamma_derivative_stack(x);
for i in 0..4 {
assert_eq!(
lg[i + 1],
dg[i],
"ln_gamma_stack[{}] != digamma_stack[{}] at x={x}",
i + 1,
i
);
}
}
}
}
#[cfg(test)]
mod contraction_symmetry_tests {
use super::*;
struct Rng(u64);
impl Rng {
fn u(&mut self) -> f64 {
self.0 = self
.0
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(self.0 >> 11) as f64 / (1u64 << 53) as f64
}
fn s(&mut self) -> f64 {
(self.u() - 0.5) * 4.0
}
}
fn rand_sym4<const K: usize>(r: &mut Rng) -> Tower4<K> {
let mut t = Tower4::<K>::zero();
t.v = r.s();
for i in 0..K {
t.g[i] = r.s();
}
for a in 0..K {
for b in a..K {
let v2 = r.s();
t.h[a][b] = v2;
t.h[b][a] = v2;
for c in b..K {
let v3 = r.s();
for p in perms3([a, b, c]) {
t.t3[p[0]][p[1]][p[2]] = v3;
}
for d in c..K {
let v4 = r.s();
for p in perms4([a, b, c, d]) {
t.t4[p[0]][p[1]][p[2]][p[3]] = v4;
}
}
}
}
}
t
}
fn perms3(idx: [usize; 3]) -> [[usize; 3]; 6] {
let [a, b, c] = idx;
[
[a, b, c],
[a, c, b],
[b, a, c],
[b, c, a],
[c, a, b],
[c, b, a],
]
}
fn perms4(idx: [usize; 4]) -> [[usize; 4]; 24] {
let [a, b, c, d] = idx;
[
[a, b, c, d],
[a, b, d, c],
[a, c, b, d],
[a, c, d, b],
[a, d, b, c],
[a, d, c, b],
[b, a, c, d],
[b, a, d, c],
[b, c, a, d],
[b, c, d, a],
[b, d, a, c],
[b, d, c, a],
[c, a, b, d],
[c, a, d, b],
[c, b, a, d],
[c, b, d, a],
[c, d, a, b],
[c, d, b, a],
[d, a, b, c],
[d, a, c, b],
[d, b, a, c],
[d, b, c, a],
[d, c, a, b],
[d, c, b, a],
]
}
fn third_full<const K: usize>(t: &Tower4<K>, dir: &[f64; K]) -> [[f64; K]; K] {
let mut out = [[0.0; K]; K];
for a in 0..K {
for b in 0..K {
let mut acc = 0.0;
for c in 0..K {
acc += t.t3[a][b][c] * dir[c];
}
out[a][b] = acc;
}
}
out
}
fn fourth_full<const K: usize>(t: &Tower4<K>, u: &[f64; K], w: &[f64; K]) -> [[f64; K]; K] {
let mut out = [[0.0; K]; K];
for i in 0..K {
for j in 0..K {
let mut acc = 0.0;
for k in 0..K {
for l in 0..K {
acc += t.t4[i][j][k][l] * u[k] * w[l];
}
}
out[i][j] = acc;
}
}
out
}
fn check_bit_identical<const K: usize>(seed: u64, n: usize) -> usize {
let mut r = Rng(seed);
let mut checks = 0usize;
for _ in 0..n {
let t = rand_sym4::<K>(&mut r);
let dir: [f64; K] = std::array::from_fn(|_| r.s());
let u: [f64; K] = std::array::from_fn(|_| r.s());
let w: [f64; K] = std::array::from_fn(|_| r.s());
let t3_sym = t.third_contracted(&dir);
let t3_full = third_full(&t, &dir);
let t4_sym = t.fourth_contracted(&u, &w);
let t4_full = fourth_full(&t, &u, &w);
for a in 0..K {
for b in 0..K {
assert_eq!(
t3_sym[a][b].to_bits(),
t3_full[a][b].to_bits(),
"third K={K} [{a}][{b}]"
);
assert_eq!(
t4_sym[a][b].to_bits(),
t4_full[a][b].to_bits(),
"fourth K={K} [{a}][{b}]"
);
checks += 2;
}
}
}
checks
}
#[test]
fn contraction_symmetry_is_bit_identical_to_full_nest() {
let checks = check_bit_identical::<2>(0x0000_0002_C0FF_EE01, 1000)
+ check_bit_identical::<3>(0x0000_0003_C0FF_EE01, 800)
+ check_bit_identical::<4>(0x0000_0004_C0FF_EE01, 600)
+ check_bit_identical::<9>(0x0000_0009_C0FF_EE01, 300);
assert_eq!(checks, 8000 + 14400 + 19200 + 48600);
}
#[test]
fn contraction_symmetry_speedup_is_reported() {
use crate::paired_timing::{SpeedGate, paired_interleaved};
const K: usize = 9;
let mut r = Rng(0xC0FF_EE99_1234_5678);
let towers: Vec<Tower4<K>> = (0..512).map(|_| rand_sym4::<K>(&mut r)).collect();
let dir: [f64; K] = std::array::from_fn(|_| r.s());
let u: [f64; K] = std::array::from_fn(|_| r.s());
let w: [f64; K] = std::array::from_fn(|_| r.s());
if cfg!(debug_assertions) {
return;
}
let mut gate = SpeedGate::open("CONTRACTION-SYMMETRY-932");
let timing = paired_interleaved(
15,
20,
0x9320_5E11,
|nudge| {
let mut dir = dir;
dir[0] += nudge;
let mut u = u;
u[0] += nudge;
let mut sink = 0.0f64;
for t in &towers {
let o3 = t.third_contracted(&dir);
let o4 = t.fourth_contracted(&u, &w);
sink += o3[0][K - 1] + o4[0][K - 1];
}
sink
},
|nudge| {
let mut dir = dir;
dir[0] += nudge;
let mut u = u;
u[0] += nudge;
let mut sink = 0.0f64;
for t in &towers {
let o3 = third_full(t, &dir);
let o4 = fourth_full(t, &u, &w);
sink += o3[0][K - 1] + o4[0][K - 1];
}
sink
},
);
gate.faster(
&format!("K={K} towers={}", towers.len()),
&timing,
"symmetric",
"full_nest",
);
gate.finish();
}
}