CODEWITHSUNDEEP
pythonnumpynumpy-ndarray
Rachel W
Rachel W

Reputation: 137

numpy perform conditional operation across matrices

I have this bit of code that works on a single value that I am trying to perform over an entire matrix. I am starting with the matrix sst and using that to get the matrix pb_opt:

    if sst < -10:
        pb_opt = 0
    elif sst < -1:
        pb_opt = 1.13
    elif sst > 28.5:
        pb_opt = 4
    else:
        pb_opt = 1.2956 + (2.749e-1 * sst) + 6.17e-2 * (sst ** 2) - 2.05e-2 * (sst ** 3)
        +2.462e-3 * (sst ** 4) - 1.348e-4 * (sst ** 5) + 3.4132e-6 * (sst ** 6)
        -3.27e-8 * (sst ** 7)

The first couple if/else statements I have gotten to work by running:

    pb_opt = np.full(sst.shape, np.nan)
    pb_opt[sst < -10] = 0
    pb_opt[(-10 <= sst) & (sst < -1)] = 1.13
    pb_opt[sst > 28.5] = 4

I am struggling to the the last statement to work, however. I have tried several things and the closest I think I have gotten is this:

    pb_opt[(-1 >= sst) & (sst <= 28)] = (
        1.2956
        + (0.2749 * sst[(-1 >= sst) & (sst <= 28)])
        + 0.0617 * (sst[(-1 >= sst) & (sst <= 28)] ** 2)
        - 0.0205 * (sst[(-1 >= sst) & (sst <= 28)] ** 3)
        + 2.462e-3 * (sst[(-1 >= sst) & (sst <= 28)] ** 4)
        - 1.348e-4 * (sst[(-1 >= sst) & (sst <= 28)] ** 5)
        + 3.4132e-6 * (sst[(-1 >= sst) & (sst <= 28)] ** 6)
        - 3.27e-8 * (sst[(-1 >= sst) & (sst <= 28)] ** 7)
    )

Which, besides being a little ugly, gives me the error IndexError: 2-dimensional boolean indexing is not supported.. Does anyone know how I could accomplish this?

Upvotes: 1

Views: 63

Answers (1)

Hammad Ahmed
Hammad Ahmed

Reputation: 885

just do this, it will be a bit slower if the size of sst is vely large. if not then this is an elegant solution

pb_opt  = 1.2956                    \
        + 2.749e-1  * (sst ** 1)    \
        + 6.17e-2   * (sst ** 2)    \
        - 2.05e-2   * (sst ** 3)    \
        + 2.462e-3  * (sst ** 4)    \
        - 1.348e-4  * (sst ** 5)    \
        + 3.4132e-6 * (sst ** 6)    \
        - 3.27e-8   * (sst ** 7)


pb_opt[sst < -1  ] = 1.13
pb_opt[sst < -10 ] = 0
pb_opt[sst > 28.5] = 4

Upvotes: 2

Related Questions