File size: 5,009 Bytes
9425aed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
// PIRTM recurrence step — ported from foundry-intel/crates/z-mos/pirtm/src/lib.rs
// + jordan_contraction matching jordan_block.f90

use nalgebra::DMatrix;
use ndarray::{Array1, Array2};
use num_complex::Complex64;
use thiserror::Error;

const PHI: f64 = 1.6180339887498948482;
pub const PHI_INV: f64 = 1.0 / PHI;    // ≈ 0.618 — Fibonacci contraction rate

#[derive(Error, Debug)]
pub enum PirtmError {
    #[error("Dimension mismatch: expected {expected}, got {found}")]
    DimensionMismatch { expected: usize, found: usize },
    #[error("Contractivity violation: c_lambda={c_lambda} >= 1-epsilon={bound}")]
    ContractivityViolation { c_lambda: f64, bound: f64 },
}

pub struct PirtmMetadata {
    pub q_t: f64,
    pub c_lambda: f64,
    pub margin: f64,
}

/// PIRTM recurrence: X_{t+1} = (1-λ_m)X_t + λ_m P(Ξ_t X_t + Λ_t sigmoid(X_t) + G_t)
/// Ported verbatim from foundry-intel z-mos/pirtm.
pub fn step(
    x_t: &Array1<f64>,
    xi_t: &Array2<f64>,
    lambda_t: &Array2<f64>,
    g_t: Option<&Array1<f64>>,
    epsilon: f64,
    lambda_m: f64,
    l_t: f64,
) -> Result<(Array1<f64>, PirtmMetadata), PirtmError> {
    let dim = x_t.len();
    if xi_t.nrows() != dim || xi_t.ncols() != dim {
        return Err(PirtmError::DimensionMismatch { expected: dim, found: xi_t.nrows() });
    }
    if lambda_t.nrows() != dim || lambda_t.ncols() != dim {
        return Err(PirtmError::DimensionMismatch { expected: dim, found: lambda_t.nrows() });
    }
    let g = match g_t {
        Some(g) if g.len() == dim => g.clone(),
        Some(g) => return Err(PirtmError::DimensionMismatch { expected: dim, found: g.len() }),
        None => Array1::zeros(dim),
    };

    let term1 = xi_t.dot(x_t);
    let tx_t = x_t.mapv(|v| 1.0 / (1.0 + (-v).exp()));
    let term2 = lambda_t.dot(&tx_t);
    let y_t = term1 + term2 + g;
    let py_t = y_t.mapv(|v| v.clamp(-1.0, 1.0));
    let x_next = if (lambda_m - 1.0).abs() < f64::EPSILON {
        py_t
    } else {
        (1.0 - lambda_m) * x_t + lambda_m * py_t
    };

    let n_xi = spectral_norm_real(xi_t);
    let n_lam = spectral_norm_real(lambda_t);
    let c_lambda = (1.0 - lambda_m) + lambda_m * (n_xi + n_lam * l_t);
    let margin = (1.0 - epsilon) - c_lambda;

    Ok((x_next, PirtmMetadata { q_t: n_xi + n_lam * l_t, c_lambda, margin }))
}

/// Power-iteration spectral norm (largest singular value) for real matrices.
pub fn spectral_norm_real(m: &Array2<f64>) -> f64 {
    let dim = m.ncols();
    let mut v = Array1::from_elem(dim, 1.0 / (dim as f64).sqrt());
    let ata = m.t().dot(m);
    for _ in 0..10 {
        let w = ata.dot(&v);
        let norm = w.dot(&w).sqrt();
        if norm < 1e-10 {
            return 0.0;
        }
        v = w / norm;
    }
    (v.dot(&ata.dot(&v))).sqrt()
}

/// Jordan contraction matching jordan_block.f90:
///   ρ' = φ⁻¹ · (U ρ U†) + φ⁻² · ρ
///
/// Contraction rate φ⁻¹ ≈ 0.618 < 1 guarantees Banach fixed-point convergence.
pub fn jordan_contraction(
    rho: &DMatrix<Complex64>,
    u: &DMatrix<Complex64>,
    phi_inv: f64,
) -> DMatrix<Complex64> {
    let u_dag = u.adjoint();
    let evolved = u * rho * &u_dag;
    evolved * Complex64::new(phi_inv, 0.0) + rho * Complex64::new(phi_inv * phi_inv, 0.0)
}

/// Iterates jordan_contraction to fixed point ρ* where T(ρ*) = ρ*.
pub fn jordan_fixpoint(
    rho0: &DMatrix<Complex64>,
    u: &DMatrix<Complex64>,
    max_iter: usize,
    tol: f64,
) -> DMatrix<Complex64> {
    let mut rho = rho0.clone();
    for _ in 0..max_iter {
        let rho_next = jordan_contraction(&rho, u, PHI_INV);
        let diff = (&rho_next - &rho).norm();
        rho = rho_next;
        if diff < tol {
            break;
        }
    }
    rho
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    #[test]
    fn test_step_basic() {
        let x = array![1.0, 0.5];
        let xi = array![[0.8, 0.0], [0.0, 0.8]];
        let lam = array![[0.1, 0.0], [0.0, 0.1]];
        let (x_next, meta) = step(&x, &xi, &lam, None, 0.05, 1.0, 1.0).unwrap();
        assert_eq!(x_next.len(), 2);
        assert!(meta.margin > 0.0);
    }

    #[test]
    fn test_spectral_norm_diagonal() {
        let m = array![[2.0, 0.0], [0.0, 3.0]];
        let norm = spectral_norm_real(&m);
        assert!((norm - 3.0).abs() < 1e-5);
    }

    #[test]
    fn test_jordan_contracts() {
        let n = 2;
        let rho = DMatrix::<Complex64>::identity(n, n) * Complex64::new(0.5, 0.0);
        let u = DMatrix::<Complex64>::identity(n, n);
        let rho2 = jordan_contraction(&rho, &u, PHI_INV);
        // trace should be less than original (contraction)
        let tr0: f64 = (0..n).map(|i| rho[(i, i)].re).sum();
        let tr1: f64 = (0..n).map(|i| rho2[(i, i)].re).sum();
        assert!(tr1 < tr0, "contraction should reduce trace: {} -> {}", tr0, tr1);
    }
}