SpectraRust/src/math/lineqs.rs

//! 线性方程组求解。
//!
//! 重构自 TLUSTY `lineqs.f`。
//!
//! 使用高斯消元法（带部分选主元）求解线性方程组 A*X = B。

// ============================================================================
// LINEQS - 高斯消元法求解线性方程组
// ============================================================================

/// 使用高斯消元法（带部分选主元）求解线性方程组 A*X = B。
///
/// 这是简化版本，假设物理维度等于逻辑维度。
/// 对于物理维度大于逻辑维度的情况，使用 `lineqs_nr`。
///
/// # 参数
///
/// - `a` - 系数矩阵，大小为 n×n（列优先存储，会被修改）
/// - `b` - 右端向量，大小为 n（会被修改）
/// - `x` - 解向量，大小为 n（输出）
/// - `n` - 实际方程数
pub fn lineqs(a: &mut [f64], b: &mut [f64], x: &mut [f64], n: usize) {
    lineqs_nr(a, b, x, n, n);
}

/// 使用高斯消元法（带部分选主元）求解线性方程组 A*X = B。
///
/// 支持物理维度大于逻辑维度的情况。
///
/// # 参数
///
/// - `a` - 系数矩阵，物理大小为 nr×nr（列优先存储，会被修改）
/// - `b` - 右端向量，物理大小为 nr（会被修改）
/// - `x` - 解向量，物理大小为 nr（输出）
/// - `n` - 实际方程数（逻辑维度）
/// - `nr` - 物理维度
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE LINEQS(A,B,X,N,NR)
///   DIMENSION A(NR,NR),B(NR),X(NR),D(MLEVEL),IP(MLEVEL)
///   ...
/// END
/// ```
///
/// # 注意
///
/// - 矩阵 A 和向量 B 在求解过程中会被修改
pub fn lineqs_nr(a: &mut [f64], b: &mut [f64], x: &mut [f64], n: usize, nr: usize) {
    // 特殊情况：2×2 系统，直接求解
    if n == 2 {
        let a11 = a[0];
        let a12 = a[nr]; // Fortran 列优先: A(1,2) = A[0 + nr*1]
        let a21 = a[1]; // Fortran 列优先: A(2,1) = A[1 + nr*0]
        let a22 = a[nr + 1];

        let det = a11 * a22 - a12 * a21;
        x[0] = (a22 * b[0] - a12 * b[1]) / det;
        x[1] = (b[1] - a21 * x[0]) / a22;
        return;
    }

    // 工作数组
    let mut d = vec![0.0; n];
    let mut ip = vec![0usize; n];

    // LU 分解（带部分选主元）
    for i in 0..n {
        // 复制第 i 列到 d
        for j in 0..n {
            d[j] = a[j + i * nr];
        }

        // 前向消元
        if i >= 1 {
            for j in 0..i {
                let it = ip[j];
                a[j + i * nr] = d[it];
                d[it] = d[j];

                for k in (j + 1)..n {
                    d[k] = d[k] - a[k + j * nr] * a[j + i * nr];
                }
            }
        }

        // 选主元
        let mut am = d[i].abs();
        ip[i] = i;
        for k in i..n {
            if am < d[k].abs() {
                ip[i] = k;
                am = d[k].abs();
            }
        }

        // 交换行
        let it = ip[i];
        a[i + i * nr] = d[it];
        d[it] = d[i];

        // 计算乘数
        if i + 1 < n {
            for k in (i + 1)..n {
                a[k + i * nr] = d[k] / a[i + i * nr];
            }
        }
    }

    // 前向替换（处理右端向量）
    for i in 0..n {
        let it = ip[i];
        x[i] = b[it];
        b[it] = b[i];

        if i + 1 < n {
            for j in (i + 1)..n {
                b[j] = b[j] - a[j + i * nr] * x[i];
            }
        }
    }

    // 后向替换
    for i in 0..n {
        let k = n - 1 - i;
        let mut sum = 0.0;

        if k + 1 < n {
            for j in (k + 1)..n {
                sum = sum + a[k + j * nr] * x[j];
            }
        }

        x[k] = (x[k] - sum) / a[k + k * nr];
    }
}

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

    #[test]
    fn test_lineqs_2x2() {
        // 2×2 系统
        // [2, 1] [x1]   [5]
        // [1, 3] [x2] = [7]
        // 解：x1 = 1.6, x2 = 1.8
        let mut a = vec![2.0, 1.0, 1.0, 3.0]; // 列优先
        let mut b = vec![5.0, 7.0];
        let mut x = vec![0.0; 2];

        lineqs(&mut a, &mut b, &mut x, 2);

        assert!((x[0] - 1.6).abs() < 1e-10);
        assert!((x[1] - 1.8).abs() < 1e-10);
    }

    #[test]
    fn test_lineqs_3x3() {
        // 3×3 系统
        // [1, 2, 3] [x1]   [6]
        // [4, 5, 6] [x2] = [15]
        // [7, 8, 10] [x3]  [25]
        // 解：x1 = 1, x2 = 1, x3 = 1
        // Fortran 列优先存储
        let mut a = vec![1.0, 4.0, 7.0, 2.0, 5.0, 8.0, 3.0, 6.0, 10.0];
        let mut b = vec![6.0, 15.0, 25.0];
        let mut x = vec![0.0; 3];

        lineqs(&mut a, &mut b, &mut x, 3);

        assert!((x[0] - 1.0).abs() < 1e-10);
        assert!((x[1] - 1.0).abs() < 1e-10);
        assert!((x[2] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_lineqs_identity() {
        // 单位矩阵
        let mut a = vec![1.0, 0.0, 0.0, 1.0];
        let mut b = vec![3.0, 4.0];
        let mut x = vec![0.0; 2];

        lineqs(&mut a, &mut b, &mut x, 2);

        assert!((x[0] - 3.0).abs() < 1e-10);
        assert!((x[1] - 4.0).abs() < 1e-10);
    }

    #[test]
    fn test_lineqs_diagonal() {
        // 对角矩阵
        let mut a = vec![2.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 4.0];
        let mut b = vec![6.0, 9.0, 16.0];
        let mut x = vec![0.0; 3];

        lineqs(&mut a, &mut b, &mut x, 3);

        assert!((x[0] - 3.0).abs() < 1e-10);
        assert!((x[1] - 3.0).abs() < 1e-10);
        assert!((x[2] - 4.0).abs() < 1e-10);
    }

    #[test]
    fn test_lineqs_pivoting() {
        // 需要选主元的系统
        // [0.001, 1] [x1]   [1]
        // [1, 1]     [x2] = [2]
        // 解：x1 ≈ 1.001, x2 ≈ 0.999
        let mut a = vec![0.001, 1.0, 1.0, 1.0];
        let mut b = vec![1.0, 2.0];
        let mut x = vec![0.0; 2];

        lineqs(&mut a, &mut b, &mut x, 2);

        // 直接验证 A*X = B
        let res1 = 0.001 * x[0] + 1.0 * x[1];
        let res2 = 1.0 * x[0] + 1.0 * x[1];
        assert!((res1 - 1.0).abs() < 1e-10, "First equation: {} != 1", res1);
        assert!((res2 - 2.0).abs() < 1e-10, "Second equation: {} != 2", res2);
    }
}