numpy.select

numpy.select(condlist, choicelist, default=0)

Return an array drawn from elements in choicelist, depending on conditions.

Parameters:

condlist : list of N boolean arrays of length M

The conditions C_0 through C_(N-1) which determine from which vector the output elements are taken.

choicelist : list of N arrays of length M

Th vectors V_0 through V_(N-1), from which the output elements are chosen.

Returns:

output : 1-dimensional array of length M

The output at position m is the m-th element of the first vector V_n for which C_n[m] is non-zero. Note that the output depends on the order of conditions, since the first satisfied condition is used.

Notes

Equivalent to:

output = []
for m in range(M):
    output += [V[m] for V,C in zip(values,cond) if C[m]]
              or [default]

Examples

>>> t = np.arange(10)
>>> s = np.arange(10)*100
>>> condlist = [t == 4, t > 5]
>>> choicelist = [s, t]
>>> np.select(condlist, choicelist)
array([  0,   0,   0,   0, 400,   0,   6,   7,   8,   9])