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::{ExecutionContext, IdxTensor};
9use num_complex::{Complex64, ComplexFloat};
10use tenferro::DType;
11use tenferro_ad::EagerTensor;
12use tenferro_linalg::EagerTensorLinalgExt;
13use thiserror::Error;
14
15/// Error type for QR operations in tensor4all-linalg.
16#[derive(Debug, Error)]
17pub enum QrError {
18    /// QR computation failed.
19    #[error("QR computation failed: {0}")]
20    ComputationError(#[from] anyhow::Error),
21    /// Invalid relative tolerance value (must be finite and non-negative).
22    #[error("Invalid rtol value: {0}. rtol must be finite and non-negative.")]
23    InvalidRtol(f64),
24}
25
26/// Options for QR decomposition with truncation control.
27///
28/// # Examples
29///
30/// ```
31/// use tensor4all_core::qr::{QrOptions, qr_with};
32/// use tensor4all_core::{DynIndex, TensorContractionLike, IdxTensor};
33///
34/// let i = DynIndex::new_dyn(3);
35/// let j = DynIndex::new_dyn(3);
36/// let data: Vec<f64> = (0..9).map(|x| x as f64).collect();
37/// let tensor = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
38///
39/// let opts = QrOptions::new().with_rtol(1e-10);
40/// let (q, r) = qr_with::<f64>(&tensor, &[i], &opts).unwrap();
41///
42/// // Q * R recovers the original tensor
43/// let recovered = q.contract_pair(&r).unwrap();
44/// assert!(tensor.distance(&recovered).unwrap() < 1e-12);
45/// ```
46#[derive(Debug, Clone, Copy)]
47pub struct QrOptions {
48    /// Relative tolerance for QR row-norm truncation.
49    /// If `None`, uses the global default.
50    pub rtol: Option<f64>,
51    truncate: bool,
52}
53
54impl Default for QrOptions {
55    fn default() -> Self {
56        Self::new()
57    }
58}
59
60impl QrOptions {
61    /// Create new QR options with no overrides.
62    #[must_use]
63    pub fn new() -> Self {
64        Self {
65            rtol: None,
66            truncate: true,
67        }
68    }
69
70    /// Set the QR truncation tolerance.
71    #[must_use]
72    pub fn with_rtol(mut self, rtol: f64) -> Self {
73        self.rtol = Some(rtol);
74        self
75    }
76
77    pub(crate) fn full_rank() -> Self {
78        Self {
79            rtol: None,
80            truncate: false,
81        }
82    }
83}
84
85// Global default rtol using the unified GlobalDefault type
86// Default value: 1e-15 (very strict, near machine precision)
87static DEFAULT_QR_RTOL: GlobalDefault = GlobalDefault::new(1e-15);
88
89/// Get the global default rtol for QR truncation.
90///
91/// The default value is 1e-15 (very strict, near machine precision).
92pub fn default_qr_rtol() -> f64 {
93    DEFAULT_QR_RTOL.get()
94}
95
96/// Set the global default rtol for QR truncation.
97///
98/// # Arguments
99/// * `rtol` - Relative tolerance (must be finite and non-negative)
100///
101/// # Errors
102/// Returns `QrError::InvalidRtol` if rtol is not finite or is negative.
103pub fn set_default_qr_rtol(rtol: f64) -> Result<(), QrError> {
104    DEFAULT_QR_RTOL
105        .set(rtol)
106        .map_err(|e| QrError::InvalidRtol(e.0))
107}
108
109fn compute_retained_rank_qr_from_dense<T>(
110    r_full: &[T],
111    k: usize,
112    n: usize,
113    rtol: f64,
114) -> Result<usize, QrError>
115where
116    T: ComplexFloat,
117    <T as ComplexFloat>::Real: Into<f64>,
118{
119    if k == 0 || n == 0 {
120        return Ok(1);
121    }
122
123    let max_diag = k.min(n);
124
125    // Compute row norms of R (upper triangular: row i has entries from column i..n).
126    // Use relative comparison against the maximum row norm, matching
127    // compute_retained_rank_qr. The previous implementation compared diagonal
128    // elements absolutely and broke at the first small value, which is incorrect
129    // for non-pivoted QR where diagonal elements are not necessarily in
130    // decreasing order.
131    let row_norms: Vec<f64> = (0..max_diag)
132        .map(|i| {
133            let mut norm_sq: f64 = 0.0;
134            for j in i..n {
135                let val: f64 = r_full[i + j * k].abs().into();
136                norm_sq += val * val;
137            }
138            norm_sq.sqrt()
139        })
140        .collect();
141
142    Ok(retained_rank_from_row_norms(&row_norms, rtol))
143}
144
145/// Select the retained QR rank from per-row norms of the upper-triangular factor.
146///
147/// Shared by the host dense path and the device path (which reduces the norms
148/// on-device and reads back only this `k`-element decision vector).
149fn retained_rank_from_row_norms(row_norms: &[f64], rtol: f64) -> usize {
150    let max_row_norm = row_norms.iter().cloned().fold(0.0_f64, f64::max);
151    if max_row_norm == 0.0 {
152        return 1;
153    }
154
155    let threshold = rtol * max_row_norm;
156    let r = row_norms.iter().filter(|&&norm| norm >= threshold).count();
157    r.max(1)
158}
159
160/// Read the `k` upper-triangular row norms of a resident `R` factor.
161///
162/// All arithmetic stays in the operand's owning runtime; only the `k`-element
163/// decision vector crosses the explicit readback boundary through `context`.
164#[cfg(feature = "tenferro-cuda")]
165fn resident_row_norms(
166    r_inner: &EagerTensor,
167    k: usize,
168    context: &ExecutionContext,
169) -> Result<Vec<f64>, QrError> {
170    let upper = r_inner
171        .triu(0)
172        .map_err(|error| QrError::ComputationError(anyhow::Error::new(error)))?;
173    let magnitudes = upper
174        .abs()
175        .map_err(|error| QrError::ComputationError(anyhow::Error::new(error)))?;
176    let sum_squares = magnitudes
177        .reduce_sum_squares(&[1])
178        .map_err(|error| QrError::ComputationError(anyhow::Error::new(error)))?;
179    if sum_squares.shape() != [k] {
180        return Err(QrError::ComputationError(anyhow::anyhow!(
181            "device QR row-norm reduction returned shape {:?}, expected [{k}]",
182            sum_squares.shape()
183        )));
184    }
185    let decision = IdxTensor::from_inner(vec![DynIndex::new_dyn(k)], sum_squares)
186        .map_err(QrError::ComputationError)?;
187    let values = decision
188        .read_decision_data(context)
189        .map_err(|error| QrError::ComputationError(anyhow::Error::new(error)))?;
190    Ok(values.into_iter().map(|value| value.sqrt()).collect())
191}
192
193/// Compute QR decomposition of a tensor with arbitrary rank, returning (Q, R).
194///
195/// This function uses the global default rtol for truncation.
196/// See `qr_with` for per-call rtol control.
197///
198/// This function computes the thin QR decomposition, where for an unfolded matrix A (m×n),
199/// we return Q (m×k) and R (k×n) with k = min(m, n).
200///
201/// The input tensor can have any rank >= 2, and indices are split into left and right groups.
202/// The tensor is unfolded into a matrix by grouping left indices as rows and right indices as columns.
203///
204/// Truncation is performed based on R's row norms: rows whose norm is below
205/// `rtol * max_row_norm` are discarded.
206///
207/// For the mathematical convention:
208/// \[ A = Q * R \]
209/// where Q is orthogonal (or unitary for complex) and R is upper triangular.
210///
211/// # Arguments
212/// * `t` - Input tensor with DenseF64 or DenseC64 storage
213/// * `left_inds` - Indices to place on the left (row) side of the unfolded matrix
214///
215/// # Returns
216/// A tuple `(Q, R)` where:
217/// - `Q` is a tensor with indices `[left_inds..., bond_index]` and dimensions `[left_dims..., r]`
218/// - `R` is a tensor with indices `[bond_index, right_inds...]` and dimensions `[r, right_dims...]`
219///
220///   where `r` is the retained rank (≤ min(m, n)) determined by rtol truncation.
221///
222/// # Errors
223/// Returns `QrError` if:
224/// - The tensor rank is < 2
225/// - Storage is not DenseF64 or DenseC64
226/// - `left_inds` is empty or contains all indices
227/// - `left_inds` contains indices not in the tensor or duplicates
228/// - The QR computation fails
229///
230/// # Examples
231///
232/// ```
233/// use tensor4all_core::{IdxTensor, DynIndex, qr};
234///
235/// // Create a 4x3 matrix
236/// let i = DynIndex::new_dyn(4);
237/// let j = DynIndex::new_dyn(3);
238/// // Identity-like data (4x3 column-major)
239/// let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
240/// let t = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
241///
242/// let (q, r) = qr::<f64>(&t, &[i.clone()]).unwrap();
243///
244/// // Q has shape (4, bond) and R has shape (bond, 3)
245/// assert_eq!(q.dims()[0], 4);
246/// assert_eq!(r.dims()[r.dims().len() - 1], 3);
247/// ```
248pub fn qr<T>(t: &IdxTensor, left_inds: &[DynIndex]) -> Result<(IdxTensor, IdxTensor), QrError> {
249    qr_with::<T>(t, left_inds, &QrOptions::default())
250}
251
252/// Compute QR decomposition of a tensor with arbitrary rank, returning (Q, R).
253///
254/// This function allows per-call control of the truncation tolerance via `QrOptions`.
255/// If `options.rtol` is `None`, uses the global default rtol.
256///
257/// This function computes the thin QR decomposition, where for an unfolded matrix A (m×n),
258/// we return Q (m×k) and R (k×n) with k = min(m, n).
259///
260/// The input tensor can have any rank >= 2, and indices are split into left and right groups.
261/// The tensor is unfolded into a matrix by grouping left indices as rows and right indices as columns.
262///
263/// Truncation is performed based on R's row norms: rows whose norm is below
264/// `rtol * max_row_norm` are discarded.
265///
266/// For the mathematical convention:
267/// \[ A = Q * R \]
268/// where Q is orthogonal (or unitary for complex) and R is upper triangular.
269///
270/// # Arguments
271/// * `t` - Input tensor with DenseF64 or DenseC64 storage
272/// * `left_inds` - Indices to place on the left (row) side of the unfolded matrix
273/// * `options` - QR options including rtol for truncation control
274///
275/// # Returns
276/// A tuple `(Q, R)` where:
277/// - `Q` is a tensor with indices `[left_inds..., bond_index]` and dimensions `[left_dims..., r]`
278/// - `R` is a tensor with indices `[bond_index, right_inds...]` and dimensions `[r, right_dims...]`
279///
280///   where `r` is the retained rank (≤ min(m, n)) determined by rtol truncation.
281///
282/// # Errors
283/// Returns `QrError` if:
284/// - The tensor rank is < 2
285/// - Storage is not DenseF64 or DenseC64
286/// - `left_inds` is empty or contains all indices
287/// - `left_inds` contains indices not in the tensor or duplicates
288/// - The QR computation fails
289/// - `options.rtol` is invalid (not finite or negative)
290pub fn qr_with<T>(
291    t: &IdxTensor,
292    left_inds: &[DynIndex],
293    options: &QrOptions,
294) -> Result<(IdxTensor, IdxTensor), QrError> {
295    if t.is_cuda_resident() {
296        return Err(QrError::ComputationError(anyhow::anyhow!(
297            "CUDA-resident input requires qr_with_in with the owning execution context"
298        )));
299    }
300    // Unfold tensor into an eager rank-2 tensor so linalg AD nodes stay connected.
301    let (matrix_inner, _, m, n, left_indices, right_indices) = unfold_split_inner(t, left_inds)
302        .map_err(|e| anyhow::anyhow!("Failed to unfold tensor: {}", e))
303        .map_err(QrError::ComputationError)?;
304    let k = m.min(n);
305    let (q_inner, r_inner) = matrix_inner
306        .qr()
307        .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
308
309    let r = qr_retained_rank(&r_inner, k, n, options, None)?;
310    qr_assemble(q_inner, r_inner, k, r, left_indices, right_indices)
311}
312
313/// Compute QR decomposition in a caller-owned execution context.
314///
315/// The factorization, truncation slicing, and result assembly all execute in
316/// `context`. With a CUDA context the retained-rank decision reads back only
317/// the `k`-element row-norm vector through [`IdxTensor::read_decision_data`];
318/// with a CPU context the rank decision uses the same host read as [`qr_with`].
319///
320/// # Arguments
321///
322/// * `t` - Input tensor, which must belong to `context`.
323/// * `left_inds` - Indices to place on the left (row) side of the unfolded matrix.
324/// * `options` - QR options including rtol for truncation control.
325/// * `context` - Caller-owned execution context owning the input and results.
326///
327/// # Examples
328///
329/// ```
330/// use std::sync::Arc;
331/// use tensor4all_core::qr::{QrOptions, qr_with_in};
332/// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor, TensorContractionLike};
333/// use tensor4all_tensorbackend::CpuExecutionContext;
334/// use tenferro_cpu::CpuBackend;
335///
336/// let context = ExecutionContext::Cpu(Arc::new(
337///     CpuExecutionContext::from_backend(CpuBackend::new()),
338/// ));
339/// let i = DynIndex::new_dyn(4);
340/// let j = DynIndex::new_dyn(3);
341/// let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
342/// let t = IdxTensor::from_dense_in(&context, vec![i.clone(), j.clone()], data)?;
343/// let (q, r) = qr_with_in::<f64>(&t, &[i.clone()], &QrOptions::default(), &context)?;
344/// let recovered = q.contract_pair(&r)?;
345/// let expected = t.to_vec::<f64>()?;
346/// let actual = recovered.to_vec::<f64>()?;
347/// let residual = expected
348///     .iter()
349///     .zip(actual.iter())
350///     .map(|(a, b)| (a - b).abs())
351///     .fold(0.0_f64, f64::max);
352/// assert!(residual < 1e-12, "QR reconstruction residual {residual}");
353/// # Ok::<(), Box<dyn std::error::Error>>(())
354/// ```
355///
356/// # Errors
357///
358/// Returns `QrError` when the tensor does not belong to `context`, when the
359/// indices, storage, or options are invalid, or when the factorization or
360/// explicit decision readback fails.
361pub fn qr_with_in<T>(
362    t: &IdxTensor,
363    left_inds: &[DynIndex],
364    options: &QrOptions,
365    context: &ExecutionContext,
366) -> Result<(IdxTensor, IdxTensor), QrError> {
367    t.validate_context(context)
368        .map_err(|error| QrError::ComputationError(anyhow::Error::new(error)))?;
369    let (matrix_inner, _, m, n, left_indices, right_indices) = unfold_split_inner(t, left_inds)
370        .map_err(|e| anyhow::anyhow!("Failed to unfold tensor: {}", e))
371        .map_err(QrError::ComputationError)?;
372    let k = m.min(n);
373    let (q_inner, r_inner) = matrix_inner
374        .qr()
375        .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
376
377    let r = qr_retained_rank(&r_inner, k, n, options, Some(context))?;
378    qr_assemble(q_inner, r_inner, k, r, left_indices, right_indices)
379}
380
381/// Select the retained QR rank, reading the decision payload through `context`
382/// for device-resident factors and directly from host memory otherwise.
383fn qr_retained_rank(
384    r_inner: &EagerTensor,
385    k: usize,
386    n: usize,
387    options: &QrOptions,
388    context: Option<&ExecutionContext>,
389) -> Result<usize, QrError> {
390    if !options.truncate {
391        return Ok(k);
392    }
393    // Determine rtol to use
394    let rtol = options.rtol.unwrap_or(default_qr_rtol());
395    if !rtol.is_finite() || rtol < 0.0 {
396        return Err(QrError::InvalidRtol(rtol));
397    }
398
399    #[cfg(feature = "tenferro-cuda")]
400    if let Some(context) = context {
401        if matches!(context, ExecutionContext::Cuda(_)) {
402            let row_norms = resident_row_norms(r_inner, k, context)?;
403            return Ok(retained_rank_from_row_norms(&row_norms, rtol));
404        }
405    }
406    #[cfg(not(feature = "tenferro-cuda"))]
407    let _ = context;
408
409    let value = r_inner
410        .value()
411        .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
412    match r_inner.dtype() {
413        DType::F64 => {
414            let values = value
415                .as_slice::<f64>()
416                .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
417            compute_retained_rank_qr_from_dense(values, k, n, rtol)
418        }
419        DType::C64 => {
420            let values = value
421                .as_slice::<Complex64>()
422                .map_err(|source| QrError::ComputationError(anyhow::Error::new(source)))?;
423            compute_retained_rank_qr_from_dense(values, k, n, rtol)
424        }
425        other => Err(QrError::ComputationError(anyhow::anyhow!(
426            "native QR returned unsupported scalar type {other:?}"
427        ))),
428    }
429}
430
431/// Slice the thin factors to the retained rank and wrap them as [`IdxTensor`]s.
432fn qr_assemble(
433    mut q_inner: EagerTensor,
434    mut r_inner: EagerTensor,
435    k: usize,
436    r: usize,
437    left_indices: Vec<DynIndex>,
438    right_indices: Vec<DynIndex>,
439) -> Result<(IdxTensor, IdxTensor), QrError> {
440    if r < k {
441        let keep: Vec<usize> = (0..r).collect();
442        q_inner = q_inner
443            .take_axis(1, &keep)
444            .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
445        r_inner = r_inner
446            .take_axis(0, &keep)
447            .map_err(|e| QrError::ComputationError(anyhow::anyhow!("{e}")))?;
448    }
449
450    let bond_index = DynIndex::new_bond(r)
451        .map_err(|e| anyhow::anyhow!("Failed to create Link index: {:?}", e))
452        .map_err(QrError::ComputationError)?;
453
454    let mut q_indices = left_indices.clone();
455    q_indices.push(bond_index.clone());
456    let q_dims: Vec<usize> = q_indices.iter().map(|idx| idx.dim).collect();
457    let q_reshaped = q_inner.reshape(&q_dims).map_err(|e| {
458        QrError::ComputationError(anyhow::anyhow!("eager QR Q reshape failed: {e}"))
459    })?;
460    let q = IdxTensor::from_inner(q_indices, q_reshaped).map_err(QrError::ComputationError)?;
461
462    let mut r_indices = vec![bond_index.clone()];
463    r_indices.extend_from_slice(&right_indices);
464    let r_dims: Vec<usize> = r_indices.iter().map(|idx| idx.dim).collect();
465    let r_reshaped = r_inner.reshape(&r_dims).map_err(|e| {
466        QrError::ComputationError(anyhow::anyhow!("eager QR R reshape failed: {e}"))
467    })?;
468    let r = IdxTensor::from_inner(r_indices, r_reshaped).map_err(QrError::ComputationError)?;
469
470    Ok((q, r))
471}
472
473#[cfg(test)]
474mod tests;
tensor4all_core/defaults/qr.rs

tensor4all_core/defaults/
qr.rs
