sparse_ir_basis/sve/

compute.rs

//! Main SVE computation functions

use crate::error::{Error, require_accuracy, require_nonzero_size};
use crate::fpu_check::FpuGuard;
use crate::kernel::{AbstractKernel, CentrosymmKernel, KernelProperties, SVEHints};
use crate::matrix::Mat;
use crate::numeric::CustomNumeric;

use super::result::SVEResult;
use super::strategy::{CentrosymmSVE, NonCentrosymmSVE, SVEStrategy};
use super::types::{TworkType, safe_epsilon};

/// Default relative cutoff for singular value truncation: `2 * T::epsilon()`
///
/// This is `2^-51` (about 4.44e-16) for `f64` and `2^-104` (about 4.93e-32)
/// for `Df64`. Both SVE paths use it when `cutoff` is `None`.
/// Convention-matched with libsparseir, whose `pre_postprocess` in
/// `backend/cxx/include/sparseir/impl/sve_impl.ipp` (commit 4bc58ea) uses
/// `T(2) * std::numeric_limits<T>::epsilon()` when the cutoff is NaN.
fn default_cutoff<T: CustomNumeric>() -> T {
    T::from_f64_unchecked(2.0) * T::epsilon()
}

/// Release unused memory back to the OS.
///
/// SVE computation allocates large temporary buffers for SVD.
/// After computation, these are freed but the allocator may retain them.
/// This function asks the allocator to return unused memory to the OS.
///
/// Supported platforms:
/// - macOS: uses `malloc_zone_pressure_relief`
/// - Linux (glibc): uses `malloc_trim`
/// - Other platforms: no-op (memory is still freed, just not returned to OS immediately)
#[inline]
fn release_unused_memory() {
    #[cfg(target_os = "macos")]
    {
        unsafe extern "C" {
            fn malloc_zone_pressure_relief(zone: *mut std::ffi::c_void, goal: usize) -> usize;
        }
        unsafe { malloc_zone_pressure_relief(std::ptr::null_mut(), 0) };
    }

    // Only use malloc_trim on Linux with glibc (not musl)
    #[cfg(all(target_os = "linux", target_env = "gnu"))]
    {
        unsafe extern "C" {
            fn malloc_trim(pad: usize) -> i32;
        }
        unsafe { malloc_trim(0) };
    }
}

/// Main SVE computation function for centrosymmetric kernels
///
/// Automatically chooses the appropriate SVE strategy based on kernel properties
/// and working precision based on epsilon.
///
/// # Arguments
///
/// * `kernel` - The centrosymmetric kernel to expand
/// * `epsilon` - Required accuracy, in (0, 1). `None` selects the best
///   accuracy of the working precision (about 1.6e-16 in Float64X2).
/// * `cutoff` - Relative tolerance for singular value truncation: singular
///   values smaller than `cutoff` times the largest singular value are
///   discarded. `None` selects `2 * machine epsilon` of the working precision,
///   about 4.44e-16 for Float64 and 4.93e-32 for Float64X2 (the libsparseir
///   default). The SVD of each even/odd block already discards singular values
///   below `2 * machine epsilon` times that block's largest singular value, so
///   a smaller `cutoff` has little effect. Must be in [0, 1].
/// * `max_num_svals` - Maximum number of singular values to keep
/// * `twork` - Working precision type (Auto for automatic selection)
///
/// # Returns
///
/// SVEResult containing singular functions and values
///
/// # Errors
///
/// * [`Error::InvalidParameter`] if `epsilon` is not in (0, 1), `cutoff` is
///   not in [0, 1], or `max_num_svals` is `Some(0)`; checked before any work
/// * [`Error::NonFiniteInput`] if the discretized kernel has a NaN or
///   infinite entry (e.g. a cutoff Λ so small that 1/Λ overflows)
/// * [`Error::DecompositionFailed`] if an SVD fails
///
/// # FPU State Warning
///
/// This function checks for dangerous FPU settings (Flush-to-Zero and Denormals-Are-Zero)
/// that can cause incorrect results. If detected, it temporarily corrects the FPU state
/// and prints a warning. If you see this warning, add `-fp-model precise` flag when
/// compiling with Intel Fortran.
pub fn compute_sve<K>(
    kernel: K,
    epsilon: Option<f64>,
    cutoff: Option<f64>,
    max_num_svals: Option<usize>,
    twork: TworkType,
) -> Result<SVEResult, Error>
where
    K: CentrosymmKernel + KernelProperties + Clone + 'static,
{
    check_sve_parameters(epsilon, cutoff, max_num_svals)?;

    // Protect computation from dangerous FPU settings (FZ/DAZ)
    // This temporarily disables FZ/DAZ and restores them after computation
    let _fpu_guard = FpuGuard::new_protect_computation();

    // Determine safe epsilon and working precision
    let (safe_epsilon, twork_actual) = safe_epsilon(epsilon, twork);

    // Dispatch based on working precision
    let result = match twork_actual {
        TworkType::Float64 => {
            compute_sve_with_precision::<f64, K>(kernel, safe_epsilon, cutoff, max_num_svals)
        }
        TworkType::Float64X2 => compute_sve_with_precision::<crate::Df64, K>(
            kernel,
            safe_epsilon,
            cutoff.map(crate::Df64::from),
            max_num_svals,
        ),
        _ => panic!("Invalid TworkType: {:?}", twork_actual),
    };

    // Release temporary memory back to OS after SVE computation
    release_unused_memory();

    result
}

/// Main SVE computation function for general kernels (centrosymmetric or non-centrosymmetric)
///
/// Discretizes the kernel on its full domain `[-xmax, xmax] × [-ymax, ymax]`
/// with [`NonCentrosymmSVE`], which only requires [`AbstractKernel`].
///
/// Centrosymmetric kernels are expanded correctly as well: their half-domain
/// [`SVEHints`] segments are mirrored onto the full domain, so the singular
/// values agree with [`compute_sve`] up to rounding. The even/odd block
/// structure is not exploited, however: the SVD is taken of one matrix with
/// twice as many rows and columns as each block used by [`compute_sve`]
/// (asymptotically about four times the work), the singular functions carry no
/// parity tag, and the singular functions of (nearly) degenerate singular
/// values may mix the even and odd sectors. For kernels implementing
/// [`CentrosymmKernel`], such as [`LogisticKernel`](crate::kernel::LogisticKernel)
/// and [`RegularizedBoseKernel`](crate::kernel::RegularizedBoseKernel), prefer
/// [`compute_sve`].
///
/// This function cannot select [`CentrosymmSVE`] by itself even when
/// [`AbstractKernel::is_centrosymmetric`] returns true: that strategy needs
/// the reduced kernels of [`CentrosymmKernel::compute_reduced`], and a
/// `K: AbstractKernel` bound cannot be refined to `K: CentrosymmKernel`
/// without trait specialization, which stable Rust does not provide.
///
/// # Arguments
///
/// * `kernel` - The kernel to expand (can be centrosymmetric or non-centrosymmetric)
/// * `epsilon` - Required accuracy, in (0, 1). `None` selects the best
///   accuracy of the working precision (about 1.6e-16 in Float64X2).
/// * `cutoff` - Relative tolerance for singular value truncation: singular
///   values smaller than `cutoff` times the largest singular value are
///   discarded. `None` selects `2 * machine epsilon` of the working precision,
///   about 4.44e-16 for Float64 and 4.93e-32 for Float64X2, as in
///   [`compute_sve`]. The SVD of the full-domain matrix already discards
///   singular values below `2 * machine epsilon` times the largest one, so a
///   smaller `cutoff` has no effect. Must be in [0, 1].
/// * `max_num_svals` - Maximum number of singular values to keep
/// * `twork` - Working precision type (Auto for automatic selection)
///
/// # Returns
///
/// SVEResult containing singular functions and values
///
/// # Errors
///
/// * [`Error::InvalidParameter`] if `epsilon` is not in (0, 1), `cutoff` is
///   not in [0, 1], or `max_num_svals` is `Some(0)`; checked before any work
/// * [`Error::NonFiniteInput`] if the discretized kernel has a NaN or
///   infinite entry (e.g. a cutoff Λ so small that 1/Λ overflows)
/// * [`Error::DecompositionFailed`] if an SVD fails
///
/// # FPU State Warning
///
/// This function checks for dangerous FPU settings (Flush-to-Zero and Denormals-Are-Zero)
/// that can cause incorrect results. If detected, it temporarily corrects the FPU state
/// and prints a warning. If you see this warning, add `-fp-model precise` flag when
/// compiling with Intel Fortran.
pub fn compute_sve_general<K>(
    kernel: K,
    epsilon: Option<f64>,
    cutoff: Option<f64>,
    max_num_svals: Option<usize>,
    twork: TworkType,
) -> Result<SVEResult, Error>
where
    K: AbstractKernel + KernelProperties + Clone + 'static,
{
    check_sve_parameters(epsilon, cutoff, max_num_svals)?;

    // Protect computation from dangerous FPU settings (FZ/DAZ)
    // This temporarily disables FZ/DAZ and restores them after computation
    let _fpu_guard = FpuGuard::new_protect_computation();

    // Determine safe epsilon and working precision
    let (safe_epsilon, twork_actual) = safe_epsilon(epsilon, twork);

    // Dispatch based on working precision
    let result = match twork_actual {
        TworkType::Float64 => compute_sve_general_with_precision::<f64, K>(
            kernel,
            safe_epsilon,
            cutoff,
            max_num_svals,
        ),
        TworkType::Float64X2 => compute_sve_general_with_precision::<crate::Df64, K>(
            kernel,
            safe_epsilon,
            cutoff.map(crate::Df64::from),
            max_num_svals,
        ),
        _ => panic!("Invalid TworkType: {:?}", twork_actual),
    };

    // Release temporary memory back to OS after SVE computation
    release_unused_memory();

    result
}

/// Check the parameters shared by [`compute_sve`] and [`compute_sve_general`]
fn check_sve_parameters(
    epsilon: Option<f64>,
    cutoff: Option<f64>,
    max_num_svals: Option<usize>,
) -> Result<(), Error> {
    require_accuracy("epsilon", epsilon)?;
    if let Some(c) = cutoff {
        if !(0.0..=1.0).contains(&c) {
            return Err(Error::InvalidParameter {
                name: "cutoff",
                value: format!("{c:?}"),
                reason: "must be in [0, 1]".to_string(),
            });
        }
    }
    require_nonzero_size("max_num_svals", max_num_svals)
}

/// Compute SVE with specific precision type
fn compute_sve_with_precision<T, K>(
    kernel: K,
    epsilon: f64,
    cutoff: Option<T>,
    max_num_svals: Option<usize>,
) -> Result<SVEResult, Error>
where
    T: CustomNumeric + Send + Sync + Clone + 'static,
    K: CentrosymmKernel + KernelProperties + Clone + 'static,
    K::SVEHintsType<T>: SVEHints<T> + Clone,
{
    // 1. Determine SVE strategy (automatically chooses CentrosymmSVE for centrosymmetric kernels)
    let sve = determine_sve::<T, K>(kernel, epsilon)?;

    // 2. Compute matrices
    let matrices = sve.matrices();

    // 3. Compute SVD for each matrix
    let mut u_list = Vec::new();
    let mut s_list = Vec::new();
    let mut v_list = Vec::new();

    for matrix in matrices.iter() {
        let (u, s, v) = crate::tsvd::compute_svd_dtensor(matrix)?;
        u_list.push(u);
        s_list.push(s);
        v_list.push(v);
    }

    // 4. Truncate based on cutoff (default: 2 * T::epsilon(), see default_cutoff)
    let rtol_t = cutoff.unwrap_or_else(default_cutoff::<T>);
    let (u_trunc, s_trunc, v_trunc) = truncate(u_list, s_list, v_list, rtol_t, max_num_svals);

    // 5. Post-process to create SVEResult
    sve.postprocess(u_trunc, s_trunc, v_trunc)
}

/// Compute SVE with specific precision type for general kernels
fn compute_sve_general_with_precision<T, K>(
    kernel: K,
    epsilon: f64,
    cutoff: Option<T>,
    max_num_svals: Option<usize>,
) -> Result<SVEResult, Error>
where
    T: CustomNumeric + Send + Sync + Clone + 'static,
    K: AbstractKernel + KernelProperties + Clone + 'static,
    K::SVEHintsType<T>: SVEHints<T> + Clone,
{
    // 1. Determine SVE strategy (full-domain NonCentrosymmSVE)
    let sve = determine_sve_general::<T, K>(kernel, epsilon)?;

    // 2. Compute matrices
    let matrices = sve.matrices();

    // 3. Compute SVD for each matrix
    let mut u_list = Vec::new();
    let mut s_list = Vec::new();
    let mut v_list = Vec::new();

    for matrix in matrices.iter() {
        let (u, s, v) = crate::tsvd::compute_svd_dtensor(matrix)?;
        u_list.push(u);
        s_list.push(s);
        v_list.push(v);
    }

    // 4. Truncate based on cutoff (default: 2 * T::epsilon(), see default_cutoff)
    let rtol_t = cutoff.unwrap_or_else(default_cutoff::<T>);
    let (u_trunc, s_trunc, v_trunc) = truncate(u_list, s_list, v_list, rtol_t, max_num_svals);

    // 5. Post-process to create SVEResult
    sve.postprocess(u_trunc, s_trunc, v_trunc)
}

/// Determine the appropriate SVE strategy
///
/// For centrosymmetric kernels, uses CentrosymmSVE for efficient computation
/// by exploiting even/odd symmetry.
fn determine_sve<T, K>(kernel: K, epsilon: f64) -> Result<Box<dyn SVEStrategy<T>>, Error>
where
    T: CustomNumeric + Send + Sync + Clone + 'static,
    K: CentrosymmKernel + KernelProperties + Clone + 'static,
    K::SVEHintsType<T>: SVEHints<T> + Clone,
{
    // CentrosymmKernel trait implies centrosymmetric
    Ok(Box::new(CentrosymmSVE::new(kernel, epsilon)?))
}

/// Determine the SVE strategy for general kernels
///
/// Always uses [`NonCentrosymmSVE`], which handles centrosymmetric kernels by
/// mirroring their half-domain hint segments onto the full domain.
/// [`CentrosymmSVE`] would require `K: CentrosymmKernel`, which cannot be
/// recovered from the `K: AbstractKernel` bound (see [`compute_sve_general`]).
fn determine_sve_general<T, K>(kernel: K, epsilon: f64) -> Result<Box<dyn SVEStrategy<T>>, Error>
where
    T: CustomNumeric + Send + Sync + Clone + 'static,
    K: AbstractKernel + KernelProperties + Clone + 'static,
    K::SVEHintsType<T>: SVEHints<T> + Clone,
{
    Ok(Box::new(NonCentrosymmSVE::new(kernel, epsilon)?))
}

/// Truncate SVD results based on cutoff and maximum size
///
/// # Arguments
///
/// * `u_list` - List of U matrices
/// * `s_list` - List of singular value vectors
/// * `v_list` - List of V matrices
/// * `rtol` - Relative tolerance for truncation
/// * `max_num_svals` - Maximum number of singular values to keep
///
/// # Returns
///
/// Tuple of (truncated_u_list, truncated_s_list, truncated_v_list)
pub(crate) fn truncate<T: CustomNumeric>(
    u_list: Vec<Mat<T>>,
    s_list: Vec<Vec<T>>,
    v_list: Vec<Mat<T>>,
    rtol: T,
    max_num_svals: Option<usize>,
) -> (Vec<Mat<T>>, Vec<Vec<T>>, Vec<Mat<T>>) {
    let zero = T::zero();

    // Validate
    if let Some(max) = max_num_svals {
        if max == 0 {
            panic!("max_num_svals must be positive");
        }
    }
    if rtol < zero || rtol > T::from_f64_unchecked(1.0) {
        panic!("rtol must be in [0, 1]");
    }

    // Find global maximum singular value
    let mut all_svals = Vec::new();
    for s in &s_list {
        all_svals.extend(s.iter().copied());
    }

    let max_sval = all_svals
        .iter()
        .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .copied()
        .unwrap_or(zero);

    // Determine cutoff
    let cutoff = if let Some(max_count) = max_num_svals {
        if max_count < all_svals.len() {
            let mut sorted = all_svals.clone();
            sorted.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
            let nth = sorted[max_count - 1];
            if rtol * max_sval > nth {
                rtol * max_sval
            } else {
                nth
            }
        } else {
            rtol * max_sval
        }
    } else {
        rtol * max_sval
    };

    // Truncate each result
    let mut u_trunc = Vec::new();
    let mut s_trunc = Vec::new();
    let mut v_trunc = Vec::new();

    for i in 0..s_list.len() {
        let s = &s_list[i];
        let u = &u_list[i];
        let v = &v_list[i];

        // Count singular values above cutoff
        let mut n_keep = 0;
        for &val in s.iter() {
            if val >= cutoff {
                n_keep += 1;
            }
        }

        // Slice U: keep first n_keep columns. Preserve one output block per
        // input block even when all singular values in that block are removed.
        let u_shape = *u.shape();
        let u_sliced = Mat::<T>::from_fn([u_shape.0, n_keep], |idx| u[[idx[0], idx[1]]]);
        u_trunc.push(u_sliced);

        s_trunc.push(s[..n_keep].to_vec());

        // Slice V: keep first n_keep columns
        let v_shape = *v.shape();
        let v_sliced = Mat::<T>::from_fn([v_shape.0, n_keep], |idx| v[[idx[0], idx[1]]]);
        v_trunc.push(v_sliced);
    }

    (u_trunc, s_trunc, v_trunc)
}

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

    /// The default cutoff is exactly `2 * machine epsilon` of each working
    /// precision (issue #249); `Df64` epsilon is `f64::EPSILON^2 / 2 = 2^-105`.
    #[test]
    fn test_default_cutoff_is_two_machine_epsilon() {
        assert_eq!(default_cutoff::<f64>(), 2.0 * f64::EPSILON);
        assert_eq!(default_cutoff::<f64>(), 2f64.powi(-51));

        let df64 = default_cutoff::<crate::Df64>();
        assert_eq!(df64, crate::Df64::from(2f64.powi(-104)));
        assert_eq!((df64.hi(), df64.lo()), (2f64.powi(-104), 0.0));
    }

    #[test]
    fn test_truncate_by_rtol() {
        let u = vec![Mat::<f64>::from_elem([3, 3], 1.0)];
        let s = vec![vec![10.0, 5.0, 0.1]];
        let v = vec![Mat::<f64>::from_elem([3, 3], 1.0)];

        // rtol = 0.1, max_sval = 10.0, cutoff = 1.0
        // Keep values >= 1.0: [10.0, 5.0]
        let (_, s_trunc, _) = truncate(u, s, v, 0.1, None);

        assert_eq!(s_trunc[0].len(), 2);
        assert_eq!(s_trunc[0][0], 10.0);
        assert_eq!(s_trunc[0][1], 5.0);
    }

    #[test]
    fn test_truncate_by_max_size() {
        let u = vec![Mat::<f64>::from_elem([3, 3], 1.0)];
        let s = vec![vec![10.0, 5.0, 2.0]];
        let v = vec![Mat::<f64>::from_elem([3, 3], 1.0)];

        // max_num_svals = 2
        let (_, s_trunc, _) = truncate(u, s, v, 0.0, Some(2));

        assert_eq!(s_trunc[0].len(), 2);
    }

    #[test]
    #[should_panic(expected = "max_num_svals must be positive")]
    fn test_truncate_rejects_zero_max_size() {
        let u = vec![Mat::<f64>::from_elem([1, 1], 1.0)];
        let s = vec![vec![1.0]];
        let v = vec![Mat::<f64>::from_elem([1, 1], 1.0)];

        truncate(u, s, v, 0.0, Some(0));
    }

    #[test]
    fn test_truncate_preserves_empty_blocks() {
        let u = vec![
            Mat::<f64>::from_elem([2, 1], 1.0),
            Mat::<f64>::from_elem([2, 1], 2.0),
        ];
        let s = vec![vec![10.0], vec![1.0]];
        let v = vec![
            Mat::<f64>::from_elem([2, 1], 1.0),
            Mat::<f64>::from_elem([2, 1], 2.0),
        ];

        let (u_trunc, s_trunc, v_trunc) = truncate(u, s, v, 0.5, None);

        assert_eq!(s_trunc.len(), 2);
        assert_eq!(s_trunc[0], vec![10.0]);
        assert!(s_trunc[1].is_empty());
        assert_eq!(*u_trunc[1].shape(), (2, 0));
        assert_eq!(*v_trunc[1].shape(), (2, 0));
    }

    #[test]
    #[should_panic(expected = "rtol must be in [0, 1]")]
    fn test_truncate_invalid_rtol() {
        let u = vec![Mat::<f64>::from_elem([1, 1], 1.0)];
        let s = vec![vec![1.0]];
        let v = vec![Mat::<f64>::from_elem([1, 1], 1.0)];

        truncate(u, s, v, 1.5, None);
    }
}