tensor4all_core/defaults/
qr.rs

1//! QR decomposition for tensors.
2//!
3//! This module works with concrete types (`DynIndex`, `IdxTensor`) only.
4
5use crate::defaults::idx_tensor::unfold_split_inner;
6use crate::defaults::DynIndex;
7use crate::global_default::GlobalDefault;
8use crate::IdxTensor;
9use num_complex::{Complex64, ComplexFloat};
10use tenferro::DType;
11use tenferro_linalg::EagerTensorLinalgExt;
12use thiserror::Error;
13
14/// Error type for QR operations in tensor4all-linalg.
15#[derive(Debug, Error)]
16pub enum QrError {
17    /// QR computation failed.
18    #[error("QR computation failed: {0}")]
19    ComputationError(#[from] anyhow::Error),
20    /// Invalid relative tolerance value (must be finite and non-negative).
21    #[error("Invalid rtol value: {0}. rtol must be finite and non-negative.")]
22    InvalidRtol(f64),
23}
24
25/// Options for QR decomposition with truncation control.
26///
27/// # Examples
28///
29/// ```
30/// use tensor4all_core::qr::{QrOptions, qr_with};
31/// use tensor4all_core::{DynIndex, TensorContractionLike, IdxTensor};
32///
33/// let i = DynIndex::new_dyn(3);
34/// let j = DynIndex::new_dyn(3);
35/// let data: Vec<f64> = (0..9).map(|x| x as f64).collect();
36/// let tensor = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
37///
38/// let opts = QrOptions::new().with_rtol(1e-10);
39/// let (q, r) = qr_with::<f64>(&tensor, &[i], &opts).unwrap();
40///
41/// // Q * R recovers the original tensor
42/// let recovered = q.contract_pair(&r).unwrap();
43/// assert!(tensor.distance(&recovered).unwrap() < 1e-12);
44/// ```
45#[derive(Debug, Clone, Copy)]
46pub struct QrOptions {
47    /// Relative tolerance for QR row-norm truncation.
48    /// If `None`, uses the global default.
49    pub rtol: Option<f64>,
50    truncate: bool,
51}
52
53impl Default for QrOptions {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59impl QrOptions {
60    /// Create new QR options with no overrides.
61    #[must_use]
62    pub fn new() -> Self {
63        Self {
64            rtol: None,
65            truncate: true,
66        }
67    }
68
69    /// Set the QR truncation tolerance.
70    #[must_use]
71    pub fn with_rtol(mut self, rtol: f64) -> Self {
72        self.rtol = Some(rtol);
73        self
74    }
75
76    pub(crate) fn full_rank() -> Self {
77        Self {
78            rtol: None,
79            truncate: false,
80        }
81    }
82}
83
84// Global default rtol using the unified GlobalDefault type
85// Default value: 1e-15 (very strict, near machine precision)
86static DEFAULT_QR_RTOL: GlobalDefault = GlobalDefault::new(1e-15);
87
88/// Get the global default rtol for QR truncation.
89///
90/// The default value is 1e-15 (very strict, near machine precision).
91pub fn default_qr_rtol() -> f64 {
92    DEFAULT_QR_RTOL.get()
93}
94
95/// Set the global default rtol for QR truncation.
96///
97/// # Arguments
98/// * `rtol` - Relative tolerance (must be finite and non-negative)
99///
100/// # Errors
101/// Returns `QrError::InvalidRtol` if rtol is not finite or is negative.
102pub fn set_default_qr_rtol(rtol: f64) -> Result<(), QrError> {
103    DEFAULT_QR_RTOL
104        .set(rtol)
105        .map_err(|e| QrError::InvalidRtol(e.0))
106}
107
108fn compute_retained_rank_qr_from_dense<T>(
109    r_full: &[T],
110    k: usize,
111    n: usize,
112    rtol: f64,
113) -> Result<usize, QrError>
114where
115    T: ComplexFloat,
116    <T as ComplexFloat>::Real: Into<f64>,
117{
118    if k == 0 || n == 0 {
119        return Ok(1);
120    }
121
122    let max_diag = k.min(n);
123
124    // Compute row norms of R (upper triangular: row i has entries from column i..n).
125    // Use relative comparison against the maximum row norm, matching
126    // compute_retained_rank_qr. The previous implementation compared diagonal
127    // elements absolutely and broke at the first small value, which is incorrect
128    // for non-pivoted QR where diagonal elements are not necessarily in
129    // decreasing order.
130    let row_norms: Vec<f64> = (0..max_diag)
131        .map(|i| {
132            let mut norm_sq: f64 = 0.0;
133            for j in i..n {
134                let val: f64 = r_full[i + j * k].abs().into();
135                norm_sq += val * val;
136            }
137            norm_sq.sqrt()
138        })
139        .collect();
140
141    let max_row_norm = row_norms.iter().cloned().fold(0.0_f64, f64::max);
142    if max_row_norm == 0.0 {
143        return Ok(1);
144    }
145
146    let threshold = rtol * max_row_norm;
147    let r = row_norms.iter().filter(|&&norm| norm >= threshold).count();
148    Ok(r.max(1))
149}
150
151/// Compute QR decomposition of a tensor with arbitrary rank, returning (Q, R).
152///
153/// This function uses the global default rtol for truncation.
154/// See `qr_with` for per-call rtol control.
155///
156/// This function computes the thin QR decomposition, where for an unfolded matrix A (m×n),
157/// we return Q (m×k) and R (k×n) with k = min(m, n).
158///
159/// The input tensor can have any rank >= 2, and indices are split into left and right groups.
160/// The tensor is unfolded into a matrix by grouping left indices as rows and right indices as columns.
161///
162/// Truncation is performed based on R's row norms: rows whose norm is below
163/// `rtol * max_row_norm` are discarded.
164///
165/// For the mathematical convention:
166/// \[ A = Q * R \]
167/// where Q is orthogonal (or unitary for complex) and R is upper triangular.
168///
169/// # Arguments
170/// * `t` - Input tensor with DenseF64 or DenseC64 storage
171/// * `left_inds` - Indices to place on the left (row) side of the unfolded matrix
172///
173/// # Returns
174/// A tuple `(Q, R)` where:
175/// - `Q` is a tensor with indices `[left_inds..., bond_index]` and dimensions `[left_dims..., r]`
176/// - `R` is a tensor with indices `[bond_index, right_inds...]` and dimensions `[r, right_dims...]`
177///
178///   where `r` is the retained rank (≤ min(m, n)) determined by rtol truncation.
179///
180/// # Errors
181/// Returns `QrError` if:
182/// - The tensor rank is < 2
183/// - Storage is not DenseF64 or DenseC64
184/// - `left_inds` is empty or contains all indices
185/// - `left_inds` contains indices not in the tensor or duplicates
186/// - The QR computation fails
187///
188/// # Examples
189///
190/// ```
191/// use tensor4all_core::{IdxTensor, DynIndex, qr};
192///
193/// // Create a 4x3 matrix
194/// let i = DynIndex::new_dyn(4);
195/// let j = DynIndex::new_dyn(3);
196/// // Identity-like data (4x3 column-major)
197/// let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
198/// let t = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
199///
200/// let (q, r) = qr::<f64>(&t, &[i.clone()]).unwrap();
201///
202/// // Q has shape (4, bond) and R has shape (bond, 3)
203/// assert_eq!(q.dims()[0], 4);
204/// assert_eq!(r.dims()[r.dims().len() - 1], 3);
205/// ```
206pub fn qr<T>(t: &IdxTensor, left_inds: &[DynIndex]) -> Result<(IdxTensor, IdxTensor), QrError> {
207    qr_with::<T>(t, left_inds, &QrOptions::default())
208}
209
210/// Compute QR decomposition of a tensor with arbitrary rank, returning (Q, R).
211///
212/// This function allows per-call control of the truncation tolerance via `QrOptions`.
213/// If `options.rtol` is `None`, uses the global default rtol.
214///
215/// This function computes the thin QR decomposition, where for an unfolded matrix A (m×n),
216/// we return Q (m×k) and R (k×n) with k = min(m, n).
217///
218/// The input tensor can have any rank >= 2, and indices are split into left and right groups.
219/// The tensor is unfolded into a matrix by grouping left indices as rows and right indices as columns.
220///
221/// Truncation is performed based on R's row norms: rows whose norm is below
222/// `rtol * max_row_norm` are discarded.
223///
224/// For the mathematical convention:
225/// \[ A = Q * R \]
226/// where Q is orthogonal (or unitary for complex) and R is upper triangular.
227///
228/// # Arguments
229/// * `t` - Input tensor with DenseF64 or DenseC64 storage
230/// * `left_inds` - Indices to place on the left (row) side of the unfolded matrix
231/// * `options` - QR options including rtol for truncation control
232///
233/// # Returns
234/// A tuple `(Q, R)` where:
235/// - `Q` is a tensor with indices `[left_inds..., bond_index]` and dimensions `[left_dims..., r]`
236/// - `R` is a tensor with indices `[bond_index, right_inds...]` and dimensions `[r, right_dims...]`
237///
238///   where `r` is the retained rank (≤ min(m, n)) determined by rtol truncation.
239///
240/// # Errors
241/// Returns `QrError` if:
242/// - The tensor rank is < 2
243/// - Storage is not DenseF64 or DenseC64
244/// - `left_inds` is empty or contains all indices
245/// - `left_inds` contains indices not in the tensor or duplicates
246/// - The QR computation fails
247/// - `options.rtol` is invalid (not finite or negative)
248pub fn qr_with<T>(
249    t: &IdxTensor,
250    left_inds: &[DynIndex],
251    options: &QrOptions,
252) -> Result<(IdxTensor, IdxTensor), QrError> {
253    // Unfold tensor into an eager rank-2 tensor so linalg AD nodes stay connected.
254    let (matrix_inner, _, m, n, left_indices, right_indices) = unfold_split_inner(t, left_inds)
255        .map_err(|e| anyhow::anyhow!("Failed to unfold tensor: {}", e))
256        .map_err(QrError::ComputationError)?;
257    let k = m.min(n);
258    let (mut q_inner, mut r_inner) = matrix_inner
259        .qr()
260        .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
261
262    let r = if options.truncate {
263        // Determine rtol to use
264        let rtol = options.rtol.unwrap_or(default_qr_rtol());
265        if !rtol.is_finite() || rtol < 0.0 {
266            return Err(QrError::InvalidRtol(rtol));
267        }
268
269        let value = r_inner
270            .value()
271            .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
272        match r_inner.dtype() {
273            DType::F64 => {
274                let values = value
275                    .as_slice::<f64>()
276                    .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
277                compute_retained_rank_qr_from_dense(values, k, n, rtol)?
278            }
279            DType::C64 => {
280                let values = value
281                    .as_slice::<Complex64>()
282                    .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
283                compute_retained_rank_qr_from_dense(values, k, n, rtol)?
284            }
285            other => {
286                return Err(QrError::ComputationError(anyhow::anyhow!(
287                    "native QR returned unsupported scalar type {other:?}"
288                )));
289            }
290        }
291    } else {
292        k
293    };
294    if r < k {
295        let keep: Vec<usize> = (0..r).collect();
296        q_inner = q_inner
297            .take_axis(1, &keep)
298            .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
299        r_inner = r_inner
300            .take_axis(0, &keep)
301            .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
302    }
303
304    let bond_index = DynIndex::new_bond(r)
305        .map_err(|e| anyhow::anyhow!("Failed to create Link index: {:?}", e))
306        .map_err(QrError::ComputationError)?;
307
308    let mut q_indices = left_indices.clone();
309    q_indices.push(bond_index.clone());
310    let q_dims: Vec<usize> = q_indices.iter().map(|idx| idx.dim).collect();
311    let q_reshaped = q_inner.reshape(&q_dims).map_err(|e| {
312        QrError::ComputationError(anyhow::anyhow!("eager QR Q reshape failed: {e}"))
313    })?;
314    let q = IdxTensor::from_inner(q_indices, q_reshaped).map_err(QrError::ComputationError)?;
315
316    let mut r_indices = vec![bond_index.clone()];
317    r_indices.extend_from_slice(&right_indices);
318    let r_dims: Vec<usize> = r_indices.iter().map(|idx| idx.dim).collect();
319    let r_reshaped = r_inner.reshape(&r_dims).map_err(|e| {
320        QrError::ComputationError(anyhow::anyhow!("eager QR R reshape failed: {e}"))
321    })?;
322    let r = IdxTensor::from_inner(r_indices, r_reshaped).map_err(QrError::ComputationError)?;
323
324    Ok((q, r))
325}
326
327#[cfg(test)]
328mod tests;
tensor4all_core/defaults/qr.rs

tensor4all_core/defaults/
qr.rs
