I recently completed 10 Days of Statistics on HackerRank. Yay!

In the process, I needed some matrix operations for a medium-difficulty problem. And here they are, code style be damned :

func transpose(_ matrix: [[Double]]) -> [[Double]] {
    let rowCount = matrix.count
    let colCount = matrix[0].count
    var transposed : [[Double]] = Array(repeating: Array(repeating: 0.0, count: rowCount), count: colCount)
    for rowPos in 0..<matrix.count {
        for colPos in 0..<matrix[0].count {
            transposed[colPos][rowPos] = matrix[rowPos][colPos]
        }
    }
    return transposed
}

func multiply(_ A: [[Double]], _ B: [[Double]]) -> [[Double]] {
    let rowCount = A.count
    let colCount = B[0].count
    var product : [[Double]] = Array(repeating: Array(repeating: 0.0, count: colCount), count: rowCount)
    for rowPos in 0..<rowCount {
        for colPos in 0..<colCount {
            for i in 0..<B.count {
                product[rowPos][colPos] += A[rowPos][i] * B[i][colPos]
            }
        }
    }
    return product
}

// gauss jordan inversion
func inverse(_ matrix: [[Double]]) -> [[Double]] {
    // augment matrix
    var matrix = matrix
    var idrow = Array(repeating: 0.0, count: matrix.count)
    idrow[0] = 1.0
    for row in 0..<matrix.count {
        matrix[row] += idrow
        idrow.insert(0.0, at:0)
        idrow.removeLast()
    }
    
    // partial pivot
    for row1 in 0..<matrix.count {
        for row2 in row1..<matrix.count {
            if abs(matrix[row1][row1]) < abs(matrix[row2][row2]) {
                (matrix[row1],matrix[row2]) = (matrix[row2],matrix[row1])
            }
        }
    }
    
    // forward elimination
    for pivot in 0..<matrix.count {
        // multiply
        let arg = 1.0 / matrix[pivot][pivot]
        for col in pivot..<matrix[pivot].count {
            matrix[pivot][col] *= arg
        }
        
        // multiply-add
        for row in (pivot+1)..<matrix.count {
            let arg = matrix[row][pivot] / matrix[pivot][pivot]
            for col in pivot..<matrix[row].count {
                matrix[row][col] -= arg * matrix[pivot][col]
            }
        }
    }
    
    // backward elimination
    for pivot in (0..<matrix.count).reversed() {
        // multiply-add
        for row in 0..<pivot {
            let arg = matrix[row][pivot] / matrix[pivot][pivot]
            for col in pivot..<matrix[row].count {
                matrix[row][col] -= arg * matrix[pivot][col]
            }
        }
    }
    
    // remove identity
    for row in 0..<matrix.count {
        for _ in 0..<matrix.count {
            matrix[row].remove(at:0)
        }
    }
    
    return matrix
}

let X = [ [1.0, 2.0, 3.0], [4.0, 5.0, 11.0], [7.0, 8.0, 9.0] ]
let XI = inverse(X)
let I = multiply(X,XI)
print(I)

That’s the identity matrix popping out at the end, which validates my implementation.

But what’s this? A 60-line method? Uncle Bob would not be pleased.

Comments should not take the place of good variable/method names. Those section comments give clues as to where my methods should be :

func augment(_ matrix: [[Double]]) -> [[Double]] {
    var augmented = matrix
    var idrow = Array(repeating: 0.0, count: matrix.count)
    idrow[0] = 1.0
    for row in 0..<matrix.count {
        augmented[row] += idrow
        idrow.insert(0.0, at:0)
        idrow.removeLast()
    }
    return augmented
}

func deaugment(_ matrix: [[Double]]) -> [[Double]] {
    var deaugmented = matrix
    
    for row in 0..<matrix.count {
        for _ in 0..<matrix.count {
            deaugmented[row].remove(at:0)
        }
    }
    return deaugmented
}

func partialPivot(_ matrix: inout [[Double]]) {
    for row1 in 0..<matrix.count {
        for row2 in row1..<matrix.count {
            if abs(matrix[row1][row1]) < abs(matrix[row2][row2]) {
                (matrix[row1],matrix[row2]) = (matrix[row2],matrix[row1])
            }
        }
    }
}

func scaleRow(_ matrix: inout [[Double]], row: Int, scale: Double) {
    for col in 0..<matrix[row].count {
        matrix[row][col] *= scale
    }
}

func addRow(_ matrix: inout [[Double]], row: Int, scaledBy: Double, toRow: Int) {
    for col in 0..<matrix[row].count {
        matrix[toRow][col] += scaledBy * matrix[row][col]
    }
}

func pivot(_ matrix: inout [[Double]], row pivotRow: Int, col pivotCol: Int, forward: Bool) {
    let scale = 1.0 / matrix[pivotRow][pivotCol]
    scaleRow(&matrix, row: pivotRow, scale: scale)
    
    if forward {
        for toRow in (pivotRow+1)..<matrix.count {
            let scaleBy = -1.0 * matrix[toRow][pivotCol]
            addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow)
        }
    } else {
        for toRow in (0..<pivotRow).reversed() {
            let scaleBy = -1.0 * matrix[toRow][pivotCol]
            addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow)
        }
    }
}

func gaussJordanInverse(_ matrix: [[Double]]) -> [[Double]] {
    var matrix = augment(matrix)
    partialPivot(&matrix)
    
    for p in 0..<matrix.count {
        pivot(&matrix, row: p, col: p, forward: true)
    }
    
    for p in (0..<matrix.count).reversed() {
        pivot(&matrix, row: p, col: p, forward: false)
    }

    matrix = deaugment(matrix)
    
    return matrix
}

let X = [ [1.0, 2.0, 3.0], [4.0, 5.0, 11.0], [7.0, 8.0, 9.0] ]
let XI = gaussJordanInverse(X)
let I = multiply(X,XI)
print(I)

Better. Uncle Bob would be proud (or give me credit for trying, anyway).

One more thing : Thanks to StackOverflow user Alexander, I have an even better way to express that pivot loop :

func pivot(_ matrix: inout [[Double]], row pivotRow: Int, col pivotCol: Int, forward: Bool) {
    let scale = 1.0 / matrix[pivotRow][pivotCol]
    scaleRow(&matrix, row: pivotRow, scale: scale)

    let range = forward ? AnyCollection((pivotRow+1)..<matrix.count) : AnyCollection((0..<pivotRow).reversed())

    for toRow in range {
        let scaleBy = -1.0 * matrix[toRow][pivotCol]
        addRow(&matrix, row: pivotRow, scaledBy: scaleBy, toRow: toRow)
    }
}