I have a slightly complex function that assigns a quality level to given data by a pre-defined step-wise logic (dependent on fixed borders and also on relative borders based on the real value). The function 'get_quality()' below does this for each row and using pandas DataFrame.apply is quite slow for huge datasets. So I'd like to vectorize this calculation. Obviously I could do something like df.groupby(pd.cut(df.ground_truth, [-np.inf, 10.0, 20.0, 50.0, np.inf]))
for the inner if-logic and then apply a similar sub-grouping within each group (based on the borders of each group), but how would I do that for the last bisect that depends on the given real/ground_truth value in each row?
Using df['quality'] = np.vectorize(get_quality)(df['measured'], df['ground_truth'])
is a lot faster already, but is there a real vectorized way to calculate the same 'quality' column?
import pandas as pd
import numpy as np
from bisect import bisect
quality_levels = ['WayTooLow', 'TooLow', 'OK', 'TooHigh', 'WayTooHigh']
# Note: to make the vertical borders always lead towards the 'better' score we use a small epsilon around them
eps = 0.000001
def get_quality(measured_value, real_value):
diff = measured_value - real_value
if real_value <= 10.0:
i = bisect([-4.0-eps, -2.0-eps, 2.0+eps, 4.0+eps], diff)
return quality_levels[i]
elif real_value <= 20.0:
i = bisect([-14.0-eps, -6.0-eps, 6.0+eps, 14.0+eps], diff)
return quality_levels[i]
elif real_value <= 50.0:
i = bisect([-45.0-eps, -20.0-eps, 20.0+eps, 45.0+eps], diff)
return quality_levels[i]
else:
i = bisect([-0.5*real_value-eps, -0.25*real_value-eps,
0.25*real_value+eps, 0.5*real_value+eps], diff)
return quality_levels[i]
N = 100000
df = pd.DataFrame({'ground_truth': np.random.randint(0, 100, N),
'measured': np.random.randint(0, 100, N)})
df['quality'] = df.apply(lambda row: get_quality((row['measured']), (row['ground_truth'])), axis=1)
print(df.head())
print(df.quality2.value_counts())
# ground_truth measured quality
#0 51 1 WayTooLow
#1 7 25 WayTooHigh
#2 38 95 WayTooHigh
#3 76 32 WayTooLow
#4 0 18 WayTooHigh
#OK 30035
#WayTooHigh 24257
#WayTooLow 18998
#TooLow 14593
#TooHigh 12117