ENH: Update piecewise to handle multivariate functions

[pbrod]

Resurrected gh-5187 in this pullrequest since the former pullrequest crashed

Made piecewise function more general using a class implementation as suggested in #5187 (comment).

The new piecewise function allows to evaluate conditionally multivariate functions
of the type f(x1, x2,... ,xn), i.e., functions that have arbitrary
number N of input arguments.
For a single variate function the call is the same as before (for N=1):

np.piecewise(
    arr_like,
    [cond1, cond2, cond3],
    [func1, func2, func3, func_default],
    *func_args, **func_kwargs)

whereas for a multivariate function with N variables the call is updated to (N>1):

np.piecewise[N](
    arr_like_1,
    arr_like_2,
    ....
    arr_like_N,
    [cond1, cond2, cond3],
    [func1, func2, func3, func_default],
    *func_args, **func_kwargs)

The difference in behaviour between numpy.select and numpy.piecewise as mentioned in xref #16475 is documented

Examples

Examples that work both in previous and current version

>>> x = np.linspace(-2,2,5)
>>> np.piecewise(x,[x < 0, x >= 0], [lambda x: -x, lambda x: x])
array([ 2.,  1.,  0.,  1.,  2.])

or

>>> np.piecewise(x, [x < 0,], [lambda x: -x, lambda x: x])
array([ 2.,  1.,  0.,  1.,  2.])

Examples demonstrating the new functionality

>>> X,Y = np.meshgrid(x,x)
>>> np.piecewise[2](X,Y, [X*Y<0,], [lambda x,y: -x*y, lambda x,y: x*y])
array([[ 4.,  2., -0.,  2.,  4.],
       [ 2.,  1., -0.,  1.,  2.],
       [-0., -0.,  0.,  0.,  0.],
       [ 2.,  1.,  0.,  1.,  2.],
       [ 4.,  2.,  0.,  2.,  4.]])

>>> np.piecewise[2](X,Y,[X*Y<-0.5, X*Y>0.5], [lambda x,y: -x*y, lambda x,y: x*y, np.nan])
array([[  4.,   2.,  nan,   2.,   4.],
       [  2.,   1.,  nan,   1.,   2.],
       [ nan,  nan,  nan,  nan,  nan],
       [  2.,   1.,  nan,   1.,   2.],
       [  4.,   2.,  nan,   2.,   4.]])

Discussion

See also:
http://www.mail-archive.com/numpy-discussion@scipy.org/msg46304.html
gh-5187

The old API for single a variate function is piecewise(x, condlist, funclist, *args, **kw) where

• x is array-like or scalar,
• condlist is a list of bool arrays or bool scalars and
• funclist is a list of callables, f(x,*args,**kw), or scalars.

The current implementation of piecewise use call method of the class to keep the old API for single variate functions and get_item method for multivarite functions. This has the benefit that the namespace is not cluttered with an additional function to handle the multivariate case and that the old API can be kept without change.

The new API for multivariate function is piecewise[N](x_0,x_1, ..., x_N, condlist, funclist, *args, **kw) where
x_0, x_1,...x_N are arrays or scalars. This gives no additional names to the current name space

However, using a class to define a dynamic piecewise callable at runtime means that the dispatch mechanism
for the array__function protocol (described in NEP-18) must be abondoned for this implementation of piecewise since it is not implemented for callable objects generated at runtime.
If that is acceptable, then this pullrequest is ready.

Another solution could be to change the API slightly like this:
piecewise((x_1, x_2, ..., x_N), condlist, funclist, *args, **kw)
where xi = (x_1, x_2, ..., x_N) is a tuple of ndarrays or scalars.
This was done in gh-5187. The problem then was that piecewise((x1, x2, x3), ...) is already valid syntax, and it is the same as piecewise(np.array([x1, x2, x3])). We can't change that behavior without at least a deprecation cycle.

So perhaps the best way forward would be to introduce a FutureWarning when x is a tuple, and then maybe in numpy version 1.22 this new behavior can be introduced.

[InessaPawson]

@pbrod, thank you for working on this! Your PR was discussed at length at the latest triage meeting (numpy/archive@487b11b).
We are not quite sure what the API should look like for this. Would you mind sharing the proposed solution on the mailing list? It would be helpful to have input from the wider NumPy community.

[InessaPawson]

@eric-wieser Any thoughts on the proposed?