rustytorch/crates/specialized/rtx-fea/tests/sparse_matrix_methods_tdd.rs

//! TDD Tests for SparseMatrix additional methods
//! Following strict Red-Green-Refactor cycle
//! No mocks, stubs, or TODOs - only full implementations

#[cfg(test)]
mod sparse_matrix_methods_tests {
    use nalgebra::{DMatrix, DVector};
    use rtx_fea::assembly::sparse_matrix::SparseMatrix;
    use rtx_fea::error::FeaResult;

    #[test]
    fn test_transpose_multiply_vector() {
        // RED: Test transpose multiplication A^T * x
        let mut matrix = SparseMatrix::new(3, 2);

        // Build a 3x2 matrix
        // [1, 2]
        // [3, 4]
        // [5, 6]
        matrix.add_entry(0, 0, 1.0).unwrap();
        matrix.add_entry(0, 1, 2.0).unwrap();
        matrix.add_entry(1, 0, 3.0).unwrap();
        matrix.add_entry(1, 1, 4.0).unwrap();
        matrix.add_entry(2, 0, 5.0).unwrap();
        matrix.add_entry(2, 1, 6.0).unwrap();
        matrix.finalize().unwrap();

        let x = DVector::from_vec(vec![1.0, 2.0, 3.0]);

        // GREEN: A^T * x should produce a 2-element vector
        // A^T = [1, 3, 5]  * [1]   = [1+6+15]  = [22]
        //       [2, 4, 6]    [2]     [2+8+18]    [28]
        //                    [3]
        let result = matrix.transpose_multiply_vector(&x).unwrap();
        assert_eq!(result.len(), 2);
        assert!((result[0] - 22.0).abs() < 1e-10);
        assert!((result[1] - 28.0).abs() < 1e-10);
    }

    #[test]
    fn test_diagonal_extraction() {
        // RED: Test extracting diagonal from matrix
        let mut matrix = SparseMatrix::new(3, 3);

        // Build a 3x3 matrix
        // [2, 1, 0]
        // [1, 3, 1]
        // [0, 1, 4]
        matrix.add_entry(0, 0, 2.0).unwrap();
        matrix.add_entry(0, 1, 1.0).unwrap();
        matrix.add_entry(1, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 3.0).unwrap();
        matrix.add_entry(1, 2, 1.0).unwrap();
        matrix.add_entry(2, 1, 1.0).unwrap();
        matrix.add_entry(2, 2, 4.0).unwrap();
        matrix.finalize().unwrap();

        // GREEN: Diagonal should be [2, 3, 4]
        let diagonal = matrix.diagonal().unwrap();
        assert_eq!(diagonal.len(), 3);
        assert!((diagonal[0] - 2.0).abs() < 1e-10);
        assert!((diagonal[1] - 3.0).abs() < 1e-10);
        assert!((diagonal[2] - 4.0).abs() < 1e-10);
    }

    #[test]
    fn test_forward_solve() {
        // RED: Test forward substitution L * x = b
        let mut matrix = SparseMatrix::new(3, 3);

        // Build lower triangular matrix L
        // [2, 0, 0]
        // [1, 3, 0]
        // [2, 1, 4]
        matrix.add_entry(0, 0, 2.0).unwrap();
        matrix.add_entry(1, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 3.0).unwrap();
        matrix.add_entry(2, 0, 2.0).unwrap();
        matrix.add_entry(2, 1, 1.0).unwrap();
        matrix.add_entry(2, 2, 4.0).unwrap();
        matrix.finalize().unwrap();

        let b = DVector::from_vec(vec![4.0, 7.0, 16.0]);

        // GREEN: Solve L * x = b
        // x[0] = 4.0 / 2.0 = 2.0
        // x[1] = (7.0 - 1.0*2.0) / 3.0 = 5.0/3.0
        // x[2] = (16.0 - 2.0*2.0 - 1.0*5.0/3.0) / 4.0 = (16.0 - 4.0 - 5.0/3.0) / 4.0
        let x = matrix.forward_solve(&b).unwrap();
        assert_eq!(x.len(), 3);
        assert!((x[0] - 2.0).abs() < 1e-10);
        assert!((x[1] - 5.0 / 3.0).abs() < 1e-10);
        // Verify full solution
        let result = matrix.multiply_vector(&x).unwrap();
        for i in 0..3 {
            assert!((result[i] - b[i]).abs() < 1e-10);
        }
    }

    #[test]
    fn test_backward_solve() {
        // RED: Test backward substitution U * x = b
        let mut matrix = SparseMatrix::new(3, 3);

        // Build upper triangular matrix U
        // [2, 1, 2]
        // [0, 3, 1]
        // [0, 0, 4]
        matrix.add_entry(0, 0, 2.0).unwrap();
        matrix.add_entry(0, 1, 1.0).unwrap();
        matrix.add_entry(0, 2, 2.0).unwrap();
        matrix.add_entry(1, 1, 3.0).unwrap();
        matrix.add_entry(1, 2, 1.0).unwrap();
        matrix.add_entry(2, 2, 4.0).unwrap();
        matrix.finalize().unwrap();

        let b = DVector::from_vec(vec![10.0, 7.0, 8.0]);

        // GREEN: Solve U * x = b using backward substitution
        // x[2] = 8.0 / 4.0 = 2.0
        // x[1] = (7.0 - 1.0*2.0) / 3.0 = 5.0/3.0
        // x[0] = (10.0 - 1.0*5.0/3.0 - 2.0*2.0) / 2.0
        let x = matrix.backward_solve(&b).unwrap();
        assert_eq!(x.len(), 3);
        // Verify full solution
        let result = matrix.multiply_vector(&x).unwrap();
        for i in 0..3 {
            assert!((result[i] - b[i]).abs() < 1e-10);
        }
    }

    #[test]
    fn test_solve_vector() {
        // RED: Test general solve A * x = b
        let mut matrix = SparseMatrix::new(3, 3);

        // Build a symmetric positive definite matrix
        // [4, 1, 0]
        // [1, 3, 1]
        // [0, 1, 2]
        matrix.add_entry(0, 0, 4.0).unwrap();
        matrix.add_entry(0, 1, 1.0).unwrap();
        matrix.add_entry(1, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 3.0).unwrap();
        matrix.add_entry(1, 2, 1.0).unwrap();
        matrix.add_entry(2, 1, 1.0).unwrap();
        matrix.add_entry(2, 2, 2.0).unwrap();
        matrix.finalize().unwrap();

        let b = DVector::from_vec(vec![5.0, 6.0, 4.0]);

        // GREEN: Solve A * x = b
        let x = matrix.solve_vector(&b).unwrap();
        assert_eq!(x.len(), 3);

        // Verify solution A * x = b
        let result = matrix.multiply_vector(&x).unwrap();
        for i in 0..3 {
            assert!((result[i] - b[i]).abs() < 1e-9);
        }
    }
}