Напишите общую функцию c для вычисления условия c на основе столбца на основе другого значения столбца в pandas

Question

У меня есть df, как показано ниже.

Date                t_factor     time_in_days
2020-02-01             5             1
2020-02-06             14            6
2020-02-09             23            9
2020-02-03             23            3       
2020-03-11             38            40         
2020-02-20             29            20               
2020-02-13             30            13           
2020-02-29             100           29           
2020-03-26             70            55 

Из этого я хотел бы создать функцию, которая будет вычислять столбец с именем t_function на основе значений в time_in_days.

Условия для t_function:

if T0 <= time_in_days <T1:
      t_function = (a1*time_in_days) + a0
else if T1 < time_in_days <= T2:
    t_function = 14
else:
    t_function = a2(time_in_days)**2 + (a1*time_in_days) + a0 

Входные данные функции будут кадром данных, time_in_days, T0, T1, T2, a0, a1, a2)

Ожидаемый результат:

if T0 =1 , T1=4, T2= 35, a0=3, a1=2, a2=1


Date                t_factor     time_in_days     t_function
2020-02-01             5             1            5
2020-02-06             14            6            14
2020-02-09             23            9            14
2020-02-03             23            3            9 
2020-03-11             38            40           1683           
2020-02-20             29            20           14     
2020-02-13             30            13           14  
2020-02-29             100           29           14
2020-03-26             70            55           3138

Я пробовал код ниже

def t_function_df( df, time_in_days, T0, T1, T2, a0, a1, a2):
     df = df.copy()
     df['t_function'] = np.select( (df['time_in_days'].ge(T0) & df['time_in_days'].le(T1),
                             df['time_in_days'].gt(T1) & df['time_in_days'].le(T2)),
                            (df['time_in_days']*a1+a0, 14),   
                            (a2*df['time_in_days']**2) + df['time_in_days']*a1 + a0)

     return df[['Date', 'time_in_days', 't_function']]

Answer 1

np.select - очевидное решение здесь, но поскольку здесь у вас есть псевдокод, почему бы просто не запустить его!

if T0 <= time_in_days < T1:
      t_function = (a1 * time_in_days) + a0
else if T1 < time_in_days <= T2:
    t_function = 14
else:
    t_function = (a2 * time_in_days) ** 2 + (a1 * time_in_days) + a0 

Теперь станет,

T0 = 1; T1 = 4; T2 = 35; a0 = 3; a1 = 2; a2 = 1
condlist = ["@T0 <= time_in_days < @T1", "@T1 < time_in_days <= @T2"]
choicelist = ["(@a1 * time_in_days) + @a0", "14"]
default = "(@a2 * time_in_days) ** 2 + (@a1 * time_in_days) + @a0"

Символ "@" символ используется для поиска фактических переменных в памяти. Теперь вы можете комбинировать возможности numpy и pandas следующим образом:

np.select(condlist=[df.eval(c) for c in condlist], 
          choicelist=[df.eval(q) for q in choicelist], 
          default=df.eval(default))  
# array([   5,   14,   14,    9, 1683,   14,   14,   14, 3138], dtype=int64)

df['t_function_actual'] = np.select(
              condlist=[df.eval(c) for c in condlist], 
              choicelist=[df.eval(q) for q in choicelist], 
              default=df.eval(default))  
df     
         Date  t_factor  time_in_days  t_function  t_function_actual
0  2020-02-01         5             1           5                  5
1  2020-02-06        14             6          14                 14
2  2020-02-09        23             9          14                 14
3  2020-02-03        23             3           9                  9
4  2020-03-11        38            40        1683               1683
5  2020-02-20        29            20          14                 14
6  2020-02-13        30            13          14                 14
7  2020-02-29       100            29          14                 14
8  2020-03-26        70            55        3138               3138

Подробнее о eval читайте в моем сообщении здесь: Dynami c Оценка выражений в pandas используя pd.eval () .

Answer 2

Использовать пользовательскую функцию fx с np.select:

def fx(days, T0, T1, T2, a0, a1, a2):
    return np.select([days.ge(T0) & days.lt(T1), days.gt(T1) & days.le(T2)],
                     [a1*days + a0, 14], a2*(days)**2 + (a1*days) + a0)

df['t_function'] = fx(df['time_in_days'], T0 =1, T1=4, T2= 35, a0=3, a1=2, a2=1)

Результат:

         Date  t_factor  time_in_days  t_function
0  2020-02-01         5             1           5
1  2020-02-06        14             6          14
2  2020-02-09        23             9          14
3  2020-02-03        23             3           9
4  2020-03-11        38            40        1683
5  2020-02-20        29            20          14
6  2020-02-13        30            13          14
7  2020-02-29       100            29          14
8  2020-03-26        70            55        3138