CODEWITHSUNDEEP
pythonarraysnumpy
InsertCreativityHere
InsertCreativityHere

Reputation: 41

Numpy: Imposing row dependent maximum on array

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'?

Upvotes: 3

Views: 1197

Answers (3)

keepAlive
keepAlive

Reputation: 6655

With

>>> a
array([[1, 4, 2, 3],
       [3, 1, 5, 4],
       [4, 3, 1, 2]])

Say you have 

>>> maxs = np.array([[3],[4],[3]])
>>> maxs
array([[3],
       [4],
       [3]])

What about doing

>>> a.clip(max=maxs)
array([[1, 3, 2, 3],
       [3, 1, 4, 4],
       [3, 3, 1, 2]])

Upvotes: 1

Brad Solomon
Brad Solomon

Reputation: 40878

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.

Upvotes: 3

Divakar
Divakar

Reputation: 221524

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]])

Upvotes: 3

Related Questions