[SciPy-dev] Nifty matrix functions that others may find helpful

[Jonathan Taylor]

I have been looking for the indices function for a while without knowing
that was the name I was looking for.  I am glad I finally found it. 
Following are two special cases (inspired by R) that I find immensely
useful for getting row and column indices out of a matrix.  In
particular I use them often to extract the upper triangle from a matrix
via:

upper_triangle = m[row(m)>col(m)]

Maybe there is something better in scipy for this anyways that I am
missing?

Cheers.
Jon.

def col(m):
    """col(m) returns a matrix of the same size of m where each element
    contains an integer denoting which column it is in.  For example,

    >>> m = eye(3)
    >>> m
    array([[1, 0, 0],
           [0, 1, 0],
           [0, 0, 1]])
    >>> col(m)
    array([[0, 1, 2],
           [0, 1, 2],
           [0, 1, 2]])
    """
    assert len(m.shape) == 2, "should be a matrix"
    return indices(m.shape)[1]

def row(m):
    """row(m) returns a matrix of the same size of m where each element
    contains an integer denoting which row it is in.  For example,

    >>> m = eye(3)
    >>> m
    array([[1, 0, 0],
           [0, 1, 0],
           [0, 0, 1]])
    >>> row(m)
    array([[0, 0, 0],
           [1, 1, 1],
           [2, 2, 2]])

    """
    assert len(m.shape) == 2, "should be a matrix"
    return indices(m.shape)[0]