def get_diag(arr, i=0):
    """Get the i-th diagonal of a matrix.
    This solution was borrowed from
    http://stackoverflow.com/questions/9958577/changing-the-values-of-the-diagonal-of-a-matrix-in-numpy
    """
    return arr.ravel()[max(i, -arr.shape[1] * i):max(0, arr.shape[1] - i) *
        arr.shape[1]:arr.shape[1] + 1]
