cram_codec/
tt.rs

1//! Tensor trains: TT-SVD with error-budget / Gavish–Donoho truncation, and
2//! full contraction.
3//!
4//! A tensor `X(i₁,…,i_d)` is written `G₁(i₁)·G₂(i₂)⋯G_d(i_d)` with core
5//! `G_k` of shape `r_{k−1} × n_k × r_k` (`r₀ = r_d = 1`), stored row-major
6//! as the `(r_{k−1}·n_k) × r_k` matrix. [`tt_svd`] returns the **canonical
7//! left-orthogonal gauge**: every core but the last has orthonormal columns
8//! as a `(r_{k−1}n_k) × r_k` matrix, so the whole norm `‖X‖_F` lives in
9//! `G_d`, and the TT is unique up to an orthogonal `r_k × r_k` rotation per
10//! bond (fixed here by the SVD's singular-vector basis, i.e. the rotation that
11//! diagonalises each left-unfolding Gram).
12
13use crate::linalg::{matmul, matmul_tn, svd_left};
14
15/// A tensor train in row-major core storage.
16#[derive(Clone, Debug, PartialEq)]
17pub struct TensorTrain {
18    /// Mode sizes `n_k`.
19    pub modes: Vec<usize>,
20    /// Bond ranks `r₀..r_d` (`r₀ = r_d = 1`).
21    pub ranks: Vec<usize>,
22    /// Cores, `G_k` as `(r_{k−1}·n_k) × r_k` row-major.
23    pub cores: Vec<Vec<f64>>,
24}
25
26impl TensorTrain {
27    /// Number of stored parameters.
28    pub fn num_params(&self) -> usize {
29        self.cores.iter().map(|c| c.len()).sum()
30    }
31
32    /// Contract to the full tensor (length `∏ n_k`), left to right.
33    pub fn full(&self) -> Vec<f64> {
34        contract(&self.modes, &self.ranks, &self.cores)
35    }
36}
37
38/// Contract cores to the full tensor. Each step is one GEMM:
39/// `P (M × r_{k−1}) · G_k (r_{k−1} × n_k r_k)`, and the row-major result
40/// *is* `(M n_k) × r_k`.
41pub fn contract(modes: &[usize], ranks: &[usize], cores: &[Vec<f64>]) -> Vec<f64> {
42    let d = modes.len();
43    let mut p = cores[0].clone(); // (1·n₁) × r₁
44    let mut rows = modes[0];
45    for k in 1..d {
46        let r = ranks[k];
47        let cols = modes[k] * ranks[k + 1];
48        p = matmul(&p, &cores[k], rows, r, cols);
49        rows *= modes[k];
50    }
51    p
52}
53
54/// Rank-selection rule for each unfolding in [`tt_svd`].
55#[derive(Clone, Copy, Debug, PartialEq)]
56pub struct Truncation {
57    /// Relative Frobenius budget `ε`: each of the `d−1` SVDs discards at most
58    /// `δ = ε‖X‖/√(d−1)` (Oseledets 2011), so `‖X − X̂‖ ≤ ε‖X‖`.
59    pub rel_eps: f64,
60    /// Additionally cap each rank at the Gavish–Donoho optimal hard
61    /// threshold `ω(β)·σ_med` (unknown-noise form).
62    pub gavish_donoho: bool,
63    /// Hard rank cap.
64    pub max_rank: usize,
65}
66
67impl Default for Truncation {
68    fn default() -> Self {
69        Self {
70            rel_eps: 1e-12,
71            gavish_donoho: false,
72            max_rank: usize::MAX,
73        }
74    }
75}
76
77/// Gavish–Donoho (2014) optimal hard threshold coefficient for unknown noise,
78/// `ω(β) ≈ 0.56β³ − 0.95β² + 1.82β + 1.43`, `β = min(m,n)/max(m,n)`.
79pub fn gd_omega(beta: f64) -> f64 {
80    0.56 * beta.powi(3) - 0.95 * beta.powi(2) + 1.82 * beta + 1.43
81}
82
83/// Number of singular values above the Gavish–Donoho threshold for an
84/// `m × n` matrix with spectrum `s` (descending, length `min(m,n)`).
85///
86/// The rule needs a bulk of noise singular values to estimate the median
87/// from; QTT's leading unfoldings are `2 × 2^{d−1}`, `4 × …` (two or four
88/// singular values), where the median is a *signal* value and the cap
89/// collapses the rank to 1 (measured: Miranda 64×128×128 at a 60 dB target
90/// fell to 12 dB). Unfoldings with `min(m,n) < GD_MIN_DIM` are left uncapped.
91pub fn gd_rank(s: &[f64], m: usize, n: usize) -> usize {
92    if s.is_empty() {
93        return 0;
94    }
95    if m.min(n) < GD_MIN_DIM {
96        return s.len();
97    }
98    let beta = m.min(n) as f64 / m.max(n) as f64;
99    let mut sorted = s.to_vec();
100    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
101    let len = sorted.len();
102    let med = if len % 2 == 1 {
103        sorted[len / 2]
104    } else {
105        0.5 * (sorted[len / 2 - 1] + sorted[len / 2])
106    };
107    let tau = gd_omega(beta) * med;
108    s.iter().filter(|&&x| x > tau).count().max(1)
109}
110
111/// Smallest unfolding side on which [`gd_rank`] applies its cap.
112pub const GD_MIN_DIM: usize = 32;
113
114/// Smallest rank whose discarded tail energy `Σ_{j≥r} σ_j²` is ≤ `delta2`.
115pub fn tail_rank(s: &[f64], delta2: f64) -> usize {
116    let mut tail = 0.0;
117    let mut r = s.len();
118    while r > 1 {
119        let t = tail + s[r - 1] * s[r - 1];
120        if t > delta2 {
121            break;
122        }
123        tail = t;
124        r -= 1;
125    }
126    r
127}
128
129/// TT-SVD (Oseledets 2011) of `x` reshaped to `modes`, left-orthogonal gauge.
130pub fn tt_svd(x: &[f64], modes: &[usize], trunc: Truncation) -> TensorTrain {
131    let d = modes.len();
132    let n: usize = modes.iter().product();
133    assert_eq!(x.len(), n, "tt_svd: data length != product of modes");
134    let norm2: f64 = x.iter().map(|v| v * v).sum();
135    let delta2 = if d > 1 {
136        trunc.rel_eps * trunc.rel_eps * norm2 / (d - 1) as f64
137    } else {
138        0.0
139    };
140    let mut ranks = vec![1usize];
141    let mut cores = Vec::with_capacity(d);
142    let mut c = x.to_vec();
143    let mut r_prev = 1;
144    let mut cols = n;
145    for &nk in modes.iter().take(d - 1) {
146        let m = r_prev * nk;
147        cols /= nk;
148        let (_s, u, r) = svd_left(&c, m, cols, |s| {
149            let mut r = tail_rank(s, delta2);
150            if trunc.gavish_donoho {
151                r = r.min(gd_rank(s, m, cols));
152            }
153            r.min(trunc.max_rank).max(1)
154        });
155        // Remainder Uᵀ C (r × cols) = Σ Vᵀ.
156        c = matmul_tn(&u, &c, m, r, cols);
157        cores.push(u);
158        ranks.push(r);
159        r_prev = r;
160    }
161    cores.push(c);
162    ranks.push(1);
163    TensorTrain {
164        modes: modes.to_vec(),
165        ranks,
166        cores,
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    fn rel_err(a: &[f64], b: &[f64]) -> f64 {
175        let num: f64 = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum();
176        let den: f64 = a.iter().map(|x| x * x).sum();
177        (num / den).sqrt()
178    }
179
180    #[test]
181    fn exact_rank_of_exponential_and_sine() {
182        // exp(a i) has QTT rank 1; sin(a i) rank 2 (Khoromskij).
183        let n = 1 << 12;
184        let modes = vec![2; 12];
185        let e: Vec<f64> = (0..n).map(|i| (0.001 * i as f64).exp()).collect();
186        // ε = 1e-6: the Gram-based SVD resolves σ only to ~1e-7·σ₁ (crate docs).
187        let tt = tt_svd(
188            &e,
189            &modes,
190            Truncation {
191                rel_eps: 1e-6,
192                ..Default::default()
193            },
194        );
195        assert!(tt.ranks.iter().all(|&r| r == 1), "{:?}", tt.ranks);
196        let s: Vec<f64> = (0..n).map(|i| (0.01 * i as f64 + 0.3).sin()).collect();
197        let tt = tt_svd(
198            &s,
199            &modes,
200            Truncation {
201                rel_eps: 1e-6,
202                ..Default::default()
203            },
204        );
205        assert!(tt.ranks.iter().all(|&r| r <= 2), "{:?}", tt.ranks);
206        assert!(rel_err(&s, &tt.full()) < 1e-6);
207    }
208
209    #[test]
210    fn error_budget_holds_and_gauge_is_left_orthogonal() {
211        let modes = vec![2usize; 14];
212        let n = 1 << 14;
213        let mut seed = 1u64;
214        let x: Vec<f64> = (0..n)
215            .map(|i| {
216                seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
217                let t = i as f64 / n as f64;
218                (20.0 * t).sin() * (-3.0 * t).exp() + 0.01 * ((seed >> 33) as f64 / 2e9 - 1.0)
219            })
220            .collect();
221        for eps in [1e-1, 1e-2, 1e-3] {
222            let tt = tt_svd(
223                &x,
224                &modes,
225                Truncation {
226                    rel_eps: eps,
227                    ..Default::default()
228                },
229            );
230            assert!(rel_err(&x, &tt.full()) <= eps * (1.0 + 1e-9));
231            for k in 0..modes.len() - 1 {
232                let (m, r) = (tt.ranks[k] * modes[k], tt.ranks[k + 1]);
233                let g = crate::linalg::gram_cols(&tt.cores[k], m, r);
234                for i in 0..r {
235                    for j in 0..r {
236                        let want = if i == j { 1.0 } else { 0.0 };
237                        assert!((g[i * r + j] - want).abs() < 1e-8);
238                    }
239                }
240            }
241            let last = tt.cores.last().unwrap();
242            let nl: f64 = last.iter().map(|v| v * v).sum();
243            let nx: f64 = tt.full().iter().map(|v| v * v).sum();
244            assert!((nl - nx).abs() < 1e-8 * nx, "norm lives in the last core");
245        }
246    }
247
248    #[test]
249    fn gavish_donoho_cuts_noise_rank() {
250        // rank-1 signal + noise in a 64×64 unfolding: GD keeps ~1.
251        let modes = vec![64usize, 64];
252        let mut seed = 3u64;
253        let x: Vec<f64> = (0..4096)
254            .map(|i| {
255                seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
256                let (a, b) = ((i / 64) as f64, (i % 64) as f64);
257                (a / 10.0).sin() * (b / 7.0).cos() * 10.0
258                    + 0.1 * ((seed >> 33) as f64 / 2147483648.0 - 0.5)
259            })
260            .collect();
261        let tt = tt_svd(
262            &x,
263            &modes,
264            Truncation {
265                rel_eps: 1e-12,
266                gavish_donoho: true,
267                ..Default::default()
268            },
269        );
270        assert_eq!(tt.ranks[1], 1);
271    }
272}
cram_codec/tt.rs

cram_codec/
tt.rs
