Aaalex11
Aaalex11

Reputation: 15

Numpy array condition and apply a formula

I have a default numpy array (speed, pressure or temperature data) like this:

a=[[ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9.]
   [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9.]
   [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9.]
   [30. 31. 32. 33. 34. 35. 36. 37. 38. 39.]
   [40. 41. 42. 43. 44. 45. 46. 47. nan 49.]]

I need to apply the following conditions and then use the corresponding formula

a<5 (a*5)+4
a>5 (a**2.)-2

I tried using:

a[a<5]=(a[a<5]*5.)+4.

but it does not work and I have also used the method creating Boolean matrices and then multiplying them by the formulas corresponding to the condition, like this:

les=(a<5.).astype(float)
mayor=(a>5.).astype(float)

les=les*((a*5)+4)
mayor=mayor*((a**2.)-2)

b=les+mayor

This works but it uses many lines of code and I think it is impractical and I would like to know if there is an easier way to do this.

Upvotes: 1

Views: 212

Answers (2)

natalie
natalie

Reputation: 204

Try using a nested list comprehension

answer = [[(x*5)+4 if x<5 else (x**2.)-2 for x in row] for row in a]

This will essentially go row by row creating a new list for each row using the conditions you have defined to convert each element

Upvotes: 3

kuzand
kuzand

Reputation: 9806

As suggested by hpaulj in comments, you could use np.where:

>>> np.where(a<5, (a*5)+4, (a**2)-2)

array([[4., 9., 14., 19., 24., 23., 34., 47., 62., 79.],
       [4., 9., 14., 19., 24., 23., 34., 47., 62., 79.],
       [4., 9., 14., 19., 24., 23., 34., 47., 62., 79.],
       [898., 959., 1022., 1087., 1154., 1223., 1294., 1367., 1442., 1519.],
       [1598., 1679., 1762., 1847., 1934., 2023., 2114., 2207., nan, 2399.]])

However, according to the conditions that you provided:

a<5: (a*5)+4
a>5: (a**2.)-2

for a = 5 the value of 5 should be kept unchanged. Here's one way to do it:

b = a.copy()
b[a<5] = a[a<5]*5 + 4
b[a>5] = a[a>5]**2 - 2

Result:

array([[4., 9., 14., 19., 24., 5., 34., 47., 62., 79.],
       [4., 9., 14., 19., 24., 5., 34., 47., 62., 79.],
       [4., 9., 14., 19., 24., 5., 34., 47., 62., 79.],
       [898., 959., 1022., 1087., 1154., 1223., 1294., 1367., 1442., 1519.],
       [1598., 1679., 1762., 1847., 1934., 2023., 2114., 2207., nan, 2399.]])

Upvotes: 0

Related Questions