added support for DGTSV. (#2)

Author: dastrobu

AccelerateArray.podspec

@@ -16,7 +16,7 @@ Pod::Spec.new do |s|
   #
 
   s.name         = "AccelerateArray"
-  s.version      = "0.0.2"
+  s.version      = "0.1.0"
   s.summary      = "Swift Array Extensions for the Apple Accelerate Framework"
 
   # This description is used to generate tags and improve search results.

Sources/AccelerateArray/lapack.swift

@@ -174,10 +174,55 @@ public extension Array where Element == Double {
         assert(ldb >= Swift.max(1, n), "\(ldb) >= max(1, \(n))")
 
         var info: Int32 = 0
-
         dgesv_(&n, &nrhs, &self, &lda, &ipiv, &B, &ldb, &info)
         if info != 0 {
             throw LapackError.dgesv(info)
         }
     }
+
+    /// DGTSV  solves the equation
+    ///
+    ///    A*X = B,
+    ///
+    /// where A is an n by n tridiagonal matrix, by Gaussian elimination with
+    /// partial pivoting.
+    ///
+    /// Note that the equation  A**T*X = B  may be solved by interchanging the
+    /// order of the arguments DU and DL.
+    ///
+    /// This array represents the diagonal of A.
+    ///
+    /// http://www.netlib.org/lapack/explore-html/d4/d62/group__double_g_tsolve_ga2bf93f2ddefa5e671866eb2191dc19d4.html#ga2bf93f2ddefa5e671866eb2191dc19d4
+    ///
+    /// - Parameters:
+    ///     - nrhs: The number of right hand sides, i.e., the number of columns of the matrix B.  NRHS >= 0.
+    ///     - dl: On entry, DL must contain the (n-1) sub-diagonal elements of A.
+    ///           On exit, DL is overwritten by the (n-2) elements of the
+    ///           second super-diagonal of the upper triangular matrix U from
+    ///           the LU factorization of A, in DL(1), ..., DL(n-2).
+    ///     - du: On entry, DU must contain the (n-1) super-diagonal elements of A.
+    ///           On exit, DU is overwritten by the (n-1) elements of the first
+    ///           super-diagonal of U.
+    ///     - B:   On entry, the N by NRHS matrix of right hand side matrix B.
+    ///            On exit, if no error was thrown, the N by NRHS solution matrix X.
+    ///
+    mutating func gtsv(nrhs: Int, dl: inout [Double], du: inout [Double], B: inout [Double]) throws {
+        assert(count - 1 == dl.count, "\(count) - 1 == \(dl.count)")
+        assert(count - 1 == du.count, "\(count) - 1 == \(du.count)")
+        var n = Int32(count)
+
+        var nrhs = Int32(B.count / Int(n))
+        assert(nrhs >= 1, "\(nrhs) >= 1")
+        assert(B.count == Int(nrhs) * count, "\(B.count) == \(nrhs) * \(n)")
+
+        var ldb = Int32(B.count / Int(nrhs))
+        assert(ldb * nrhs == B.count, "\(ldb) * \(nrhs) == \(B.count)")
+        assert(ldb >= Swift.max(1, n), "\(ldb) >= max(1, \(n))")
+
+        var info: Int32 = 0
+        dgtsv_(&n, &nrhs, &dl, &self, &du, &B, &ldb, &info)
+        if info != 0 {
+            throw LapackError.dgesv(info)
+        }
+    }
 }

Tests/AccelerateArrayTests/lapack.swift

@@ -38,7 +38,7 @@ class LapackTests: XCTestCase {
 
         let n = 4
         var P: [Double] = Array(repeating: 0, count: n * n)
-        for i in 0 ..< p.count {
+        for i in 0..<p.count {
             // i iterates columns of P (in row major)
             // p[i] indicates which element in the column must be set to one, to create the permutation matrix
             P[i + p[i] * n] = 1.0
@@ -72,6 +72,9 @@ class LapackTests: XCTestCase {
         ]
         // convert A to col major
         var At = A.mtrans(m: 2, n: 2)
+        let b: [Double] = [
+            1.0, 1.0
+        ]
         // B in row major
         let B: [Double] = [
             1.0, 2.0, 3.0,
@@ -84,21 +87,65 @@ class LapackTests: XCTestCase {
             -2.0, 1.0,
             1.5, -0.5,
         ]
-        // x1 is in row major
-        let x1 = Ainv.mmul(B: B, m: 2, n: 3, p: 2)
+        let x1 = Ainv.mmul(B: b, m: 2, n: 1, p: 2)
+        // X1 is in row major
+        let X1 = Ainv.mmul(B: B, m: 2, n: 3, p: 2)
 
-        // solution is stored row major in Bt
+        // solution is stored col major in Bt
         try At.gesv(B: &Bt)
-        let x2 = Bt.mtrans(m: 2, n: 3)
+        let X2 = Bt.mtrans(m: 2, n: 3)
+
+        XCTAssertEqual(x1[0], X1[0], accuracy: 1e-15)
+        XCTAssertEqual(x1[1], X1[3], accuracy: 1e-15)
+        XCTAssertEqual(X1, X2, accuracy: 1e-15)
+    }
+
+    func testGtsvDouble() throws {
+        // A in row major
+        let A: [Double] = [
+            1.0, 1.0, 0.0, 0.0,
+            -1.0, 2.0, 2.0, 0.0,
+            0.0, -2.0, 3.0, 3.0,
+            0.0, 0.0, -3.0, 4.0,
+        ]
+        // convert A to col major
+        var At = A.mtrans(m: 4, n: 4)
+
+        // diagonals of A
+        var d: [Double] = [1.0, 2.0, 3.0, 4.0, ]
+        var du: [Double] = [1.0, 2.0, 3.0, ]
+        var dl: [Double] = [-1.0, -2.0, -3.0, ]
+
+        // B in row major
+        let B: [Double] = [
+            1.0, 2.0,
+            1.0, 2.0,
+            1.0, 2.0,
+            1.0, 2.0,
+        ]
+        // B in col major
+        var Bt = B.mtrans(m: 2, n: 4)
+        // make a copy of Bt
+        var Ct = Bt
+
+        // solve with general solver
+        // solution is stored col major in Ct
+        try At.gesv(B: &Ct)
+        let X1 = Ct.mtrans(m: 4, n: 2)
+
+        // solution is stored col major in Bt
+        try d.gtsv(nrhs: 2, dl: &dl, du: &du, B: &Bt)
+        let X2 = Bt.mtrans(m: 4, n: 2)
 
-        XCTAssertEqual(x1, x2, accuracy: 1e-15)
+        XCTAssertEqual(X1, X2, accuracy: 1e-15)
     }
 
     static var allTests: [(String, (LapackTests) -> () throws -> Void)] {
         return [
             ("testGetrfDouble", testGetrfDouble),
             ("testGetriDouble", testGetriDouble),
             ("testGesvDouble", testGesvDouble),
+            ("testGtsvDouble", testGtsvDouble),
         ]
     }
 }