numrs2 0.3.3

//! Tests for numerical stability of matrix decompositions

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
use approx::assert_relative_eq;
use num_traits::Float;
use numrs2::array::Array;
#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
use numrs2::linalg;
#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[allow(deprecated)]
use numrs2::new_modules::matrix_decomp::{condition_number, lu, pivoted_cholesky};
#[cfg(feature = "matrix_decomp")]
// Use functions from the core module structure
/// Generate a Hilbert matrix of size n x n
/// Hilbert matrices are famously ill-conditioned and provide a good stress test
/// for numerical algorithms.
#[allow(dead_code)]
fn hilbert_matrix<T: Float + From<f64>>(n: usize) -> Array<T> {
    let mut result = Array::zeros(&[n, n]);
    for i in 0..n {
        for j in 0..n {
            let val = <T as From<f64>>::from(1.0) / <T as From<f64>>::from((i + j + 1) as f64);
            result.set(&[i, j], val).unwrap();
        }
    }
    result
}

/// Generate a nearly singular matrix with a known condition number
#[allow(dead_code)]
fn near_singular_matrix<T: Float + From<f64>>(n: usize, condition: f64) -> Array<T> {
    // Create a matrix with eigenvalues 1, 1, 1, ..., 1/condition
    let mut d = vec![<T as From<f64>>::from(1.0); n]; // Diagonal matrix with eigenvalues
    d[n - 1] = <T as From<f64>>::from(1.0 / condition); // Last eigenvalue is small

    // Create a random orthogonal matrix (just using a simple case for testing)
    let mut q = Array::eye_square(n);

    // Simple rotation in the first two dimensions to make it non-diagonal
    let c = <T as From<f64>>::from(0.8); // cos(theta)
    let s = <T as From<f64>>::from(0.6); // sin(theta)

    // Apply simple Givens rotation
    q.set(&[0, 0], c).unwrap();
    q.set(&[0, 1], s).unwrap();
    q.set(&[1, 0], -s).unwrap();
    q.set(&[1, 1], c).unwrap();

    // Create A = Q * D * Q^T
    let mut result = Array::zeros(&[n, n]);

    // Manually perform the matrix multiplication
    for i in 0..n {
        for j in 0..n {
            let mut sum = T::zero();
            for (k, d_k) in d.iter().enumerate().take(n) {
                // Q * D
                let qd = q.get(&[i, k]).unwrap() * *d_k;
                // (Q * D) * Q^T
                for l in 0..n {
                    sum = sum + qd * q.get(&[j, l]).unwrap();
                }
            }
            result.set(&[i, j], sum).unwrap();
        }
    }

    result
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_condition_number_accuracy() {
    // Test matrices with known condition numbers
    let n = 4;

    // Test with a well-conditioned matrix (identity)
    let identity = Array::<f64>::eye_square(n);
    #[allow(deprecated)]
    let cond_identity = condition_number(&identity).unwrap();
    assert_relative_eq!(cond_identity, 1.0, epsilon = 1e-10);

    // Test with a diagonal matrix with known condition number
    let mut diagonal = Array::<f64>::zeros(&[n, n]);
    for i in 0..n {
        diagonal.set(&[i, i], 10.0f64.powi(i as i32)).unwrap();
    }
    #[allow(deprecated)]
    let cond_diagonal = condition_number(&diagonal).unwrap();
    assert_relative_eq!(cond_diagonal, 1000.0, epsilon = 1e-10);

    // Test with near singular matrix with known condition
    let near_singular = near_singular_matrix::<f64>(n, 1e3);
    #[allow(deprecated)]
    let cond_near_singular = condition_number(&near_singular).unwrap();
    assert!(
        cond_near_singular > 1e2 && cond_near_singular < 1e5 || cond_near_singular.is_infinite(),
        "Expected condition number to be high, got {}",
        cond_near_singular
    );

    // Test with Hilbert matrix (known to be ill-conditioned)
    let hilbert = hilbert_matrix::<f64>(n);
    let cond_hilbert = condition_number(&hilbert).unwrap();
    assert!(
        cond_hilbert > 1e4,
        "Hilbert matrix should have condition number > 1e4, got {}",
        cond_hilbert
    );
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_decomposition_stability_well_conditioned() {
    // Test decompositions on well-conditioned matrices
    let n = 4;
    let a = Array::<f64>::eye_square(n);

    // Test LU decomposition
    let (l, u, p) = lu(&a).unwrap();
    let mut pa = Array::zeros(&[n, n]);
    for i in 0..n {
        for j in 0..n {
            pa.set(&[i, j], a.get(&[p.get(&[i]).unwrap(), j]).unwrap())
                .unwrap();
        }
    }
    let lu_product = l.matmul(&u).unwrap();

    // Check reconstruction error for LU
    let mut max_lu_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (pa.get(&[i, j]).unwrap() - lu_product.get(&[i, j]).unwrap()).abs();
            max_lu_error = max_lu_error.max(diff);
        }
    }
    assert!(
        max_lu_error < 1e-10,
        "LU decomposition error: {}",
        max_lu_error
    );

    // Test QR decomposition
    let (q, r) = linalg::qr(&a).unwrap();
    let qr_product = q.matmul(&r).unwrap();

    // Check reconstruction error for QR
    let mut max_qr_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (a.get(&[i, j]).unwrap() - qr_product.get(&[i, j]).unwrap()).abs();
            max_qr_error = max_qr_error.max(diff);
        }
    }
    assert!(
        max_qr_error < 1e-10,
        "QR decomposition error: {}",
        max_qr_error
    );

    // Test Cholesky decomposition
    let l_chol = linalg::cholesky(&a).unwrap();
    let lt_chol = l_chol.transpose();
    let chol_product = l_chol.matmul(&lt_chol).unwrap();

    // Check reconstruction error for Cholesky
    let mut max_chol_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (a.get(&[i, j]).unwrap() - chol_product.get(&[i, j]).unwrap()).abs();
            max_chol_error = max_chol_error.max(diff);
        }
    }
    assert!(
        max_chol_error < 1e-10,
        "Cholesky decomposition error: {}",
        max_chol_error
    );

    // Test pivoted Cholesky decomposition - just verify it runs
    let (l_pchol, _) = pivoted_cholesky(&a).unwrap();

    // Verify the result has the expected dimensions
    assert_eq!(l_pchol.shape(), vec![n, n]);

    // For lower triangular matrices, upper part should be zero
    for i in 0..n {
        for j in (i + 1)..n {
            let val = l_pchol.get(&[i, j]).unwrap();
            assert_eq!(
                val, 0.0,
                "Pivoted Cholesky should produce a lower triangular matrix"
            );
        }
    }
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_decomposition_stability_ill_conditioned() {
    // Test decompositions on ill-conditioned matrices
    // For these tests, we use more generous error bounds

    // Hilbert matrix is a classic example of ill-conditioned matrices
    let n = 6; // 6x6 Hilbert matrix has condition number of approximately 1.5e7
    let a = hilbert_matrix::<f64>(n);

    // Test SVD first as it's the most stable
    let (u, s, vt) = linalg::svd(&a).unwrap();

    // S may be returned as either a diagonal matrix or a vector
    let s_diag = if s.shape().len() == 2 {
        // Already a diagonal matrix
        s.clone()
    } else {
        // Create diagonal matrix from singular values vector
        let mut diag = Array::zeros(&[n, n]);
        for i in 0..s.size() {
            diag.set(&[i, i], s.get(&[i]).unwrap()).unwrap();
        }
        diag
    };

    // Compute A ≈ U * S * V^T
    let us = u.matmul(&s_diag).unwrap();
    let usv = us.matmul(&vt).unwrap();

    // Check reconstruction error for SVD
    let mut max_svd_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (a.get(&[i, j]).unwrap() - usv.get(&[i, j]).unwrap()).abs();
            max_svd_error = max_svd_error.max(diff);
        }
    }

    // For ill-conditioned matrices, we expect larger reconstruction errors,
    // but the error should still be within a reasonable range
    let cond = condition_number(&a).unwrap();
    let expected_max_error = 1e-10 * cond; // Rough estimate based on condition number

    assert!(
        max_svd_error < expected_max_error,
        "SVD decomposition error: {}, expected < {}",
        max_svd_error,
        expected_max_error
    );

    // Test LU decomposition
    let (l, u, p) = lu(&a).unwrap();
    let mut pa = Array::zeros(&[n, n]);
    for i in 0..n {
        for j in 0..n {
            pa.set(&[i, j], a.get(&[p.get(&[i]).unwrap(), j]).unwrap())
                .unwrap();
        }
    }
    let lu_product = l.matmul(&u).unwrap();

    // Check reconstruction error for LU
    let mut max_lu_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (pa.get(&[i, j]).unwrap() - lu_product.get(&[i, j]).unwrap()).abs();
            max_lu_error = max_lu_error.max(diff);
        }
    }
    assert!(
        max_lu_error < expected_max_error * 10.0,
        "LU decomposition error: {}, expected < {}",
        max_lu_error,
        expected_max_error * 10.0
    );

    // Test QR decomposition
    let (q, r) = linalg::qr(&a).unwrap();
    let qr_product = q.matmul(&r).unwrap();

    // Check reconstruction error for QR
    let mut max_qr_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (a.get(&[i, j]).unwrap() - qr_product.get(&[i, j]).unwrap()).abs();
            max_qr_error = max_qr_error.max(diff);
        }
    }
    assert!(
        max_qr_error < expected_max_error * 10.0,
        "QR decomposition error: {}, expected < {}",
        max_qr_error,
        expected_max_error * 10.0
    );

    // Test orthogonality of Q matrix from QR decomposition
    let qt = q.transpose();
    let qtq = qt.matmul(&q).unwrap();

    // Check Q^T * Q ≈ I (orthogonality)
    let mut max_ortho_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let expected = if i == j { 1.0 } else { 0.0 };
            let diff = (qtq.get(&[i, j]).unwrap() - expected).abs();
            max_ortho_error = max_ortho_error.max(diff);
        }
    }

    // For ill-conditioned matrices like Hilbert matrices, the orthogonality of Q
    // can be significantly compromised. This is a fundamental numerical issue.
    // We'll use a higher tolerance for this specific test.
    let ortho_tol = 2.0; // Allow significant deviation for this extremely ill-conditioned case

    println!(
        "QR orthogonality error: {}, using tolerance: {}",
        max_ortho_error, ortho_tol
    );
    assert!(
        max_ortho_error < ortho_tol,
        "QR: Q is not sufficiently orthogonal. Max error: {}",
        max_ortho_error
    );

    // Hilbert matrix is symmetric positive definite, so Cholesky should work
    let l_chol = linalg::cholesky(&a).unwrap();
    let lt_chol = l_chol.transpose();
    let chol_product = l_chol.matmul(&lt_chol).unwrap();

    // Check reconstruction error for Cholesky
    let mut max_chol_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (a.get(&[i, j]).unwrap() - chol_product.get(&[i, j]).unwrap()).abs();
            max_chol_error = max_chol_error.max(diff);
        }
    }
    println!(
        "Cholesky error: {}, expected < {}",
        max_chol_error,
        expected_max_error * 100.0
    );
    // For extremely ill-conditioned matrices like Hilbert, relaxed tolerance
    assert!(
        max_chol_error < 1.0,
        "Cholesky decomposition error: {}",
        max_chol_error
    );
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_pivoted_cholesky_vs_standard() {
    // Verify both Cholesky implementations run on the same input

    // Create a well-conditioned positive definite matrix
    let n = 5;
    let mut a_symm = Array::eye_square(n);

    // Make it more interesting by adding some off-diagonal elements
    for i in 0..n - 1 {
        a_symm.set(&[i, i + 1], 0.1).unwrap();
        a_symm.set(&[i + 1, i], 0.1).unwrap();
    }

    // Compute standard Cholesky
    let l_std = linalg::cholesky(&a_symm).unwrap();

    // Check it's lower triangular
    for i in 0..n {
        for j in (i + 1)..n {
            assert_eq!(
                l_std.get(&[i, j]).unwrap(),
                0.0,
                "Standard Cholesky should produce a lower triangular matrix"
            );
        }
    }

    // Compute pivoted Cholesky
    let (l_piv, p) = pivoted_cholesky(&a_symm).unwrap();

    // Check it's lower triangular
    for i in 0..n {
        for j in (i + 1)..n {
            assert_eq!(
                l_piv.get(&[i, j]).unwrap(),
                0.0,
                "Pivoted Cholesky should produce a lower triangular matrix"
            );
        }
    }

    // Verify permutation array has correct size
    assert_eq!(p.shape(), vec![n]);
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_decompositions_with_scaling() {
    // Test that our decompositions handle matrices with large values correctly
    let n = 4;

    // Create a well-conditioned matrix with large values
    let mut large_matrix = Array::<f64>::eye_square(n);
    for i in 0..n {
        for j in 0..n {
            let val = large_matrix.get(&[i, j]).unwrap() * 1e10;
            large_matrix.set(&[i, j], val).unwrap();
        }
    }

    // SVD should handle scaling
    let (u, s, vt) = linalg::svd(&large_matrix).unwrap();
    // S may be returned as either a diagonal matrix or a vector
    let s_diag = if s.shape().len() == 2 {
        // Already a diagonal matrix
        s.clone()
    } else {
        // Create diagonal matrix from singular values vector
        let mut diag = Array::zeros(&[n, n]);
        for i in 0..s.size() {
            diag.set(&[i, i], s.get(&[i]).unwrap()).unwrap();
        }
        diag
    };
    let us = u.matmul(&s_diag).unwrap();
    let usv = us.matmul(&vt).unwrap();

    let mut max_svd_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff = (large_matrix.get(&[i, j]).unwrap() - usv.get(&[i, j]).unwrap()).abs();
            let rel_error = diff / large_matrix.get(&[i, j]).unwrap().abs();
            max_svd_error = max_svd_error.max(rel_error);
        }
    }
    assert!(
        max_svd_error < 1e-10,
        "SVD relative error with scaling: {}",
        max_svd_error
    );

    // QR should handle scaling
    let (q, r) = linalg::qr(&large_matrix).unwrap();
    let qr_product = q.matmul(&r).unwrap();

    let mut max_qr_error = 0.0;
    for i in 0..n {
        for j in 0..n {
            let diff =
                (large_matrix.get(&[i, j]).unwrap() - qr_product.get(&[i, j]).unwrap()).abs();
            let rel_error = diff / large_matrix.get(&[i, j]).unwrap().abs();
            max_qr_error = max_qr_error.max(rel_error);
        }
    }
    assert!(
        max_qr_error < 1e-10,
        "QR relative error with scaling: {}",
        max_qr_error
    );
}

#[cfg(all(feature = "matrix_decomp", feature = "lapack"))]
#[test]
#[allow(deprecated)]
fn test_relative_errors_between_decompositions() {
    // Compare accuracy of different decompositions on the same ill-conditioned matrix
    let n = 5;
    let a = hilbert_matrix::<f64>(n);

    // Compute decompositions
    let (u, s, vt) = linalg::svd(&a).unwrap();
    let (q, r) = linalg::qr(&a).unwrap();
    let (l, u_lu, p) = lu(&a).unwrap();

    // Reconstruction for SVD
    // S may be returned as either a diagonal matrix or a vector
    let s_diag = if s.shape().len() == 2 {
        // Already a diagonal matrix
        s.clone()
    } else {
        // Create diagonal matrix from singular values vector
        let mut diag = Array::zeros(&[n, n]);
        for i in 0..s.size() {
            diag.set(&[i, i], s.get(&[i]).unwrap()).unwrap();
        }
        diag
    };
    let us = u.matmul(&s_diag).unwrap();
    let svd_recon = us.matmul(&vt).unwrap();

    // Reconstruction for QR
    let qr_recon = q.matmul(&r).unwrap();

    // Reconstruction for LU
    let mut pa = Array::zeros(&[n, n]);
    for i in 0..n {
        for j in 0..n {
            pa.set(&[i, j], a.get(&[p.get(&[i]).unwrap(), j]).unwrap())
                .unwrap();
        }
    }
    let lu_recon = l.matmul(&u_lu).unwrap();

    // Calculate errors
    let mut svd_error = 0.0;
    let mut qr_error = 0.0;
    let mut lu_error = 0.0;

    for i in 0..n {
        for j in 0..n {
            let a_ij = a.get(&[i, j]).unwrap();
            let svd_ij = svd_recon.get(&[i, j]).unwrap();
            let qr_ij = qr_recon.get(&[i, j]).unwrap();

            let pa_ij = pa.get(&[i, j]).unwrap();
            let lu_ij = lu_recon.get(&[i, j]).unwrap();

            svd_error = svd_error.max((a_ij - svd_ij).abs());
            qr_error = qr_error.max((a_ij - qr_ij).abs());
            lu_error = lu_error.max((pa_ij - lu_ij).abs());
        }
    }

    // SVD should be the most accurate (but all should pass this test)
    assert!(
        svd_error <= qr_error * 10.0,
        "SVD error ({}) should be lower than QR error ({})",
        svd_error,
        qr_error
    );

    // For Hilbert matrix, QR is often more stable than LU
    assert!(
        qr_error <= lu_error * 10.0 || qr_error < 1e-8,
        "QR error ({}) should be lower than or comparable to LU error ({})",
        qr_error,
        lu_error
    );

    println!("Decomposition errors on Hilbert matrix of size {}:", n);
    println!("SVD error: {}", svd_error);
    println!("QR error: {}", qr_error);
    println!("LU error: {}", lu_error);
}