Numpy "where" with multiple conditions

I try to add a new column "energy_class" to a dataframe "df_energy" which it contains the string "high" if the "consumption_energy" value > 400, "medium" if the "consumption_energy" value is between 200 and 400, and "low" if the "consumption_energy" value is under 200. I try to use np.where from numpy, but I see that numpy.where(condition[, x, y]) treat only two condition not 3 like in my case.

Any idea to help me please?

Thank you in advance

Try this: Using the setup from @Maxu

col         = 'consumption_energy'
conditions  = [ df2[col] >= 400, (df2[col] < 400) & (df2[col]> 200), df2[col] <= 200 ]
choices     = [ "high", 'medium', 'low' ]
    
df2["energy_class"] = np.select(conditions, choices, default=np.nan)


  consumption_energy energy_class
0                 459         high
1                 416         high
2                 186          low
3                 250       medium
4                 411         high
5                 210       medium
6                 343       medium
7                 328       medium
8                 208       medium
9                 223       medium

You can use a ternary:

np.where(consumption_energy > 400, 'high', 
         (np.where(consumption_energy < 200, 'low', 'medium')))

I like to keep the code clean. That's why I prefer np.vectorize for such tasks.

def conditions(x):
    if x > 400:
        return "High"
    elif x > 200:
        return "Medium"
    else:
        return "Low"

func = np.vectorize(conditions)
energy_class = func(df_energy["consumption_energy"])

Then just add numpy array as a column in your dataframe using:

df_energy["energy_class"] = energy_class

The advantage in this approach is that if you wish to add more complicated constraints to a column, it can be done easily. Hope it helps.

I would use the cut() method here, which will generate very efficient and memory-saving category dtype:

In [124]: df
Out[124]:
   consumption_energy
0                 459
1                 416
2                 186
3                 250
4                 411
5                 210
6                 343
7                 328
8                 208
9                 223

In [125]: pd.cut(df.consumption_energy,
                 [0, 200, 400, np.inf],
                 labels=['low','medium','high']
          )
Out[125]:
0      high
1      high
2       low
3    medium
4      high
5    medium
6    medium
7    medium
8    medium
9    medium
Name: consumption_energy, dtype: category
Categories (3, object): [low < medium < high]

WARNING: Be careful with NaNs

Always be careful that if your data has missing values np.where may be tricky to use and may give you the wrong result inadvertently.

Consider this situation:

df['cons_ener_cat'] = np.where(df.consumption_energy > 400, 'high', 
         (np.where(df.consumption_energy < 200, 'low', 'medium')))

# if we do not use this second line, then
#  if consumption energy is missing it would be shown medium, which is WRONG.
df.loc[df.consumption_energy.isnull(), 'cons_ener_cat'] = np.nan

Alternatively, you can use one-more nested np.where for medium versus nan which would be ugly.

IMHO best way to go is pd.cut. It deals with NaNs and easy to use.

Examples:

import numpy as np
import pandas as pd
import seaborn as sns

df = sns.load_dataset('titanic')

# pd.cut
df['age_cat'] = pd.cut(df.age, [0, 20, 60, np.inf], labels=['child','medium','old'])


# manually add another line for nans
df['age_cat2'] = np.where(df.age > 60, 'old', (np.where(df.age <20, 'child', 'medium')))
df.loc[df.age.isnull(), 'age_cat'] = np.nan

# multiple nested where
df['age_cat3'] = np.where(df.age > 60, 'old',
                         (np.where(df.age <20, 'child',
                                   np.where(df.age.isnull(), np.nan, 'medium'))))

# outptus
print(df[['age','age_cat','age_cat2','age_cat3']].head(7))
    age age_cat age_cat2 age_cat3
0  22.0  medium   medium   medium
1  38.0  medium   medium   medium
2  26.0  medium   medium   medium
3  35.0  medium   medium   medium
4  35.0  medium   medium   medium
5   NaN     NaN   medium      nan
6  54.0  medium   medium   medium