Numpy: Imposing row dependent maximum on array

Question

Suppose I have the following array:

a = [[1, 4, 2, 3]
     [3, 1, 5, 4]
     [4, 3, 1, 2]]

What I'd like to do is impose a maximum value on the array, but have that maximum vary by row. For instance if I wanted to limit the 1st and 3rd row to a maximum value of 3, and the 2nd row to a value of 4, I could create something like:

[[1, 3, 2, 3]
 [3, 1, 4, 4]
 [3, 3, 1, 2]

Is there any better way than just looping over each row individually and setting it with 'nonzero'?

Answer 1

With numpy.clip (using the method version here):

a.clip(max=np.array([3, 4, 3])[:, None]) # np.clip(a, ...)
# array([[1, 3, 2, 3],
#        [3, 1, 4, 4],
#        [3, 3, 1, 2]])

Generalized:

def clip_2d_rows(a, maxs):
    maxs = np.asanyarray(maxs)
    if maxs.ndim == 1:
        maxs = maxs[:, np.newaxis]
    return np.clip(a, a_min=None, a_max=maxs)

You might be safer using the module-level function (np.clip) rather than the class method (np.ndarray.clip). The former uses a_max as a parameter, while the latter uses the builtin max as a parameter which is never a great idea.

Answer 2

With masking -

In [50]: row_lims = np.array([3,4,3])

In [51]: np.where(a > row_lims[:,None], row_lims[:,None], a)
Out[51]: 
array([[1, 3, 2, 3],
       [3, 1, 4, 4],
       [3, 3, 1, 2]])