Correlation

correlation.swift

// Code and tests adapted from Numpy library

func dot_product(in arr1: [Double], index1: Int, in arr2: [Double], index2: Int, n: Int, result_index: Int, inout result_arr: [Double]) {
    var sum = 0.0;
    for i in 0..<n {
        sum += arr1[index1 + i] * arr2[index2 + i]
    }
    result_arr[result_index] = sum
}


func correlate(arr1: [Double], arr2: [Double], mode: String) -> [Double]? {
    var arr1 = arr1
    var arr2 = arr2
    var inverted = true
    
    var n1 = arr1.count
    var n2 = arr2.count
    if n1 < n2 {
        let _temp1 = arr1
        arr1 = arr2
        arr2 = _temp1
        
        let _temp2 = n1
        n1 = n2
        n2 = _temp2
        inverted = true
    }
    else {
        inverted = false
    }
    
    var length = n1
    var n = n2
    var n_left = 0
    var n_right = 0
    if mode == "valid" {
        length = length - n + 1
    }
    else if mode == "same" {
        n_left = Int(floor(Double(n) / 2.0))
        n_right = n - n_left - 1
    }
    else if mode == "full" {
        n_right = n - 1
        n_left = n - 1
        length = length + n - 1
    }
    else {
        print("ERROR: mode must be valid, same, or full")
        return nil
    }
    n = n - n_left
    
    var result = Array<Double>(count: length, repeatedValue: 0.0)
    var result_index = 0
    var arr1_index = 0
    var arr2_index = n_left
    
    for _ in 0..<n_left {
        dot_product(in: arr1, index1: arr1_index, in: arr2, index2: arr2_index, n: n, result_index: result_index, result_arr: &result)
        n += 1
        arr2_index -= 1
        result_index += 1
    }
    
    for _ in 0..<(n1 - n2 + 1) {
        dot_product(in: arr1, index1: arr1_index, in: arr2, index2: arr2_index, n: n, result_index: result_index, result_arr: &result)
        arr1_index += 1
        result_index += 1
    }
    
    for _ in 0..<n_right {
        n -= 1
        dot_product(in: arr1, index1: arr1_index, in: arr2, index2: arr2_index, n: n, result_index: result_index, result_arr: &result)
        arr1_index += 1
        result_index += 1
    }
    
    if inverted {
        result = result.reverse()
    }
    
    return result
}



import XCTest

class CorrelateTests: XCTestCase {
    func testValid() {
        let result = correlate([1.0, 2.0, 3.0], arr2: [0.0, 1.0, 0.5], mode: "valid")
        XCTAssertEqual(result!, [3.5], "correlate - valid")
    }
    
    func testSame() {
        let result = correlate([1.0, 2.0, 3.0], arr2: [0.0, 1.0, 0.5], mode: "same")
        XCTAssertEqual(result!, [2.0, 3.5, 3.0], "correlate - same")
    }

    func testFull() {
        let result = correlate([1.0, 2.0, 3.0], arr2: [0.0, 1.0, 0.5], mode: "full")
        XCTAssertEqual(result!, [0.5, 2.0, 3.5, 3.0, 0.0], "correlate - full")
    }
}