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

[Fernando Perez]

On Sat, Jun 6, 2009 at 12:09 AM, Robert Kern<robert.kern at gmail.com> wrote:

> +1

OK, thanks.  I'll try to get it ready.

> diag_indices() can be made more efficient, but these are fine.

Suggestion?  Right now it's not obvious to me...

A few more questions:

- Are doctests considered enough testing for numpy, or are separate
tests also required?

- Where should these go?

- Any interest in also having the stuff below?  I'm needing to build
structured random arrays a lot (symmetric, anti-symmetric, symmetric
with  a particular diagonal, etc), and these are coming in handy.  If
you want them, I'll put the whole thing together (these use the
indexing utilities from the previous suggestion).

Thanks!


f

#### Other suggested utilities.  Not fully commented yet, but if they
are wanted for numpy, will be submitted in final form.

def structured_rand_arr(size, sample_func=np.random.random,
                        ltfac=None, utfac=None, fill_diag=None):
    """Make a structured random 2-d array of shape (size,size).

    Parameters
    ----------

    size : int
      Determines the shape of the output array: (size,size).

    sample_func : function, optional.
      Must be a function which when called with a 2-tuple of ints, returns a
      2-d array of that shape.  By default, np.random.random is used, but any
      other sampling function can be used as long as matches this API.

    utfac : float, optional
      Multiplicative factor for the lower triangular part of the matrix.

    ltfac : float, optional
      Multiplicative factor for the lower triangular part of the matrix.

    fill_diag : float, optional
      If given, use this value to fill in the diagonal.  Otherwise the diagonal
      will contain random elements.
    """
    # Make a random array from the given sampling function
    mat0 = sample_func((size,size))
    # And the empty one we'll then fill in to return
    mat = np.empty_like(mat0)
    # Extract indices for upper-triangle, lower-triangle and diagonal
    uidx = triu_indices(size,1)
    lidx = tril_indices(size,-1)
    didx = diag_indices(size)
    # Extract each part from the original and copy it to the output, possibly
    # applying multiplicative factors.  We check the factors instead of
    # defaulting to 1.0 to avoid unnecessary floating point multiplications
    # which could be noticeable for very large sizes.
    if utfac:
        mat[uidx] = utfac * mat0[uidx]
    else:
        mat[uidx] = mat0[uidx]
    if ltfac:
        mat[lidx] = itfac * mat0.T[lidx]
    else:
        mat[lidx] = mat0.T[lidx]
    # If fill_diag was provided, use it; otherwise take the values in the
    # diagonal from the original random array.
    if fill_diag:
        mat[didx] = fill_diag
    else:
        mat[didx] = mat0[didx]

    return mat


def symm_rand_arr(size,sample_func=np.random.random,fill_diag=None):
    """Make a symmetric random 2-d array of shape (size,size).

    Parameters
    ----------

    n : int
      Size of the output array.

    fill_diag : float, optional
      If given, use this value to fill in the diagonal.  Useful for
      """
    return structured_rand_arr(size,sample_func,fill_diag=fill_diag)


def antisymm_rand_arr(size,sample_func=np.random.random,fill_diag=None):
    """Make an anti-symmetric random 2-d array of shape (size,size).

    Parameters
    ----------

    n : int
      Size of the output array.
      """
    return structured_rand_arr(size,sample_func,ltfac=-1.0,fill_diag=fill_diag)