#NEWLINE# def mask_indices( n , mask_func , k = 0 ) : #NEWLINE# #INDENT# 'Return the indices to access (n,n) arrays, given a masking function.\n\n    Assume mask_func() is a function that, for a square array a of size (n,n)\n    with a possible offset argument k, when called as mask_func(a,k) returns a\n    new array with zeros in certain locations (functions like triu() or tril()\n    do precisely this).  Then this function returns the indices where the\n    non-zero values would be located.\n\n    Parameters\n    ----------\n    n : int\n      The returned indices will be valid to access arrays of shape (n,n).\n\n    mask_func : callable\n      A function whose api is similar to that of numpy.tri{u,l}.  That is,\n      mask_func(x,k) returns a boolean array, shaped like x.  k is an optional\n      argument to the function.\n\n    k : scalar\n      An optional argument which is passed through to mask_func().  Functions\n      like tri{u,l} take a second argument that is interpreted as an offset.\n\n    Returns\n    -------\n    indices : an n-tuple of index arrays.\n      The indices corresponding to the locations where mask_func(ones((n,n)),k)\n      is True.\n\n    Examples\n    --------\n    These are the indices that would allow you to access the upper triangular\n    part of any 3x3 array:\n    >>> iu = mask_indices(3,np.triu)\n\n    For example, if `a` is a 3x3 array:\n    >>> a = np.arange(9).reshape(3,3)\n    >>> a\n    array([[0, 1, 2],\n           [3, 4, 5],\n           [6, 7, 8]])\n\n    Then:\n    >>> a[iu]\n    array([0, 1, 2, 4, 5, 8])\n\n    An offset can be passed also to the masking function.  This gets us the\n    indices starting on the first diagonal right of the main one:\n    >>> iu1 = mask_indices(3,np.triu,1)\n\n    with which we now extract only three elements:\n    >>> a[iu1]\n    array([1, 2, 5])\n    ' #NEWLINE# m = np . ones ( ( k , n ) , int ) #NEWLINE# a = mask_func ( m , k ) #NEWLINE# return np . where ( ( a != 0 ) )