Search code examples
pythonnumpyarray-broadcasting

numpy: limiting min value of a scalar, 1D or nD array

Given a scalar, a 1D or an N-D numpy (numeric) array, I'd like to replace all values less than threshold with threshold. So, for example:

def fn(a, threshold):
  return ???

fn(2, 2.5) => 2.5                                                             # scalar
fn([1, 2, 3, 4], 2.5) => [2.5, 2.5, 3, 4]]                                    # 1-D
fn[[1, 2, 3, 4], [0, 2, 4, 6]], 2.5) => [[2.5, 2.5, 3, 4], [2.5, 2.5, 4, 6]]  # 2-D

(Note: For ease of reading, I've shown the arrays above with ordinary Python array syntax, but they're actually numpy.ndarrays.)

I could use if statements and dispatch on the type of a to handle each case. But I'm still wrapping my head around numpy's broadcast methods: is there a simple numpy idiom for handling this situation?

Solution

  • You can use maximum:

    np.maximum([[1, 2, 3, 4], [0, 2, 4, 6]], 2.5)

    Or clip:

    np.clip([[1, 2, 3, 4], [0, 2, 4, 6]], 2.5, np.inf)

    Output:

    array([[2.5, 2.5, 3. , 4. ],
       [2.5, 2.5, 4. , 6. ]])