Search code examples
pythonnumpymultidimensional-arraynumpy-ndarrayarray-broadcasting

How can I update a 2d numpy array based on condition, but still using it absolute value?


having the following function:

def one_more_fix(matrix):
height, width = matrix.shape

for i in range(height):
    for j in range(width):
        if abs(matrix[i,j]) >= 180:
            matrix[i,j] = abs(matrix[i,j]) - 360

Example input and output:

simple_matrix = np.array([
[179, 181, -182],
[179, 181, -182],
[361, -362, -183],

])

array([[ 179, -179, -178],
       [ 179, -179, -178],
       [   1,    2, -177]])

I would like to substitute this function with some "numpy magic" to remove the loop and make it faster. Tried something like this, but can't find a way how to update the original matrix:

np.abs(simple_matrix1[np.abs(simple_matrix1) >= 180]) - 360

Any suggestions?


Solution

  • First get the indices to be changed (i.e those that satisfy condition)

    idx = (np.abs(simple_matrix) >= 180)  # this is a boolean array 
    

    Now update the values of those indices

    simple_matrix[idx] = np.abs(simple_matrix[idx])-360