gustavgans
gustavgans

Reputation: 5191

Many numpy array manipulations in one array

I need to translate a matlab code to python numpy code

I have 4 2-dimension arrays (A, B, C and D) and want to create a new one "res".

the code in matlab

index1 = find(A == 2 & B > C);
index2 = find(A == 2 & B >= D & B <= C);

res = uint8(zeros(100,100));
res(index1) = 5;
res(index2) = 4;

in Python/numpy I figured out I can create new arrays like this:

res = ((A == 2) & (B > C)) * 5
res = ((A == 2) & (B >= D) & (B <= C)) * 4

but how can I combine this 2 results in only one array? Like the matlab code does.

Upvotes: 0

Views: 50

Answers (1)

unutbu
unutbu

Reputation: 879621

The translation (in this case) is almost one-for-one:

index1 = ((A == 2) & (B > C)) 
index2 = ((A == 2) & (B >= D) & (B <= C)) 
res = np.zeros((100, 100), dtype='uint8')
res[index1] = 5
res[index2] = 4

Alternatively, you could define res as

res = np.where(A == 2, np.where(B > C, 5, np.where(B >= D, 4, 0)), 0)

though I don't think that helps readability :)


PS. As Mr. E suggests, mask1 might be a better variable name than index1 since index (usually) suggests integer values while mask suggest booleans, which is what we have here.

Upvotes: 3

Related Questions