[Numpy-discussion] Functions for indexing into certain parts of an array (2d)

[Keith Goodman]

On Sat, Jun 6, 2009 at 2:01 PM, Robert Kern <robert.kern at gmail.com> wrote:
>  There is a neat trick for accessing the diagonal of an existing array
> (a.flat[::a.shape[1]+1]), but it won't work to implement
> diag_indices().

Perfect. That's 3x faster.

def fill_diag(arr, value):
    if arr.ndim != 2:
        raise ValueError, "Input must be 2-d."
    if arr.shape[0] != arr.shape[1]:
        raise ValueError, 'Input must be square.'
    arr.flat[::arr.shape[1]+1] = value