user706838
user706838

Reputation: 5380

How can I perform the following computation in a more elegant (and Pythonic) way?

Given the object data, which is of type numpy.ndarray, how can I perform the following in one line?

VERY_LOW = np.where(data<=-2, 'VERY_LOW', 'LOL').astype('S12') 
LOW = np.where((data>-2) & (data<=-1), 'LOW', VERY_LOW)
AVERAGE = np.where((data>-1) & (data<+1), 'AVERAGE', LOW)
HIGH = np.where((data>=+1) & (data<+2), 'HIGH', AVERAGE)
VERY_HIGH = np.where(data>=+2, 'VERY_HIGH', HIGH)

Basically, what I am trying to achieve is to assign a tag to each cell depending on its value (one out of five available).

Upvotes: 0

Views: 49

Answers (2)

HYRY
HYRY

Reputation: 97281

Try this:

import numpy as np

data = np.random.randint(-5, 5, size=20)

conditions = [
    [data <= -2, 'VERY_LOW'],
    [data <= -1, 'LOW'],
    [data <   1, 'AVERAGE'],
    [data <   2, 'HIGH'],
    [True      , 'VERY_HIGH']
]

condlist, choicelist = zip(*conditions)

np.select(condlist, choicelist)

the output is something like:

array(['VERY_HIGH', 'VERY_LOW', 'VERY_HIGH', 'VERY_HIGH', 'VERY_HIGH',
       'AVERAGE', 'VERY_HIGH', 'LOW', 'LOW', 'AVERAGE', 'AVERAGE',
       'VERY_LOW', 'HIGH', 'HIGH', 'VERY_HIGH', 'VERY_HIGH', 'VERY_LOW',
       'AVERAGE', 'VERY_HIGH', 'VERY_LOW'], 
      dtype='|S11')

If you can use Pandas, (there maybe some problem about equal conditions):

import pandas as pd
pd.cut(data, [-np.inf, -2, -1, 1, 2, np.inf], 
       labels=["VERY LOW", "LOW", "AVERAGE", "HIGH", "VERY_HIGH"])

Upvotes: 0

Hannes Ovr&#233;n
Hannes Ovr&#233;n

Reputation: 21831

You could write a function that maps a value to a tag, and then use the np.vectorize function to apply it.

def map_func(x):
    if x <= -2:
        return 'VERY LOW'
    elif <= -1:
        return 'LOW'
    # Keep adding more conditions here
    else:
        return 'OTHER'

vec_map_func = np.vectorize(map_func)
tag_array = vec_map_func(data)

Upvotes: 1

Related Questions