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

时间:2016-03-17 20:38:55

标签: python function pandas seaborn

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

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

%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)

1 个答案:

答案 0 :(得分:0)

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)