Update linear regression class

Package.swift

@@ -30,9 +30,12 @@ let package = Package(
                 .product(
                     name: "RefdsShared",
                     package: "refds-shared"),
+            ],
+            swiftSettings: [
+                .define("ACCELERATE_NEW_LAPACK")
             ]),
         .testTarget(
             name: "RefdsAlgorithmTests",
-            dependencies: ["RefdsAlgorithm"]),
+            dependencies: ["RefdsAlgorithm"])
     ]
 )

Sources/RefdsAlgorithm/MachineLearning/LinearRegression.swift

@@ -1,27 +1,95 @@
 import Foundation
+import Accelerate
 
-class LinearRegression {
-    private var x: [Double]
-    private var y: [Double]
+public class LinearRegression {
+    private var features: [[NSNumber]]
+    private var targets: [NSNumber]
+    private var coefficients: [Double] = []
 
-    public init(x: [Double], y: [Double]) {
-        self.x = x
-        self.y = y
-    }
+    private let semaphore = DispatchSemaphore(value: 0)
+    private let queue = DispatchQueue(
+        label: "refds.algorithm.linearRegression",
+        qos: .background
+    )
 
-    private func average(_ input: [Double]) -> Double {
-        return input.reduce(0, +) / Double(input.count)
+    public init(
+        features: [[NSNumber]],
+        targets: [NSNumber]
+    ) {
+        self.features = features
+        self.targets = targets
+
+        queue.async { self.makeCoefficients() }
     }
+
+    private func makeCoefficients() {
+        self.coefficients = []
+        let rows = features.count
+
+        guard let cols = features[safe: 0]?.count,
+              rows == targets.count,
+              !features.isEmpty && !features[0].isEmpty
+        else {
+            semaphore.signal()
+            return
+        }
+
+        var X = features
+        for i in 0 ..< rows { X[i].insert(1.0, at: 0) }
+
+        let flatX = X.flatMap { $0 }.map { $0.doubleValue }
+        let flatY = targets.map { $0.doubleValue }
+        var coefficients = [Double](repeating: 0.0, count: cols + 1)
+
+        var XT = [Double](repeating: 0.0, count: rows * (cols + 1))
+        vDSP_mtransD(flatX, 1, &XT, 1, vDSP_Length(cols + 1), vDSP_Length(rows))
+
+        var XTX = [Double](repeating: 0.0, count: (cols + 1) * (cols + 1))
+        vDSP_mmulD(XT, 1, flatX, 1, &XTX, 1, vDSP_Length(cols + 1), vDSP_Length(cols + 1), vDSP_Length(rows))
 
-    private func multiply(_ a: [Double], _ b: [Double]) -> [Double] {
-        return zip(a, b).map(*)
+        var N = __CLPK_integer(cols + 1)
+        var pivots = [__CLPK_integer](repeating: 0, count: cols + 1)
+        var workspace = [Double](repeating: 0.0, count: cols + 1)
+        var error: __CLPK_integer = 0
+
+        let _ = withUnsafeMutablePointer(to: &N) {
+            dgetrf_($0, $0, &XTX, $0, &pivots, &error)
+        }
+
+        guard error == 0 else {
+            semaphore.signal()
+            return
+        }
+
+        let _ = withUnsafeMutablePointer(to: &N) {
+            dgetri_($0, &XTX, $0, &pivots, &workspace, $0, &error)
+        }
+
+        guard error == 0 else {
+            semaphore.signal()
+            return
+        }
+
+        var XTY = [Double](repeating: 0.0, count: cols + 1)
+        vDSP_mmulD(XT, 1, flatY, 1, &XTY, 1, vDSP_Length(cols + 1), 1, vDSP_Length(rows))
+        vDSP_mmulD(XTX, 1, XTY, 1, &coefficients, 1, vDSP_Length(cols + 1), 1, vDSP_Length(cols + 1))
+
+        self.coefficients = coefficients
+        semaphore.signal()
     }
-
-    public var result: (Double) -> Double {
-        let sum1 = average(multiply(x, y)) - average(x) * average(y)
-        let sum2 = average(multiply(x, x)) - pow(average(x), 2)
-        let slope = sum1 / sum2
-        let intercept = average(y) - slope * average(x)
-        return { x in intercept + slope * x }
+
+    public func predict(for features: [Double]) -> Double? {
+        semaphore.wait()
+        guard !coefficients.isEmpty,
+              var prediction = coefficients[safe: 0]
+        else { return nil }
+
+        for i in features.indices {
+            if let coefficient = coefficients[safe: i + 1] {
+                prediction += coefficient * features[i]
+            }
+        }
+        semaphore.signal()
+        return prediction
     }
 }

Tests/RefdsAlgorithmTests/MachineLearningTests.swift

@@ -5,8 +5,8 @@ final class MachineLearningTests: XCTestCase {
     func testLinearRegressionShouldCompletesWithValidValue() {
         let carAge: [Double] = [10, 8, 3, 3, 2, 1]
         let carPrice: [Double] = [500, 400, 7000, 8500, 11000, 10500]
-        let linearRegression = LinearRegression(x: carAge, y: carPrice)
-        let result = linearRegression.result(4)
+        let linearRegression = LinearRegression(features: carAge.map { [$0] }, targets: carPrice)
+        let result = linearRegression.predict(for: [4]) ?? .zero
         XCTAssertEqual(Int(result), 6952)
     }
 }