eaebec39创建于 2024年10月24日历史提交
package scientific.linear

import std.math.*
import std.unittest.*
import std.unittest.testmacro.*

import scientific.numbers.*

/* Compute the LU factorization of a general m-by-n matrix. */
foreign func LAPACKE_dgetrf(
    matrix_layout: IntNative,  // row or column major
    m: IntNative,              // number of rows in A
    n: IntNative,              // number of columns in A
    a: CPointer<Unit>,         // matrix A (overwritten, Float64)
    lda: IntNative,            // leading dimension of A
    ipiv: CPointer<Unit>      // pivot indices (Int32)
): Int64

public func lu(a: Matrix<Float64>): (Matrix<Float64>, Matrix<Float64>, Matrix<Float64>) {
    if (!a.isSquare()) {
        throw IllegalArgumentException("lu: input must be square matrix")
    }

    let LAPACK_COL_MAJOR: Int64 = 102
    let m = a.getRows()
    let n = a.getCols()
    let lda: Int64 = m
    let minmn = min(m,n)
    var ipiv = zeros<Int32>(minmn)
    let info = unsafe { LAPACKE_dgetrf(IntNative(LAPACK_COL_MAJOR), IntNative(m), IntNative(n), 
                        a.ptr, IntNative(lda), ipiv.ptr) }
    var l = a.copy()
    var u = a.copy()
    for (i in 0..m) {
        for (j in 0..n) {
            if (i <= j) {
                l[i,j] = 0.0
            }
            if (i == j) {
                l[i,j] = 1.0
            }
            if (i > j) {
                u[i,j] = 0.0
            }
        } 
    }
    var arr = vector<Int64>(minmn, {i => i + 1})
    for (i in 0..minmn) {
        var tmp = arr[i]
        arr[i] = arr[Int64(ipiv[i])-1]
        arr[Int64(ipiv[i])-1] = tmp
    }
    var p = zeros<Float64>(minmn, minmn)
    for (i in 0..minmn) {
        p[arr[i]-1, i] = 1.0
    }
    return (p, l, u)
}

public func luFactor(a: Matrix<Float64>): (Matrix<Float64>, Vector<Int32>) {
    if (!a.isSquare()) {
        throw IllegalArgumentException("luFactor: input must be square matrix")
    }

    let LAPACK_COL_MAJOR: Int64 = 102
    let m = a.getRows()
    let n = a.getCols()
    var acopy = a.copy()
    let lda: Int64 = m
    let minmn = min(m,n)
    var ipiv = zeros<Int32>(minmn)
    let info = unsafe { LAPACKE_dgetrf(IntNative(LAPACK_COL_MAJOR), IntNative(m), IntNative(n), 
                        acopy.ptr, IntNative(lda), ipiv.ptr) }
    return (acopy, ipiv)
}

/* Solving system of linear equations using LU-factored matrix. */
foreign func LAPACKE_dgetrs(
    matrix_layout: IntNative,  // row or column major
    trans: UInt8,              // whether to take transpose (hermitian transpose)
    n: IntNative,              // size of A
    nrhs: IntNative,           // number of right-hand sides
    a: CPointer<Unit>,         // matrix A (Float64)
    lda: IntNative,            // leading dimension of A
    ipiv: CPointer<Unit>,      // ipiv as returned by ?getrf (Int32)
    b: CPointer<Unit>,         // right-side of the equation B (Float64)
    ldb: IntNative             // leading dimension of B
): Int64

public func luSolve(lu: Matrix<Float64>, piv: Vector<Int32>, b: Matrix<Float64>): Matrix<Float64> {
    if (!lu.isSquare() || lu.getRows() != piv.size() || lu.getRows() != b.getRows()) {
        throw IllegalArgumentException("luSolve: dimension mismatch")
    }

    let LAPACK_COL_MAJOR: Int64 = 102
    let n = lu.getCols()
    let nrhs = b.getCols()
    let lda: Int64 = n
    let ldb: Int64 = n
    
    let trans: Rune = r'N'
    let ctrans: UInt32 = UInt32(trans)
    var bcopy = b.copy()
    let info = unsafe { LAPACKE_dgetrs(IntNative(LAPACK_COL_MAJOR), UInt8(ctrans), IntNative(n), IntNative(nrhs),
                        lu.ptr, IntNative(lda), piv.ptr, bcopy.ptr, IntNative(ldb))}
    return bcopy
}

public func luSolve(lu: Matrix<Float64>, piv: Vector<Int32>, b: Vector<Float64>): Vector<Float64> {
    return luSolve(lu, piv, b.toMatrix()).toVector()
}

/* Compute the inverse of an LU-factored general matrix. */
foreign func LAPACKE_dgetri(
    matrix_layout: IntNative,  // row or column major
    n: IntNative,              // size of A
    a: CPointer<Unit>,         // matrix A (Float64)
    lda: IntNative,            // leading dimension of A
    ipiv: CPointer<Unit>       // ipiv array, as returned by ?getrf (Int32)
): Int64

public func inv(a: Matrix<Float64>) {
    let LAPACK_COL_MAJOR: Int64 = 102
    let m = a.getRows()
    let n = a.getCols()
    let lda: Int64 = max(1,n)
    let minmn = min(m,n)
    var ipiv = empty<Int32>(minmn)
    var repA = a.copy()
    let info = unsafe { LAPACKE_dgetrf(IntNative(LAPACK_COL_MAJOR), IntNative(m), IntNative(n), 
                        repA.ptr, IntNative(lda), ipiv.ptr) }
    let info2 = unsafe { LAPACKE_dgetri(IntNative(LAPACK_COL_MAJOR),  IntNative(n), 
                        repA.ptr, IntNative(lda), ipiv.ptr) }
    return repA
}

public func det(a: Matrix<Float64>) {
    if (!a.isSquare()) {
        throw IllegalArgumentException("det: input must be square matrix")
    }

    let (b, ipiv) = luFactor(a)
    let sz = a.getRows()

    // Obtain number of swaps
    var arr = vector<Int64>(sz, {i => i + 1})
    var swap: Int64 = 0
    for (i in 0..sz) {
        if (i != Int64(ipiv[i]) - 1) {
            swap += 1
        }
    }

    var res: Float64 = 1.0
    for (i in 0..sz) {
        res *= b[i,i]        
    }
    if (swap % 2 == 1) {
        res *= -1.0
    }
    return res
}

func testDgetrf() {
    let LAPACK_COL_MAJOR: Int64 = 102
    var a = matrix<Float64>(
        [[2.0, 5.0, 8.0, 7.0],
         [5.0, 2.0, 2.0, 8.0],
         [7.0, 5.0, 6.0, 6.0],
         [5.0, 4.0, 4.0, 8.0]]
    )
    let m = a.getRows()
    let n = a.getCols()
    let lda: Int64 = m
    // Int32 is needed for LAPACK
    let minmn = min(m,n)
    var ipiv = zeros<Int32>(minmn)
    let info = unsafe { LAPACKE_dgetrf(IntNative(LAPACK_COL_MAJOR), IntNative(m), IntNative(n), 
                        a.ptr, IntNative(lda), ipiv.ptr) }
    var l = a.copy()
    var u = a.copy()
    for (i in 0..m) {
        for (j in 0..n) {
            if (i <= j) {
                l[i,j] = 0.0
            }
            if (i == j) {
                l[i,j] = 1.0
            }
            if (i > j) {
                u[i,j] = 0.0
            }
        } 
    }

    var arr = vector<Int64>(min(m,n), {i => i + 1})

    for (i in 0..minmn) {
        var tmp = arr[i]
        arr[i] = arr[Int64(ipiv[i])-1]
        arr[Int64(ipiv[i])-1] = tmp
    }

    var p = zeros<Float64>(minmn, minmn)
    for (i in 0..minmn) {
        p[arr[i]-1,i] =1.0
    }

    var plu = p * l * u
    var expected_plu = matrix<Float64>(
        [[2.0, 5.0, 8.0, 7.0],
         [5.0, 2.0, 2.0, 8.0],
         [7.0, 5.0, 6.0, 6.0],
         [5.0, 4.0, 4.0, 8.0]]
    )
    assertApproxEqual(plu, expected_plu, atol:1e-6)
}

func testDgetrs() {
    let LAPACK_COL_MAJOR: Int64 = 102
    var a = matrix<Float64>(
        [[2.0, 5.0, 8.0, 7.0],
         [5.0, 2.0, 2.0, 8.0],
         [7.0, 5.0, 6.0, 6.0],
         [5.0, 4.0, 4.0, 8.0]]
    )
    var b = matrix<Float64>(
        [[1.0], [1.0], [1.0], [1.0]]
    )
    let n = a.getCols()
    let nrhs = b.getCols()
    let lda: Int64 = n
    let ldb: Int64 = n
    var (lu_res, piv) = luFactor(a)
    let trans: Rune = r'N'
    let ctrans: UInt32 = UInt32(trans)
    let info = unsafe { LAPACKE_dgetrs(IntNative(LAPACK_COL_MAJOR), UInt8(ctrans), IntNative(n), IntNative(nrhs), 
                        lu_res.ptr, IntNative(lda), piv.ptr, b.ptr, IntNative(ldb))}
    var expected_x = matrix<Float64>(
        [[0.051546], [-0.082474], [0.082474], [0.092784]]
    )
    assertApproxEqual(b, expected_x, atol:1e-6)
}

func testDgetri(): Unit {
    let LAPACK_COL_MAJOR: Int64 = 102
    var a = matrix<Float64>(
        [[1.80,  2.88, 2.05, -0.89],
         [5.25, -2.95, -0.95, -3.80],
         [1.58, -2.69, -2.90, -1.04],
         [-1.11, -0.66, -0.59, 0.80]]
    )
    let m = a.getRows()
    let n = a.getCols()
    let lda: Int64 = max(1,n)
    let minmn = min(m,n)
    var ipiv = zeros<Int32>(minmn)
    let info = unsafe { LAPACKE_dgetrf(IntNative(LAPACK_COL_MAJOR), IntNative(m), IntNative(n), 
                        a.ptr, IntNative(lda), ipiv.ptr) }
    let info2 = unsafe { LAPACKE_dgetri(IntNative(LAPACK_COL_MAJOR),  IntNative(n), 
                        a.ptr, IntNative(lda), ipiv.ptr) }
    
    var expected_inva = matrix<Float64>(
        [[1.771998, 0.575691, 0.084325, 4.815502],
         [-0.117466, -0.445615, 0.411363, -1.712581],
         [0.179856, 0.452662, -0.667565, 1.482400],
         [2.494382, 0.764977, -0.035954, 7.611900]]
    )
    assertApproxEqual(a, expected_inva, atol:1e-6)
}

@Test
public class TestLU {
    @TestCase
    func testLapackDgetrf(): Unit {
        testDgetrf()
    }

    @TestCase
    func testLapackDgetrs(): Unit {
        testDgetrs()
    }

    @TestCase
    func testLapackDgetri(): Unit {
        testDgetri()
    }

    @TestCase
    func testLu(): Unit {
        var a = matrix<Float64>(
            [[4.0, 2.0, 1.0, 5.0],
             [8.0, 7.0, 2.0, 10.0],
             [4.0, 8.0, 3.0, 6.0],
             [6.0, 8.0, 4.0, 9.0]]
        )
        var (p, l, u) = lu(a)
        var plu = p * l * u
        var expected_plu = matrix<Float64>(
            [[4.0, 2.0, 1.0, 5.0],
             [8.0, 7.0, 2.0, 10.0],
             [4.0, 8.0, 3.0, 6.0],
             [6.0, 8.0, 4.0, 9.0]]
        )
        @Assert(approxEqual(plu, expected_plu, atol:1e-6))
    }

    @TestCase
    func testLuFactor(): Unit {
        var a = matrix<Float64>(
            [[2.0, 5.0, 8.0, 7.0],
             [5.0, 2.0, 2.0, 8.0],
             [7.0, 5.0, 6.0, 6.0],
             [5.0, 4.0, 4.0, 8.0]]
        )
        var (l_u, piv) = luFactor(a)
        var expected_lu = matrix<Float64>(
            [[7.000000, 5.000000, 6.000000, 6.000000],
             [0.285714, 3.571429, 6.285714, 5.285714],
             [0.714286, 0.120000, -1.040000, 3.080000],
             [0.714286, -0.440000, -0.461538, 7.461538]]
        )
        var expected_piv = vector<Int32>(
            [3, 3, 4, 4]
        )
        @Assert(approxEqual(l_u, expected_lu, atol:1e-6))
        @Assert(piv == expected_piv)
    }

    @TestCase
    func testLuSolve(): Unit {
        var a = matrix<Float64>(
            [[2.0, 5.0, 8.0, 7.0],
             [5.0, 2.0, 2.0, 8.0],
             [7.0, 5.0, 6.0, 6.0],
             [5.0, 4.0, 4.0, 8.0]]
        )
        var b = vector<Float64>(
            [1.0, 1.0, 1.0, 1.0]
        )
        var (l_u, piv) = luFactor(a)
        var x = luSolve(l_u, piv, b)
        var ax = a * x
        @Assert(approxEqual(ax, b, atol:1e-6))
    }

    @TestCase
    func testInv(): Unit {
        var a = matrix<Float64>(
            [[1.80,  2.88, 2.05, -0.89],
             [5.25, -2.95, -0.95, -3.80],
             [1.58, -2.69, -2.90, -1.04],
             [-1.11, -0.66, -0.59, 0.80]]
        )
        var inva = inv(a)
        var result = a * inva
        var expected_result = matrix<Float64>(
            [[1.0, 0.0, 0.0, 0.0],
             [0.0, 1.0, 0.0, 0.0],
             [0.0, 0.0, 1.0, 0.0],
             [0.0, 0.0, 0.0, 1.0]]
        )
        @Assert(approxEqual(result, expected_result, atol:1e-6))
    }

    @TestCase
    func testDet(): Unit {
        var a = matrix<Float64>([
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0]
        ])
        @Assert(approxEqual(det(a), 0.0, atol:1e-6))
    }

    @TestCase
    func testDet2(): Unit {
        var a = matrix<Float64>([
            [0.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0]
        ])
        @Assert(approxEqual(det(a), 3.0, atol:1e-6))
    }

    @TestCase
    func testDet3(): Unit {
        var a = matrix<Float64>([
            [1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0],
            [0.0, 1.0, 0.0]
        ])
        @Assert(approxEqual(det(a), -1.0, atol:1e-6))
    }
}