如何编写一个使用 pandas 数据框变量作为输入的函数？

Question

我想编写一个函数，该函数将生成一个绘图并将一个变量作为输入。我是Python 的新手，对R 有更多经验。我见过涉及来自pandas 数据框的变量的函数通常使用apply()，但我不知道如何根据组有条件地选择行。

在下面的示例中，我想使用函数mykdeplot() 为变量X1 和X2 绘制图表。

%matplotlib inline
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns
import random as random

# generate data
group = ['a', 'b']
a = np.random.exponential(4, 10)
b = np.random.exponential(1/4, 10)
df = pd.DataFrame({'Group': np.repeat(group, 10),
                   'X1': np.random.randint(1,30,20),
                   'X2': np.concatenate((a,b)) })

# plot for X1  
width = 4
sns.kdeplot(np.array(df[df.Group == 'a'].X1), bw=width, label = "Group A")
sns.kdeplot(np.array(df[df.Group == 'b'].X1), bw=width, label = "Group B")

# plot for X2  
width = 4
sns.kdeplot(np.array(df[df.Group == 'a'].X2), bw=width, label = "Group A")
sns.kdeplot(np.array(df[df.Group == 'b'].X2), bw=width, label = "Group B")

# Write function (this is where I need help)   
def mykdeplot(var, width):
    sns.kdeplot(np.array(df[df.Group == 'a'].var), bw=width, label = "Group A")
    sns.kdeplot(np.array(df[df.Group == 'b'].var), bw=width, label = "Group B")

# call function (does not work)
mykdeplot(df.X1, 3)
mykdeplot(df.X2, 3)

Answer 1

def mykdeplot(df, var, width):
    sns.kdeplot(np.array(getattr(df[df.Group == 'a'], var)), bw=width, label = "Group A")
    sns.kdeplot(np.array(getattr(df[df.Group == 'b'], var)), bw=width, label = "Group B")

mykdeplot(df, 'X1', 3)